modeling_sentio.py

"""
Sentio: WeMM-Embedding-9B (+ a LoRA for video moment retrieval) with an audio branch, all
embedding into one 4096-d space.

  text / image / video / page -> WeMM-Embedding-9B with the vdr LoRA merged in
  audio                       -> Qwen2.5-Omni-3B audio tower + LoRA -> mean pool -> trained head
                                 (connector -> modality adapter fusion -> RMSNorm)

config.json points at the two base checkpoints, the vdr LoRA adapter (PEFT directory) and
trainable.pt (audio head + audio-tower LoRA). Preprocessing mirrors the MMEB-V3 evaluation that
produced the reported scores. Adapted from Wianmu-5B's modeling_wianmu.py.
"""
import json
import os
from typing import List, Optional, Sequence, Union

import numpy as np
import soundfile as sf
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
from safetensors import safe_open
from transformers import AutoConfig, AutoModel, AutoProcessor, WhisperFeatureExtractor
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder

AudioInput = Union[str, np.ndarray, torch.Tensor]

TARGET_MODALITIES = ("text", "audio", "video", "label")
WEMM_IMAGE_KWARGS = {"min_pixels": 64 * 32 * 32, "max_pixels": 8192 * 32 * 32}
WEMM_TEXT_MAX_LENGTH = 8192


# ============================================================================
# Trained audio head
# ============================================================================
class AudioConnector(nn.Module):
    def __init__(self, input_dim: int, output_dim: int):
        super().__init__()
        h1, h2 = 2 * input_dim, 4 * input_dim
        self.mlp = nn.Sequential(
            nn.Linear(input_dim, h1, bias=False), nn.LayerNorm(h1), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(h1, h2, bias=False), nn.LayerNorm(h2), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(h2, output_dim, bias=False),
        )

    def forward(self, x):
        return self.mlp(x)


class LightweightAdapter(nn.Module):
    def __init__(self, dim: int, reduction: int = 8):
        super().__init__()
        self.down = nn.Linear(dim, dim // reduction, bias=False)
        self.act = nn.GELU()
        self.drop = nn.Dropout(0.1)
        self.up = nn.Linear(dim // reduction, dim, bias=False)

    def forward(self, x):
        return x + self.up(self.drop(self.act(self.down(x))))


class ModalityAdapterFusion(nn.Module):
    """One adapter per target modality, mixed by a softmax gate computed from the input."""

    def __init__(self, dim: int):
        super().__init__()
        self.adapters = nn.ModuleDict({m: LightweightAdapter(dim) for m in TARGET_MODALITIES})
        self.gate = nn.Sequential(
            nn.Linear(dim, dim // 4, bias=False), nn.LayerNorm(dim // 4), nn.GELU(),
            nn.Linear(dim // 4, len(TARGET_MODALITIES)),
        )

    def forward(self, x):
        weights = F.softmax(self.gate(x), dim=-1)
        stacked = torch.stack([self.adapters[m](x) for m in TARGET_MODALITIES], dim=1)
        return (weights.unsqueeze(-1) * stacked).sum(dim=1)


class ProjectionHead(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.norm = nn.RMSNorm(dim, eps=1e-6)

    def forward(self, x):
        return self.norm(x)


class SentioAudioHead(nn.Module):
    """Pooled audio-tower output (B, audio_dim) -> L2-normalized (B, dim)."""

    def __init__(self, audio_dim: int, dim: int):
        super().__init__()
        self.audio_connector = AudioConnector(audio_dim, dim)
        self.adapter_fusion = ModalityAdapterFusion(dim)
        self.projection = ProjectionHead(dim)

    def forward(self, x):
        x = self.projection(self.adapter_fusion(self.audio_connector(x)))
        return F.normalize(x, p=2, dim=-1, eps=1e-6)


# ============================================================================
# Qwen2.5-Omni-3B audio tower (frozen base + LoRA)
# ============================================================================
def load_omni_audio_tower(qwen_omni_path: str, prefix: str = "thinker.audio_tower.") -> Qwen2_5OmniAudioEncoder:
    """Reads only the audio-tower tensors out of the full Qwen2.5-Omni checkpoint."""
    audio_config = AutoConfig.from_pretrained(qwen_omni_path).thinker_config.audio_config
    tower = Qwen2_5OmniAudioEncoder(audio_config)
    with open(os.path.join(qwen_omni_path, "model.safetensors.index.json")) as f:
        weight_map = json.load(f)["weight_map"]
    state = {}
    for shard in sorted({s for k, s in weight_map.items() if k.startswith(prefix)}):
        with safe_open(os.path.join(qwen_omni_path, shard), framework="pt") as f:
            for key in f.keys():
                if key.startswith(prefix):
                    state[key[len(prefix):]] = f.get_tensor(key)
    tower.load_state_dict(state, strict=True)
    return tower


def inject_audio_lora(tower: nn.Module, lora: dict) -> nn.Module:
    """Adds the (untrained) LoRA modules Sentio was trained with; weights come from trainable.pt."""
    from peft import LoraConfig, inject_adapter_in_model
    cfg = LoraConfig(r=lora["r"], lora_alpha=lora["alpha"], lora_dropout=lora["dropout"],
                     target_modules=list(lora["targets"]), bias="none")
    return inject_adapter_in_model(cfg, tower, adapter_name="default")


def load_trainable(trainable_path: str, audio_tower: nn.Module, audio_head: nn.Module):
    """trainable.pt holds the head blocks and the audio-tower LoRA tensors saved during training."""
    state = torch.load(trainable_path, map_location="cpu", weights_only=True)
    for block in ("audio_connector", "adapter_fusion", "projection"):
        missing, unexpected = getattr(audio_head, block).load_state_dict(state[block], strict=False)
        if missing or unexpected:
            raise RuntimeError(f"{block}: missing {missing}, unexpected {unexpected}")
    lora_state = {k[len("encoder."):] if k.startswith("encoder.") else k: v
                  for k, v in state["audio_encoder"].items() if "lora_" in k}
    missing, unexpected = audio_tower.load_state_dict(lora_state, strict=False)
    if unexpected or any("lora_" in k for k in missing):
        raise RuntimeError(f"audio-tower LoRA: unexpected {unexpected}, missing {[k for k in missing if 'lora_' in k]}")


@torch.no_grad()
def pool_audio_tower(tower: Qwen2_5OmniAudioEncoder, features: torch.Tensor, num_frames: int) -> torch.Tensor:
    """(128, T) mel features of one clip -> mean of the tower's output states over its valid frames."""
    if num_frames < 4:
        return features.new_zeros(tower.config.output_dim)
    features = features[:, :num_frames].contiguous()
    feature_lens = torch.tensor([num_frames], dtype=torch.long, device=features.device)
    aftercnn_lens = ((feature_lens - 1) // 2 + 1).clamp(min=1)
    out = tower(features, feature_lens=feature_lens, aftercnn_lens=aftercnn_lens)
    out = out[0] if isinstance(out, tuple) else out.last_hidden_state
    return out.mean(dim=0) if out.dim() == 2 else out.mean(dim=1).squeeze(0)


# ============================================================================
# Sentio
# ============================================================================
class Sentio(nn.Module):
    def __init__(self, config: dict, audio_tower, audio_head, feature_extractor, wemm, wemm_processor):
        super().__init__()
        self.config = config
        self.audio_tower = audio_tower
        self.audio_head = audio_head
        self.feature_extractor = feature_extractor
        self.wemm = wemm
        self.wemm_processor = wemm_processor

    @classmethod
    def from_pretrained(cls, path: str, device: Optional[str] = None, dtype: torch.dtype = torch.bfloat16,
                        **config_overrides):
        """config_overrides replace config.json entries, e.g. wemm_path=..., trainable_path=..."""
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        with open(os.path.join(path, "config.json")) as f:
            config = {**json.load(f), **config_overrides}

        audio_tower = inject_audio_lora(load_omni_audio_tower(config["qwen_omni_path"]), config["audio_lora"])
        audio_head = SentioAudioHead(config["audio_tower_dim"], config["embedding_dim"])
        load_trainable(config["trainable_path"], audio_tower, audio_head)
        feature_extractor = WhisperFeatureExtractor.from_pretrained(config["qwen_omni_path"])

        wemm = AutoModel.from_pretrained(
            config["wemm_path"], torch_dtype=dtype, trust_remote_code=True, attn_implementation="sdpa",
        )
        if config.get("wemm_adapter_path"):
            from peft import PeftModel
            wemm = PeftModel.from_pretrained(wemm, config["wemm_adapter_path"]).merge_and_unload()
        wemm_processor = AutoProcessor.from_pretrained(config["wemm_path"], trust_remote_code=True)

        model = cls(config, audio_tower, audio_head, feature_extractor, wemm, wemm_processor)
        return model.to(device=device, dtype=dtype).eval()

    @property
    def device(self) -> torch.device:
        return next(self.audio_head.parameters()).device

    def _autocast(self):
        return torch.autocast(device_type=self.device.type, dtype=torch.bfloat16,
                              enabled=self.device.type == "cuda")

    # ---------------------------------------------------------------- audio
    def load_audio(self, audio: AudioInput, sampling_rate: Optional[int] = None) -> np.ndarray:
        """Path or waveform -> mono float32 at 16 kHz, head-cropped to audio_max_seconds.
        Clips shorter than audio_min_seconds come back as 10 ms of silence (a zero embedding)."""
        target_sr = self.config["audio_sample_rate"]
        if isinstance(audio, str):
            wav, sampling_rate = sf.read(audio, dtype="float32")
            wav = torch.from_numpy(wav)
        else:
            if sampling_rate is None:
                raise ValueError("sampling_rate is required when passing a waveform")
            wav = torch.as_tensor(audio, dtype=torch.float32)
        if wav.dim() > 1:
            wav = wav.mean(dim=1)
        if sampling_rate != target_sr and wav.numel() > 0:
            wav = torchaudio.functional.resample(wav, sampling_rate, target_sr)
        if wav.numel() < int(self.config["audio_min_seconds"] * target_sr):
            return np.zeros(target_sr // 100, dtype=np.float32)
        return wav[: int(self.config["audio_max_seconds"] * target_sr)].numpy()

    @torch.no_grad()
    def encode_audio(self, audios: Sequence[AudioInput], sampling_rate: Optional[int] = None) -> torch.Tensor:
        """Audio file paths or 1-D/2-D (samples, channels) waveforms -> (N, 4096)."""
        pooled = []
        for audio in audios:
            wav = self.load_audio(audio, sampling_rate)
            feats = self.feature_extractor(wav, sampling_rate=self.config["audio_sample_rate"],
                                           return_attention_mask=True, return_tensors="pt")
            features = feats["input_features"][0].to(self.device, dtype=self.audio_tower.dtype)
            num_frames = min(int(feats["attention_mask"][0].sum()), features.shape[-1])
            with self._autocast():
                pooled.append(pool_audio_tower(self.audio_tower, features, num_frames))
        with self._autocast():
            return self.audio_head(torch.stack(pooled)).float()

    # ----------------------------------------------------------------- WeMM
    @torch.no_grad()
    def _encode_wemm(self, conversations: List[list], has_visual: bool) -> torch.Tensor:
        chat_texts = [
            self.wemm_processor.apply_chat_template(conv, tokenize=False, add_generation_prompt=False)
            for conv in conversations
        ]
        kwargs = dict(text=chat_texts, padding=True, return_tensors="pt")
        if has_visual:
            from qwen_vl_utils import process_vision_info
            images, _, _ = process_vision_info(conversations, image_patch_size=16,
                                               return_video_kwargs=True, return_video_metadata=True)
            kwargs.update(images=images, truncation=False)
        else:
            kwargs.update(truncation=True, max_length=WEMM_TEXT_MAX_LENGTH)
        inputs = self.wemm_processor(**kwargs).to(self.device)
        inner = getattr(self.wemm, "model", None)
        if inner is not None and hasattr(inner, "rope_deltas"):
            inner.rope_deltas = None
        with self._autocast():
            emb = self.wemm.embedding(**inputs)
        return F.normalize(emb.float(), p=2, dim=-1)

    def encode_text(self, texts: Sequence[str]) -> torch.Tensor:
        """Texts -> (N, 4096). Audio<->text was trained and evaluated on raw labels and
        captions (e.g. "dog", "chirping_birds"), with no instruction prefix."""
        conversations = [[{"role": "user", "content": [{"type": "text", "text": t.strip()}]}] for t in texts]
        return self._encode_wemm(conversations, has_visual=False)

    def encode_image(self, images: Sequence, texts: Optional[Sequence[str]] = None) -> torch.Tensor:
        """Image paths/URLs/PIL images, each with optional accompanying text -> (N, 4096).
        This is WeMM-Embedding-9B with the vdr LoRA merged; for video, call self.wemm as in WeMM's model card."""
        texts = texts or [""] * len(images)
        conversations = []
        for image, text in zip(images, texts):
            content = [{"type": "image", "image": image, **WEMM_IMAGE_KWARGS}]
            if text.strip():
                content.append({"type": "text", "text": text.strip()})
            conversations.append([{"role": "user", "content": content}])
        return self._encode_wemm(conversations, has_visual=True)