| """组合模型 9B 版 (导出包 vendored): Qwen3-Omni AuT(2048) + connector(2048→4096 投影) + Qwen3.5-9B (early fusion)。 |
| |
| 与 audio-trans work/audio_model_9b.py 逻辑一致, 路径全部参数化 |
| (llm_path / omni_path 由调用方传入, 不再硬编码 $AT)。 |
| |
| 架构: wav → WhisperFeatureExtractor(128 mel) → AuT(2048) → connector(2048→4096) → |
| 替换 <|audio_pad|> 位。connector = 列正交 W × gain (EMBED_RMS_9B/AUT_RMS/√(2048/4096), |
| 输出 per-elem RMS ≈ 9B embed 行 RMS 0.0137); log_alpha 已删 (2026-09-11 极简口径)。 |
| |
| 类名保持 Qwen35AudioForCausalLM 与 2B 包一致, pipeline 无需区分 2B/9B。 |
| """ |
| import json |
| import os |
|
|
| import torch |
| import torch.nn as nn |
|
|
| AUDIO_PAD = 248076 |
| AUDIO_DIM = 2048 |
| EMBED_RMS_9B = 0.0137 |
| AUT_RMS = 0.0143 |
|
|
|
|
| class AudioConnector9B(nn.Module): |
| """h' = W·h; W: 2048→4096 列正交 × gain。""" |
|
|
| def __init__(self, in_dim=AUDIO_DIM, out_dim=4096, |
| gain=EMBED_RMS_9B / AUT_RMS / (AUDIO_DIM / 4096) ** 0.5): |
| super().__init__() |
| self.proj = nn.Linear(in_dim, out_dim, bias=False) |
| nn.init.orthogonal_(self.proj.weight) |
| self.proj.weight.data *= gain |
|
|
| def forward(self, h): |
| return self.proj(h) |
|
|
|
|
| def load_audio_tower(omni_path, dtype=torch.bfloat16): |
| """从 Qwen3-Omni shard 1 加载 AuT(525 keys)。omni_path: 含 config.json + |
| model-00001-of-00015.safetensors 的目录 (导出包 stlm_omni/)。""" |
| from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( |
| Qwen3OmniMoeAudioEncoder, Qwen3OmniMoeAudioEncoderConfig) |
| from safetensors.torch import load_file |
| cfg = json.load(open(os.path.join(omni_path, 'config.json'))) |
| audio_cfg = Qwen3OmniMoeAudioEncoderConfig(**{k: v for k, v in |
| cfg['thinker_config']['audio_config'].items() |
| if not k.startswith('_')}) |
| enc = Qwen3OmniMoeAudioEncoder(audio_cfg).eval() |
| sd = load_file(os.path.join(omni_path, 'model-00001-of-00015.safetensors')) |
| sd = {k.replace('thinker.audio_tower.', ''): v for k, v in sd.items() |
| if k.startswith('thinker.audio_tower.')} |
| missing, unexpected = enc.load_state_dict(sd, strict=False) |
| assert not missing and not unexpected, f'missing={missing[:3]} unexpected={unexpected[:3]}' |
| return enc.to(dtype) |
|
|
|
|
| def load_feature_extractor(): |
| from transformers import WhisperFeatureExtractor |
| fe = WhisperFeatureExtractor(feature_size=128, hop_length=160, n_fft=400, |
| sampling_rate=16000, padding_value=0.0, |
| return_attention_mask=True) |
| |
| fe.n_samples = 4800000 |
| fe.nb_max_frames = 30000 |
| return fe |
|
|
|
|
| class Qwen35AudioForCausalLM(nn.Module): |
| """AuT + connector + Qwen3.5-9B, audio embedding 原位替换 <|audio_pad|>。 |
| |
| (类名与 2B 包相同; 9B 差异: LLM 4096h/32L untied, connector 为 2048→4096 投影) |
| """ |
|
|
| def __init__(self, llm_path, omni_path, dtype=torch.bfloat16, |
| freeze_llm=False, freeze_audio=False): |
| super().__init__() |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| self.tokenizer = AutoTokenizer.from_pretrained(llm_path) |
| self.llm = AutoModelForCausalLM.from_pretrained(llm_path, dtype=dtype) |
| self.audio_tower = load_audio_tower(omni_path, dtype) |
| self.fe = load_feature_extractor() |
| self.connector = AudioConnector9B(out_dim=self.llm.config.hidden_size).to(dtype) |
| assert self.llm.config.hidden_size == 4096 |
| assert self.llm.config.tie_word_embeddings is False |
| if freeze_llm: |
| for p in self.llm.parameters(): |
| p.requires_grad_(False) |
| if freeze_audio: |
| for p in self.audio_tower.parameters(): |
| p.requires_grad_(False) |
|
|
| @property |
| def device(self): |
| return next(self.llm.parameters()).device |
|
|
| def encode_audio(self, wav_paths): |
| """wav paths → list of (T_i, 4096) connector 后 audio embeddings + 各段帧数。""" |
| import librosa |
| embs = [] |
| dt = next(self.audio_tower.parameters()).dtype |
| for p in wav_paths: |
| wav, _ = librosa.load(p, sr=16000) |
| f = self.fe(wav, sampling_rate=16000, return_tensors='pt', |
| return_attention_mask=True) |
| T = int(f.attention_mask.sum(-1)[0]) |
| feats = f.input_features[0][:, :T] |
| out = self.audio_tower(input_features=feats.to(self.device, dt), |
| feature_lens=torch.tensor([T], device=self.device)) |
| h = out.last_hidden_state if hasattr(out, 'last_hidden_state') else out[0] |
| embs.append(self.connector(h.to(self.device, dt))) |
| return embs |
|
|
| def forward(self, input_ids, attention_mask=None, labels=None, |
| audio_embs=None, **kwargs): |
| embed = self.llm.get_input_embeddings() |
| inputs_embeds = embed(input_ids) |
| if audio_embs is not None: |
| pad_mask = input_ids == AUDIO_PAD |
| n_pad = int(pad_mask.sum()) |
| n_audio = sum(e.shape[0] for e in audio_embs) |
| assert n_pad == n_audio, f'pad tokens {n_pad} != audio frames {n_audio}' |
| inputs_embeds = inputs_embeds.clone() |
| inputs_embeds[pad_mask] = torch.cat(audio_embs).to(inputs_embeds.dtype) |
| return self.llm(inputs_embeds=inputs_embeds, attention_mask=attention_mask, |
| labels=labels, **kwargs) |
|
|
