Add files using upload-large-folder tool

README.md

@@ -1,6 +1,8 @@
 ---
+
 license: apache-2.0
 language:
+
 - en
 pipeline_tag: feature-extraction
 tags:
@@ -9,11 +11,28 @@ tags:
 - custom_code
 - multimodal
 base_model: mispeech/midashenglm-7b-0804-fp32
+
 ---
 
-# AudioEmb (Finetune)
+<h1 align="center">ALM2Vec-FT</h1>
+
+<p align="center">
+  <a href="https://arxiv.org/abs/xxxx.xxxxx"><img src="https://img.shields.io/badge/Paper-arXiv-b31b1b?logo=arxiv&logoColor=white" alt="Paper"></a>
+  <a href="https://caml-labs.github.io/ALM2Vec"><img src="https://img.shields.io/badge/Project-Page-1f6feb?logo=googlechrome&logoColor=white" alt="Project Page"></a>
+  <a href="https://github.com/caml-labs/ALM2Vec"><img src="https://img.shields.io/badge/Code-GitHub-181717?logo=github&logoColor=white" alt="GitHub"></a>
+</p>
+
+**ALM2Vec** is a universal audio embedding model for retrieval, derived from a pretrained large audio–language model (LALM). Instead of being optimized only for audio–caption matching like conventional contrastive dual-encoders, it transfers the audio understanding, instruction-following, and reasoning abilities of LALMs into a single unified embedding space that works across audio domains, task types, and user intents.
+
+Its key feature is **instruction-aware retrieval**: a natural-language instruction guides the embedding, so the *same* audio can be encoded differently for different needs. This supports:
+
+- **Instruction-aware retrieval** — focus the embedding on a specific aspect of the audio.
+- **Text ↔ audio retrieval** — bidirectional matching between audio and text.
+- **Audio question answering** — match an audio query plus a question against candidate answers.
+
+ALM2Vec achieves competitive results on standard audio and speech retrieval benchmarks while adding these controllable retrieval capabilities. See the [project page](https://caml-labs.github.io/ALM2Vec/) for interactive demos.
 
-Audio-text embedding model for retrieval, based on [MiDashengLM](https://huggingface.co/mispeech/midashenglm-7b-0804-fp32). Checkpoint: **finetune** stage.
+This repository hosts the **finetune** checkpoint, built on [MiDashengLM](https://huggingface.co/mispeech/midashenglm-7b-0804-fp32).
 
 Requirements: `transformers>=4.52`, `torch`, `safetensors`, and `torchaudio` for non-WAV audio. Requires a GPU (~31GB weights) and `trust_remote_code=True`.
 
@@ -23,7 +42,7 @@ Requirements: `transformers>=4.52`, `torch`, `safetensors`, and `torchaudio` for
 import torch
 from transformers import AutoModel, AutoTokenizer
 
-repo_id = "cara-ai/AudioEmb-finetune"
+repo_id = "cara-ai/ALM2Vec-FT"
 tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
 model = AutoModel.from_pretrained(
     repo_id, trust_remote_code=True, torch_dtype=torch.float32
@@ -41,3 +60,77 @@ similarity = query_emb @ doc_emb.T
 print(similarity)
 ```
 
+## Results
+
+**ALM2Vec-FT** is the checkpoint hosted in this repository; **ALM2Vec-PT** is the pretrain variant. In every table, **bold** marks the best score and <u>underline</u> the second best.
+
+### Text–audio retrieval — AudioCaps
+
+| Method | T→A R@1 | T→A R@5 | T→A R@10 | A→T R@1 | A→T R@5 | A→T R@10 |
+| --- | --- | --- | --- | --- | --- | --- |
+| LAION-CLAP | 36.1 | 71.8 | 83.9 | 46.8 | <u>82.9</u> | <u>90.7</u> |
+| MS-CLAP | 15.4 | 47.2 | 64.5 | 32.0 | 66.0 | 79.2 |
+| WavCaps-CLAP-PT | 39.7 | 74.5 | 86.1 | 51.7 | 82.3 | 90.6 |
+| WavCaps-CLAP-FT | <u>42.2</u> | <u>76.5</u> | <u>87.1</u> | <u>54.6</u> | **85.2** | **92.4** |
+| JINA-Embed.-v5 | 20.4 | 50.3 | 64.4 | 23.1 | 52.7 | 67.2 |
+| **ALM2Vec-PT** | 40.0 | 74.5 | 85.9 | 43.8 | 74.3 | 86.5 |
+| **ALM2Vec-FT** | **43.2** | **78.0** | **87.8** | **55.5** | 80.0 | 88.2 |
+
+### Text–audio retrieval — Clotho
+
+| Method | T→A R@1 | T→A R@5 | T→A R@10 | A→T R@1 | A→T R@5 | A→T R@10 |
+| --- | --- | --- | --- | --- | --- | --- |
+| LAION-CLAP | 16.1 | 38.3 | 51.1 | 22.7 | 48.5 | 60.8 |
+| MS-CLAP | 15.6 | 38.9 | 51.4 | 22.1 | 48.9 | 62.0 |
+| WavCaps-CLAP-PT | 19.5 | 45.2 | 58.2 | 23.4 | 50.9 | 63.4 |
+| WavCaps-CLAP-FT | <u>19.7</u> | <u>45.7</u> | <u>59.4</u> | <u>26.9</u> | <u>52.6</u> | <u>64.9</u> |
+| JINA-Embed.-v5 | 9.2 | 23.9 | 35.0 | 10.5 | 24.7 | 34.3 |
+| **ALM2Vec-PT** | 19.2 | 43.4 | 55.7 | 17.9 | 39.4 | 52.2 |
+| **ALM2Vec-FT** | **24.8** | **52.9** | **65.8** | **27.9** | **52.7** | **66.3** |
+
+
+### Speech retrieval — LibriSQA
+
+
+| Method         | T→S R@1 | T→S R@5 | T→S R@10 | S→T R@1 | S→T R@5 | S→T R@10 |
+| -------------- | --------------- | -------- | -------- | --------------- | -------- | -------- |
+| LAION-CLAP †   | 0.0             | 0.1      | 0.8      | 0.1             | 0.2      | 0.6      |
+| Whisper+BGE    | 83.7            | 93.3     | 94.9     | 85.2            | 93.4     | 95.3     |
+| CLSR           | **85.0**        | <u>93.4</u> | <u>95.0</u> | <u>85.5</u>        | <u>94.0</u> | <u>95.6</u> |
+| **ALM2Vec-PT** | 43.7            | 64.5     | 72.8     | 11.2            | 24.9     | 34.1     |
+| **ALM2Vec-FT** | <u>84.7</u>        | **94.1** | **95.8** | **86.0**        | **95.2** | **97.2** |
+
+
+### Audio understanding — MMAU-mini (accuracy)
+
+
+| Method             | Overall  | Music    | Sound    | Speech   |
+| ------------------ | -------- | -------- | -------- | -------- |
+| GPT-4o Audio ‡     | 60.8     | 63.2     | 64.6     | 56.3     |
+| Gemini 2.5 Pro ‡   | <u>71.6</u> | <u>75.1</u> | 71.5     | 68.3     |
+| Qwen2.5-Omni ‡     | 71.5     | 65.9     | <u>78.1</u> | <u>70.6</u> |
+| Audio Flamingo 3 ‡ | **73.1** | **76.9** | 66.1     | **73.9** |
+| **ALM2Vec-PT**     | 66.3     | 62.3     | **78.7** | 58.0     |
+| **ALM2Vec-FT**     | 63.0     | 61.7     | 74.8     | 52.6     |
+
+
+† LAION-CLAP is not trained for speech and effectively fails on LibriSQA; shown for reference.
+‡ Generative large audio–language models, listed as reference upper bounds rather than directly comparable retrieval baselines.
+
+## Citation
+
+If you find this work useful, please consider citing:
+
+```
+@article{ALM2Vec2026,
+  title={ALM2Vec: Learning Audio Embeddings for Universal
+        Audio Retrieval with Large Audio-Language Models},
+  author={TBD},
+  journal={arXiv preprint arXiv:TBD},
+  year={2026}
+}
+```
+
+## Acknowledgement
+
+ALM2Vec is built on [MiDashengLM](https://github.com/xiaomi-research/dasheng-lm) and further trained for universal audio retrieval. We thank MiDashengLM and its underlying [Dasheng](https://github.com/RicherMans/Dasheng) audio encoder for their open-source contributions.

config.json

@@ -1,12 +1,12 @@
 {
-  "model_type": "audio_emb",
+  "model_type": "alm2vec",
   "architectures": [
-    "AudioEmbModel"
+    "ALM2VecModel"
   ],
   "torch_dtype": "float32",
   "transformers_version": "5.0.0.dev0",
   "auto_map": {
-    "AutoConfig": "configuration_audio_emb.AudioEmbConfig",
-    "AutoModel": "modeling_audio_emb.AudioEmbModel"
+    "AutoConfig": "configuration_alm2vec.ALM2VecConfig",
+    "AutoModel": "modeling_alm2vec.ALM2VecModel"
   }
 }

configuration_alm2vec.py

@@ -0,0 +1,5 @@
+from transformers import PreTrainedConfig
+
+
+class ALM2VecConfig(PreTrainedConfig):
+    model_type = "alm2vec"

modeling_alm2vec.py

@@ -0,0 +1,1118 @@
+"""ALM2Vec audio-text embedding model."""
+
+import math
+from torch.utils.checkpoint import checkpoint
+import wave
+from io import BytesIO
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from urllib.parse import urlparse
+from urllib.request import urlopen
+
+from .configuration_alm2vec import ALM2VecConfig
+
+import collections
+import collections.abc
+
+from dataclasses import dataclass
+import torch
+import torch.nn as nn
+import torchaudio.functional as F
+from torch import Tensor
+from torch.nn.functional import scaled_dot_product_attention
+from typing import Any, Dict, Callable, Iterable, List, Optional, Sequence, Tuple, Union, cast
+
+from transformers import PreTrainedModel, PreTrainedConfig, GenerationMixin
+from transformers import AutoTokenizer
+
+from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
+    Qwen2_5OmniTextConfig,
+)
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
+    Qwen2_5OmniThinkerTextModel,
+)
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from transformers.utils import can_return_tuple
+
+import copy
+
+try:
+    import torchaudio
+except ImportError:
+    torchaudio = None
+
+_Tuple2 = Union[int, Tuple[int, int], Sequence[int]]
+
+TARGET_SR = 16000
+
+QUERY_INSTRUCTION = "Based on the question asked in the text query and context in the audio query, retrieve the relevant text document associated with that question."
+DOC_INSTRUCTION = "Represent the user's input."
+
+
+def _resolve_tuple2(x: _Tuple2) -> Tuple[int, int]:
+    if isinstance(x, collections.abc.Sequence):
+        assert len(x) == 2, (
+            f"Expected a sequence of length 2, got {x} with length {len(x)}"
+        )
+        return cast(Tuple[int, int], tuple(x))
+    return (x, x)
+
+
+
+
+DASHENG_ARCH_CONFIG = {
+    "audio_encoder_config": {
+        "attn_drop_rate": 0.0,
+        "center": True,
+        "depth": 32,
+        "drop_rate": 0.0,
+        "embed_dim": 1280,
+        "f_max": 8000.0,
+        "f_min": 0.0,
+        "hop_length": 160,
+        "init_values": None,
+        "input_channels": 1,
+        "mlp_ratio": 4.0,
+        "model_type": "midashenglm_dasheng_encoder",
+        "n_fft": 512,
+        "n_mels": 64,
+        "num_heads": 16,
+        "outputdim": 527,
+        "patch_size": [
+        64,
+        4
+        ],
+        "patch_stride": [
+        64,
+        4
+        ],
+        "qkv_bias": True,
+        "sample_rate": 16000,
+        "target_length": 1008,
+        "win_length": 512
+    }, 
+
+    "audio_projector_config": {
+        "in_dim": 1280,
+        "downsample_rate": 5,
+        "out_dim": 3584,
+    },
+
+    "text_config": {
+        "attention_dropout": 0.0,
+        "hidden_act": "silu",
+        "hidden_size": 3584,
+        "init_std": 0.02,
+        "initializer_range": 0.02,
+        "intermediate_size": 18944,
+        "max_position_embeddings": 32768,
+        "max_window_layers": 28,
+        "model_type": "qwen2_5_omni_text",
+        "num_attention_heads": 28,
+        "num_hidden_layers": 28,
+        "num_key_value_heads": 4,
+        "rms_norm_eps": 1e-06,
+        "rope_scaling": {
+        "mrope_section": [
+            16,
+            24,
+            24
+        ],
+        "rope_type": "default",
+        "type": "default"
+        },
+        "rope_theta": 1000000.0,
+        "sliding_window": 32768,
+        "use_cache": True,
+        "use_sliding_window": False,
+        "vocab_size": 152064
+    },
+    
+    "lite_random_decoder_config": {
+        "attention_dropout": 0.0,
+        "hidden_act": "silu",
+        "hidden_size": 576,
+        "init_std": 0.02,
+        "initializer_range": 0.02,
+        "intermediate_size": 1536,
+        "max_position_embeddings": 2048,
+        "max_window_layers": 12,
+        "model_type": "qwen2_5_omni_text",
+        "num_attention_heads": 8,
+        "num_hidden_layers": 12,
+        "num_key_value_heads": 4,
+        "rms_norm_eps": 1e-06,
+        "rope_scaling": {
+            "mrope_section": [
+                12,
+                12,
+                12
+            ],
+            "rope_type": "default",
+            "type": "default"
+        },
+        "rope_theta": 1000000.0,
+        "sliding_window": 2048,
+        "use_cache": True,
+        "use_sliding_window": False,
+        "vocab_size": 152064
+    }
+}
+
+
+class DashengConfig(PreTrainedConfig):
+    model_type = "midashenglm_dasheng_encoder"
+
+    def __init__(
+        self,
+        embed_dim: int = 768,
+        outputdim: int = 527,
+        patch_size: Union[int, Tuple[int, int]] = 16,
+        patch_stride: Union[int, Tuple[int, int]] = 16,
+        input_channels: int = 1,
+        target_length: int = 1012,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        init_values: Optional[float] = None,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        f_min: float = 0.0,
+        f_max: float = 8000.0,
+        center: bool = True,
+        win_length: int = 512,
+        hop_length: int = 160,
+        sample_rate: int = 16000,
+        n_fft: int = 512,
+        n_mels: int = 64,
+        **kwargs,
+    ):
+        self.embed_dim = embed_dim
+        self.outputdim = outputdim
+        self.patch_size = patch_size
+        self.patch_stride = patch_stride
+        self.input_channels = input_channels
+        self.target_length = target_length
+        self.depth = depth
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.qkv_bias = qkv_bias
+        self.init_values = init_values
+        self.drop_rate = drop_rate
+        self.attn_drop_rate = attn_drop_rate
+        self.f_min = f_min
+        self.f_max = f_max
+        self.center = center
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.n_mels = n_mels
+        super().__init__(**kwargs)
+
+
+class AudioPatchEmbed(nn.Module):
+    def __init__(
+        self,
+        input_size: _Tuple2 = 64,
+        patch_size: _Tuple2 = 16,
+        patch_stride: _Tuple2 = 16,
+        in_chans: int = 1,
+        embed_dim: int = 768,
+        norm_layer: Optional[Callable] = None,
+        flatten: bool = False,
+    ):
+        super().__init__()
+        self.input_size = _resolve_tuple2(input_size)
+        self.patch_size = _resolve_tuple2(patch_size)
+        self.patch_stride = _resolve_tuple2(patch_stride)
+        self.grid_size = (
+            self.input_size[0] // self.patch_stride[0],
+            self.input_size[1] // self.patch_stride[1],
+        )
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_stride,
+        )
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x)
+        if self.flatten:
+            x = torch.permute(
+                torch.flatten(x, 2, 3), (0, 2, 1)
+            )  # rearrange(x, "b c f t -> b (f t) c")
+        x = self.norm(x)
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class DashengMlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        drop: float = 0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class DashengAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None):
+        B, N, C = x.shape
+        q, k, v = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+            .unbind(0)
+        )
+        x = scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=mask[:, None, None, :] if mask is not None else None,
+        )
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class DashengBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values: Optional[float] = None,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
+        self.attn = DashengAttention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.ls1 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
+        self.mlp = DashengMlp(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            drop=drop,
+        )
+        self.ls2 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+
+    # Kwargs usually has a mask parameter that is passed to Attention
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        x = x + self.ls1(self.attn(self.norm1(x), mask))
+        x = x + self.ls2(self.mlp(self.norm2(x)))
+        return x
+
+
+class DashengFrontend(nn.Module):
+    def __init__(self, config: DashengConfig):
+        super().__init__()
+        self.config = config
+
+        spectrogram_window, melscale_fbanks = self._build_frontend_buffers()
+        self.register_buffer(
+            "spectrogram_window",
+            spectrogram_window,
+            persistent=False,
+        )
+        self.spectrogram_window: torch.Tensor
+        self.register_buffer("melscale_fbanks", melscale_fbanks, persistent=False)
+        self.melscale_fbanks: torch.Tensor
+
+    def _build_frontend_buffers(self) -> tuple[torch.Tensor, torch.Tensor]:
+        # Build on CPU explicitly: from_pretrained may construct modules on meta device.
+        with torch.device("cpu"):
+            spectrogram_window = torch.hann_window(
+                self.config.win_length,
+                dtype=torch.float32,
+            )
+            melscale_fbanks = F.melscale_fbanks(
+                n_freqs=self.config.n_fft // 2 + 1,
+                f_min=self.config.f_min,
+                f_max=self.config.f_max,
+                n_mels=self.config.n_mels,
+                sample_rate=self.config.sample_rate,
+            ).to(torch.float32)
+        return spectrogram_window, melscale_fbanks
+
+    def ensure_frontend_buffers(self, device: torch.device) -> None:
+        """Self-heal non-persistent audio frontend buffers if corrupted/uninitialized."""
+        expected_win_shape = (self.config.win_length,)
+        expected_fb_shape = (self.config.n_fft // 2 + 1, self.config.n_mels)
+
+        def _is_bad(name: str, tensor: torch.Tensor, expected_shape: tuple[int, ...]) -> bool:
+            if tensor is None:
+                return True
+            if getattr(tensor, "is_meta", False):
+                return True
+            if tuple(tensor.shape) != expected_shape:
+                return True
+            t = tensor.detach().float()
+            if not torch.isfinite(t).all().item():
+                return True
+            if t.numel() > 0 and t.abs().max().item() > 1e6:
+                return True
+            return False
+
+        win_bad = _is_bad("spectrogram_window", self.spectrogram_window, expected_win_shape)
+        fb_bad = _is_bad("melscale_fbanks", self.melscale_fbanks, expected_fb_shape)
+        if win_bad or fb_bad:
+            new_win, new_fb = self._build_frontend_buffers()
+            self.spectrogram_window = new_win.to(device=device)
+            self.melscale_fbanks = new_fb.to(device=device)
+            print(
+                f"[WARN] Rebuilt frontend buffers (win_bad={win_bad}, fb_bad={fb_bad})",
+                flush=True,
+            )
+        else:
+            if self.spectrogram_window.device != device:
+                self.spectrogram_window = self.spectrogram_window.to(device=device)
+            if self.melscale_fbanks.device != device:
+                self.melscale_fbanks = self.melscale_fbanks.to(device=device)
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        self.ensure_frontend_buffers(waveform.device)
+
+        spectrogram = F.spectrogram(
+            waveform=waveform.to(torch.float32),
+            pad=0,
+            window=self.spectrogram_window,
+            n_fft=self.config.n_fft,
+            hop_length=self.config.hop_length,
+            win_length=self.config.win_length,
+            power=2,
+            normalized=False,
+            center=self.config.center,
+        )
+        mel_spectrogram = (spectrogram.mT @ self.melscale_fbanks.to(torch.float32)).mT
+        # x has shape [batch, freq, time].
+        # F.amplitude_to_DB accepts inputs shaped as:
+        #   - [freq, time]
+        #   - [channel, freq, time]
+        #   - [..., channel, freq, time]
+        # Here we insert a channel dimension of size 1 before calling it,
+        # then remove that extra dimension afterward.
+        log_mel_spectrogram = F.amplitude_to_DB(
+            mel_spectrogram.unsqueeze(1),
+            multiplier=10,
+            amin=1e-10,
+            db_multiplier=0,
+            top_db=120,
+        ).squeeze(1)
+        return log_mel_spectrogram.to(waveform.dtype)
+
+
+class FixedAffine2d(nn.Module):
+    """
+    Per-channel fixed affine transform:
+        y = x * scale + bias
+    where scale/bias are broadcast on (B, C, H, W).
+    """
+
+    def __init__(self, scale: torch.Tensor, bias: torch.Tensor):
+        super().__init__()
+        self.register_buffer("scale", scale.reshape(1, -1, 1, 1))
+        self.register_buffer("bias", bias.reshape(1, -1, 1, 1))
+
+    @classmethod
+    def from_batchnorm2d(cls, bn: nn.BatchNorm2d) -> "FixedAffine2d":
+        if bn.running_mean is None or bn.running_var is None:
+            raise ValueError("BatchNorm2d must have running stats to be converted.")
+
+        if bn.affine:
+            gamma = bn.weight.detach()
+            beta = bn.bias.detach()
+        else:
+            gamma = torch.ones_like(bn.running_mean)
+            beta = torch.zeros_like(bn.running_mean)
+
+        running_mean = bn.running_mean.detach()
+        running_var = bn.running_var.detach()
+
+        scale = gamma / torch.sqrt(running_var + bn.eps)
+        bias = beta - running_mean * scale
+        return cls(scale=scale, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.scale + self.bias
+
+
+class DashengAudioTransformer(PreTrainedModel):
+    config_class = DashengConfig
+    supports_gradient_checkpointing = True
+
+    def __init__(self, config: DashengConfig):
+        super().__init__(config)
+
+        self.target_length = config.target_length
+        self.embed_dim = config.embed_dim
+        self.hop_length = config.hop_length
+        self.gradient_checkpointing = False
+
+        self.front_end = DashengFrontend(config)
+
+        self.init_bn = nn.BatchNorm2d(config.n_mels, momentum=0.01)
+
+        self.patch_embed = AudioPatchEmbed(
+            input_size=(config.n_mels, config.target_length),
+            embed_dim=config.embed_dim,
+            in_chans=config.input_channels,
+            patch_size=config.patch_size,
+            flatten=False,
+            patch_stride=config.patch_stride,
+        )
+
+        self.time_pos_embed = nn.Parameter(
+            torch.randn(1, config.embed_dim, 1, self.patch_embed.grid_size[1]) * 0.02
+        )
+        self.freq_pos_embed = nn.Parameter(
+            torch.randn(1, config.embed_dim, self.patch_embed.grid_size[0], 1) * 0.02
+        )
+
+        self.pos_drop = nn.Dropout(p=config.drop_rate)
+        self.blocks = nn.ModuleList(
+            DashengBlock(
+                dim=config.embed_dim,
+                num_heads=config.num_heads,
+                mlp_ratio=config.mlp_ratio,
+                qkv_bias=config.qkv_bias,
+                init_values=config.init_values,
+                drop=config.drop_rate,
+                attn_drop=config.attn_drop_rate,
+            )
+            for _ in range(config.depth)
+        )
+        self.norm = nn.LayerNorm(config.embed_dim, eps=1e-6)
+
+        self.post_init()
+
+    def replace_init_bn_with_fixed_affine(self):
+        """
+        Call this after checkpoint is loaded and before inference/export.
+        """
+        if isinstance(self.init_bn, nn.BatchNorm2d):
+            self.init_bn.eval()
+            self.init_bn = FixedAffine2d.from_batchnorm2d(self.init_bn)
+
+    def forward_features(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        t = x.shape[-1]
+        x = x + self.time_pos_embed[:, :, :, :t]
+        x = (
+            x + self.freq_pos_embed[:, :, :, :]
+        )  # Just to support __getitem__ in posembed
+        x = torch.permute(
+            torch.flatten(x, 2, 3), (0, 2, 1)
+        )  # rearrange(x, "b c f t -> b (f t) c")
+        x = self.pos_drop(x)
+        for block in self.blocks:
+            if self.gradient_checkpointing and self.training:
+                x = self._gradient_checkpointing_func(block, x, mask)
+            else:
+                x = block(x, mask)
+        x = self.norm(x)
+        return x
+
+    def _to_mask(self, lengths: torch.Tensor, max_length: int) -> torch.Tensor:
+        batch_size = len(lengths)
+        idx = torch.arange(max_length, device=lengths.device)
+        idx = idx.repeat(batch_size).view(batch_size, max_length)
+        mask = (idx < lengths.unsqueeze(-1)).bool()
+        return mask
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_length: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        x = self.front_end(x)
+        target_length_in_patches = self.target_length // 4
+        x = x.unsqueeze(1)
+        x = torch.permute(x, (0, 2, 1, 3))
+        x = self.init_bn(x)
+        x = torch.permute(x, (0, 2, 1, 3))
+
+        x = self.patch_embed(x)
+        t = x.shape[-1]
+
+        input_splits = x.split(target_length_in_patches, dim=-1)
+
+        if x_length is not None:
+            assert len(x_length) == len(x), (
+                "batchsizes of input x and x_length need to be same"
+            )
+            assert x_length.ndim == 1, "Lengths are of size (B,)"
+            scaled_lengths = (x_length / (self.hop_length * 4)).long()
+            mask = self._to_mask(max_length=t, lengths=scaled_lengths)
+            split_masks = mask.split(target_length_in_patches, dim=-1)
+        else:
+            mask = None
+            split_masks = [None] * len(input_splits)
+
+        outputs = []
+
+        for split_x, split_mask in zip(input_splits, split_masks):
+            split_x = self.forward_features(split_x, mask=split_mask)
+            outputs.append(split_x)
+        x = torch.cat(outputs, dim=1)
+
+        return x, mask
+
+
+class AudioProjectorSubsample(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        downsample_rate=5,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.k = downsample_rate
+        self.out_dim = out_dim
+        self.net = nn.Sequential(
+            nn.Linear(in_dim * self.k, out_dim, dtype=dtype),
+            nn.GELU(),
+            nn.Linear(out_dim, out_dim, dtype=dtype),
+        )
+
+    def forward(self, x, mask=None):
+        batch_size, seq_len, dim = x.shape
+        num_frames_to_discard = seq_len % self.k
+        if num_frames_to_discard > 0:
+            x = x[:, :-num_frames_to_discard, :]
+            if mask is not None:
+                mask = mask[:, :-num_frames_to_discard]
+        if mask is None:
+            mask = torch.ones(x.shape[:-1], dtype=torch.long, device=x.device)
+        x = x.reshape(
+            batch_size, -1, self.k * dim
+        )  # rearrange(x, "b (s k) d -> b s (k d)", k=self.k)
+        x = self.net(x)
+        mask = mask.reshape(
+            batch_size, -1, self.k
+        )  # rearrange(mask, "b (s k) -> b s k", k=self.k)
+        mask = mask.any(dim=-1).long()
+        return x, mask
+
+
+
+@dataclass
+class Qwen25OmniTextModelOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+class Qwen25OmniThinkerTextOnlyDecoder(PreTrainedModel, GenerationMixin):
+    config_class = Qwen2_5OmniTextConfig
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+    _supports_static_cache = True
+
+    def __init__(self, config: Qwen2_5OmniTextConfig):
+        super().__init__(config)
+        self.model = Qwen2_5OmniThinkerTextModel._from_config(config)
+        self.lm_head = nn.Linear(
+            config.hidden_size,
+            config.vocab_size,
+            bias=False,
+        )
+        self.post_init()
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, Qwen25OmniTextModelOutput]:
+        if attention_mask is not None and position_ids is None:
+            position_ids = (
+                attention_mask.long()
+                .cumsum(dim=-1)
+                .masked_fill_(attention_mask == 0, 1)
+                - 1
+            )
+
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            cache_position=cache_position,
+            return_dict=True,
+        )
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        loss = (
+            self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+            if labels is not None
+            else None
+        )
+
+        return Qwen25OmniTextModelOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+# Hardcoded architecture / training choices (not exposed in config.json).
+USE_LOGIT_SCALE = True
+HIDDEN_SIZE = 3584
+
+
+class ALM2VecModel(PreTrainedModel):
+    config_class = ALM2VecConfig
+
+    def __init__(self, config: ALM2VecConfig):
+        super().__init__(config)
+        text_config = Qwen2_5OmniTextConfig(**DASHENG_ARCH_CONFIG["text_config"])
+        decoder = Qwen25OmniThinkerTextOnlyDecoder(text_config)
+        self.model = decoder.model
+        self.dasheng = DashengAudioTransformer(
+            DashengConfig(**DASHENG_ARCH_CONFIG["audio_encoder_config"])
+        )
+        self.dasheng_down = AudioProjectorSubsample(**DASHENG_ARCH_CONFIG["audio_projector_config"])
+        self.dasheng_proj = nn.Identity()
+
+        # Placeholder shapes match exported checkpoints (overwritten by load_state_dict).
+        self.register_buffer("audio_start_token", torch.zeros(1, dtype=torch.long))
+        self.register_buffer("audio_end_token", torch.zeros(1, dtype=torch.long))
+        self.register_buffer("eos_token", torch.zeros(7, dtype=torch.long))
+
+        self.hidden_size = self.model.config.hidden_size
+        self.use_checkpointing = False
+        self.checkpoint_reentrant = False
+
+        if USE_LOGIT_SCALE:
+            init_value = math.log(1 / 0.07)
+            self.register_buffer(
+                "logit_scale", torch.tensor([init_value], dtype=torch.float32)
+            )
+        else:
+            self.logit_scale = None
+
+        self.siglip_head = nn.Linear(self.hidden_size, self.hidden_size)
+        self.dasheng.replace_init_bn_with_fixed_affine()
+        self._tokenizer = None
+        self.post_init()
+
+    def set_tokenizer(self, tokenizer) -> None:
+        """Attach a tokenizer for high-level encode APIs."""
+        self._tokenizer = tokenizer
+
+    def _resolve_tokenizer(self):
+        if self._tokenizer is not None:
+            return self._tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(self.name_or_path, trust_remote_code=True)
+        self._tokenizer = tokenizer
+        return tokenizer
+
+    @staticmethod
+    def _to_list(x: Any) -> list[Any]:
+        if x is None:
+            return []
+        if isinstance(x, (list, tuple)):
+            return list(x)
+        return [x]
+
+    @staticmethod
+    def _build_prompt(instruction: str, text: Optional[str]) -> str:
+        system_part = f"<|im_start|>system\n{instruction}<|im_end|>\n"
+        if text is None:
+            user_part = "<|im_start|>user\n"
+        else:
+            user_part = f"<|im_start|>user\n{text}"
+        return system_part + user_part
+
+    def _prepare_text_batch(
+        self,
+        tokenizer,
+        texts: list[Optional[str]],
+        *,
+        task: str,
+        instruction: Optional[str],
+        device: torch.device,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if task not in ("query", "document"):
+            raise ValueError(f"Unsupported task={task}. Use 'query' or 'document'.")
+        default_instruction = QUERY_INSTRUCTION if task == "query" else DOC_INSTRUCTION
+        instruction = instruction or default_instruction
+        prompts = [self._build_prompt(instruction, text) for text in texts]
+        encoded = tokenizer(prompts, padding=True, add_special_tokens=False, return_tensors="pt")
+        text_ids = encoded["input_ids"].to(device)
+        text_lens = encoded["attention_mask"].to(device).sum(dim=1)
+        return text_ids, text_lens
+
+    def _load_audio_path(self, path: Union[str, Path], target_sr: int) -> torch.Tensor:
+        raw_path = str(path)
+        parsed = urlparse(raw_path)
+        is_remote_url = parsed.scheme in ("http", "https")
+        local_path = None if is_remote_url else Path(path)
+        suffix_source = Path(parsed.path) if is_remote_url else local_path
+        suffix = suffix_source.suffix.lower()
+
+        if is_remote_url:
+            with urlopen(raw_path) as resp:
+                audio_bytes = resp.read()
+            source_for_torchaudio = NamedTemporaryFile(
+                suffix=suffix or ".audio",
+                delete=False,
+            )
+            source_for_torchaudio.write(audio_bytes)
+            source_for_torchaudio.flush()
+            source_for_torchaudio.close()
+        else:
+            source_for_torchaudio = None
+
+        path_for_wave = BytesIO(audio_bytes) if is_remote_url else str(local_path)
+
+        try:
+            if suffix in (".wav", ".wave"):
+                with wave.open(path_for_wave, "rb") as wf:
+                    sr = wf.getframerate()
+                    n_channels = wf.getnchannels()
+                    sample_width = wf.getsampwidth()
+                    raw = wf.readframes(wf.getnframes())
+                if sample_width == 1:
+                    audio = torch.frombuffer(bytearray(raw), dtype=torch.uint8).float()
+                    audio = (audio - 128.0) / 128.0
+                elif sample_width == 2:
+                    audio = torch.frombuffer(bytearray(raw), dtype=torch.int16).float() / 32768.0
+                elif sample_width == 4:
+                    audio = torch.frombuffer(bytearray(raw), dtype=torch.int32).float() / 2147483648.0
+                else:
+                    raise ValueError(f"Unsupported WAV sample width: {sample_width}")
+                if n_channels > 1:
+                    audio = audio.reshape(-1, n_channels).mean(dim=1)
+            else:
+                if torchaudio is None:
+                    raise ImportError("torchaudio is required for non-WAV audio paths.")
+                load_target = (
+                    source_for_torchaudio.name if is_remote_url else str(local_path)
+                )
+                waveform, sr = torchaudio.load(load_target)
+                if waveform.shape[0] > 1:
+                    waveform = waveform.mean(dim=0, keepdim=True)
+                audio = waveform.squeeze(0)
+        finally:
+            if source_for_torchaudio is not None:
+                try:
+                    Path(source_for_torchaudio.name).unlink(missing_ok=True)
+                except OSError:
+                    pass
+        if sr != target_sr:
+            if torchaudio is None:
+                raise ImportError("torchaudio is required for resampling.")
+            audio = torchaudio.functional.resample(audio.unsqueeze(0), sr, target_sr).squeeze(0)
+        return audio.float()
+
+    def _prepare_audio_batch(
+        self,
+        audio_items: list[Optional[Union[str, Path, torch.Tensor]]],
+        *,
+        target_sr: int,
+        device: torch.device,
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        tensor_list: list[torch.Tensor] = []
+        lens_list: list[int] = []
+        has_audio = False
+        for item in audio_items:
+            if item is None:
+                tensor_list.append(torch.zeros(1, dtype=torch.float32))
+                lens_list.append(0)
+                continue
+            has_audio = True
+            if isinstance(item, (str, Path)):
+                wav = self._load_audio_path(item, target_sr=target_sr)
+            elif isinstance(item, torch.Tensor):
+                wav = item.detach().float().cpu()
+                if wav.dim() == 2:
+                    wav = wav.mean(dim=0)
+                elif wav.dim() != 1:
+                    raise ValueError("Audio tensor must be 1D waveform or 2D [channels, length].")
+            else:
+                raise TypeError(f"Unsupported audio item type: {type(item)}")
+            if wav.numel() == 0:
+                wav = torch.zeros(1, dtype=torch.float32)
+                length = 0
+            else:
+                length = int(wav.numel())
+            tensor_list.append(wav)
+            lens_list.append(length)
+        if not has_audio:
+            return None, None
+        max_len = max(t.numel() for t in tensor_list)
+        padded = torch.zeros(len(tensor_list), max_len, dtype=torch.float32)
+        for i, wav in enumerate(tensor_list):
+            L = wav.numel()
+            if L > 0:
+                padded[i, :L] = wav
+        return padded.to(device), torch.tensor(lens_list, dtype=torch.long, device=device)
+
+    def encode(
+        self,
+        *,
+        text: Optional[Union[str, list[str]]] = None,
+        audio: Optional[Union[str, Path, torch.Tensor, list[Optional[Union[str, Path, torch.Tensor]]]]] = None,
+        task: str = "document",
+        instruction: Optional[str] = None,
+        normalize: bool = True,
+        device: Optional[Union[str, torch.device]] = None,
+    ) -> torch.Tensor:
+        """High-level embedding API. Accepts raw text/audio and returns embeddings."""
+        self.eval()
+        tokenizer = self._resolve_tokenizer()
+        device = torch.device(device) if device is not None else next(self.parameters()).device
+
+        text_items = self._to_list(text)
+        audio_items = self._to_list(audio)
+        batch_size = max(len(text_items), len(audio_items))
+        if batch_size == 0:
+            raise ValueError("At least one of text/audio must be provided.")
+
+        if len(text_items) == 0:
+            text_items = [None] * batch_size
+        elif len(text_items) == 1 and batch_size > 1:
+            text_items = text_items * batch_size
+        elif len(text_items) != batch_size:
+            raise ValueError("text and audio batch sizes must match (or be broadcastable length 1).")
+
+        if len(audio_items) == 0:
+            audio_items = [None] * batch_size
+        elif len(audio_items) == 1 and batch_size > 1:
+            audio_items = audio_items * batch_size
+        elif len(audio_items) != batch_size:
+            raise ValueError("text and audio batch sizes must match (or be broadcastable length 1).")
+
+        text_ids, text_lens = self._prepare_text_batch(
+            tokenizer,
+            texts=text_items,
+            task=task,
+            instruction=instruction,
+            device=device,
+        )
+        audio_tensor, audio_lens = self._prepare_audio_batch(
+            audio_items,
+            target_sr=TARGET_SR,
+            device=device,
+        )
+
+        with torch.inference_mode():
+            emb = self(
+                text_ids=text_ids,
+                text_lens=text_lens,
+                audio=audio_tensor,
+                audio_lens=audio_lens,
+            )
+        if normalize:
+            emb = torch.nn.functional.normalize(emb.float(), dim=-1)
+        return emb
+
+    @staticmethod
+    def _check_list_arg(name: str, value: Any, item_types: Tuple[type, ...]) -> None:
+        if value is None:
+            return
+        if not isinstance(value, list):
+            raise TypeError(
+                f"`{name}` must be a list, got {type(value).__name__}."
+            )
+        if len(value) == 0:
+            raise ValueError(f"`{name}` must be a non-empty list when provided.")
+        for i, item in enumerate(value):
+            if not isinstance(item, item_types):
+                allowed = ", ".join(t.__name__ for t in item_types)
+                raise TypeError(
+                    f"`{name}[{i}]` must be one of ({allowed}), "
+                    f"got {type(item).__name__}."
+                )
+
+    def encode_query(
+        self,
+        text: Optional[List[str]] = None,
+        audio: Optional[List[Union[str, Path, torch.Tensor]]] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Encode queries. Accepts text-only, audio-only, or both."""
+        self._check_list_arg("text", text, (str,))
+        self._check_list_arg("audio", audio, (str, Path, torch.Tensor))
+        if text is None and audio is None:
+            raise ValueError(
+                "encode_query requires at least one of `text` or `audio`."
+            )
+        if text is not None and audio is not None and len(text) != len(audio):
+            raise ValueError(
+                f"encode_query: `text` (len={len(text)}) and `audio` (len={len(audio)}) "
+                "must have the same length when both are provided."
+            )
+        return self.encode(text=text, audio=audio, task="query", **kwargs)
+
+    def encode_document(
+        self,
+        text: Optional[List[str]] = None,
+        audio: Optional[List[Union[str, Path, torch.Tensor]]] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Encode documents. Accepts exactly one of `text` or `audio`."""
+        self._check_list_arg("text", text, (str,))
+        self._check_list_arg("audio", audio, (str, Path, torch.Tensor))
+        if (text is None) == (audio is None):
+            raise ValueError(
+                "encode_document requires exactly one of `text` or `audio` "
+                "(not both, not neither)."
+            )
+        return self.encode(text=text, audio=audio, task="document", **kwargs)
+
+    def forward(self, text_ids, text_lens, audio=None, audio_lens=None):
+        device = text_ids.device
+        B = text_ids.size(0)
+        embed_dtype = self.model.embed_tokens.weight.dtype
+
+        if audio is not None:
+            audio_emb, audio_mask = self.dasheng(
+                audio,
+                audio_lens,
+            )
+            audio_emb, audio_mask = self.dasheng_down(audio_emb, audio_mask)
+            audio_lens = audio_mask.sum(dim=1)
+            audio_emb = self.dasheng_proj(audio_emb)
+        else:
+            audio_emb = None
+
+        text_emb = self.model.embed_tokens(text_ids)
+        audio_start_emb = self.model.embed_tokens(self.audio_start_token.clone())
+        audio_end_emb = self.model.embed_tokens(self.audio_end_token.clone())
+        eos_emb = self.model.embed_tokens(self.eos_token.clone())
+
+        input_embeds = []
+        attention_masks = []
+        last_indices = []
+
+        for i in range(B):
+            seq = [text_emb[i, : text_lens[i]]]
+            if audio_emb is not None and audio_lens[i] > 0:
+                seq.append(audio_start_emb)
+                seq.append(audio_emb[i, : audio_lens[i]])
+                seq.append(audio_end_emb)
+            seq.append(eos_emb)
+            seq = torch.cat(seq, dim=0)
+            input_embeds.append(seq)
+            attention_masks.append(torch.ones(seq.size(0), device=device))
+            last_indices.append(seq.size(0) - 1)
+
+        max_len = max(x.size(0) for x in input_embeds)
+        padded_embeds = torch.zeros(
+            B, max_len, self.hidden_size, device=device, dtype=embed_dtype
+        )
+        padded_mask = torch.zeros(B, max_len, device=device)
+        for i in range(B):
+            L = input_embeds[i].size(0)
+            padded_embeds[i, :L] = input_embeds[i]
+            padded_mask[i, :L] = attention_masks[i]
+
+        outputs = self.model(
+            inputs_embeds=padded_embeds,
+            attention_mask=padded_mask,
+        ).last_hidden_state
+
+        final_hidden = torch.stack([outputs[i, last_indices[i]] for i in range(B)])
+        out = self.siglip_head(final_hidden).squeeze(1)
+        return out
+