modeling_alm2vec.py

"""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