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