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