""" Wren-ASR model — a transformers-compatible wrapper over Qwen2.5-0.5B + Mimi input embedding tables. Designed for use with `AutoModel.from_pretrained(..., trust_remote_code=True)`. Self-contained: no imports from a `src/` folder. Sequence layout: [ | sum_q embed_q(codes[q, t]) for t in 0..T-1 | | text... | ] Audio positions feed a single summed-codebook embedding per real frame (no delay pattern). Text-token prediction uses the LLM's existing `lm_head`; no new output heads are added. """ import math from typing import Optional import torch import torch.nn as nn from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel try: from .configuration_wren_asr import WrenASRConfig # package context (HF trust_remote_code) except ImportError: import importlib WrenASRConfig = importlib.import_module("configuration_wren_asr").WrenASRConfig class WrenForASR(PreTrainedModel): config_class = WrenASRConfig base_model_prefix = "wren_asr" # Qwen2.5 ties lm_head to embed_tokens. safetensors deduplicates the tied # tensor at save time; this tells HF to re-tie it on load. Without this, # lm_head is left at random init and generation collapses to repeated token 0 # ("!" in Qwen's BPE). _tied_weights_keys = ["llm.lm_head.weight"] # Forward HF's tying machinery to the inner LLM. PreTrainedModel.tie_weights # calls self.get_output_embeddings() / get_input_embeddings(); without these # overrides those raise NotImplementedError and lm_head is never re-tied # after load — bug surfaces as generate() emitting only token 0 ("!"). def get_input_embeddings(self): return self.llm.get_input_embeddings() def set_input_embeddings(self, new_embeddings): self.llm.set_input_embeddings(new_embeddings) def get_output_embeddings(self): return self.llm.get_output_embeddings() def set_output_embeddings(self, new_embeddings): self.llm.set_output_embeddings(new_embeddings) def tie_weights(self): # Delegate fully to the inner Qwen2ForCausalLM, which knows how to tie # its own lm_head ↔ embed_tokens. Doing super().tie_weights() afterwards # is harmless but redundant since the inner tie covers it. if hasattr(self.llm, "tie_weights"): self.llm.tie_weights() super().tie_weights() def __init__(self, config: WrenASRConfig): super().__init__(config) self.k = config.k_codebooks # Build backbone from its config only. Pretrained backbone weights are # already in our state_dict; no need to re-download. llm_cfg = AutoConfig.from_pretrained(config.llm_name) llm_cfg.vocab_size = config.vocab_size self.llm = AutoModelForCausalLM.from_config(llm_cfg) hidden = self.llm.config.hidden_size # k input embedding tables (codes are inputs only — no PAD row needed). self.audio_embeds = nn.ModuleList([ nn.Embedding(config.codebook_size, hidden) for _ in range(self.k) ]) self.embed_scale = 1.0 / math.sqrt(self.k) self._mimi = None # lazy-loaded on first use # --- Mimi codec (lazy-loaded encoder for raw-waveform input) --- @property def mimi(self): if self._mimi is None: from transformers import MimiModel self._mimi = MimiModel.from_pretrained(self.config.mimi_model_name).to(self.device) self._mimi.eval() for p in self._mimi.parameters(): p.requires_grad_(False) return self._mimi @torch.no_grad() def encode_audio( self, waveform: torch.Tensor, src_sample_rate: int = 24000, ) -> torch.LongTensor: """Encode a waveform to Mimi codes [k, n_frames].""" if waveform.dim() == 1: waveform = waveform.unsqueeze(0) if src_sample_rate != self.config.sampling_rate: import torchaudio.transforms as T waveform = T.Resample(src_sample_rate, self.config.sampling_rate)(waveform) x = waveform.unsqueeze(0).to(self.device) out = self.mimi.encode(x, num_quantizers=self.k) return out.audio_codes[0].cpu() # [k, n_frames] # --- Generation --- @torch.no_grad() def generate( self, audio_codes: torch.LongTensor, # [k, T] or [B, k, T] max_new_tokens: int = 200, do_sample: bool = False, temperature: float = 1.0, top_k: int = 50, top_p: float = 1.0, eos_token_id: Optional[int] = None, pad_token_id: Optional[int] = None, **kwargs, ) -> torch.LongTensor: """Transcribe Mimi codes to text-token IDs. Returns the generated token IDs (without the audio prefix). Decode them with your tokenizer to get text — typically: ids = model.generate(audio_codes=codes) text = tokenizer.decode(ids[0], skip_special_tokens=True) """ device = next(self.parameters()).device self.eval() audio_codes = audio_codes.to(device) if audio_codes.dim() == 2: audio_codes = audio_codes.unsqueeze(0) # [1, k, T] B, k, T = audio_codes.shape assert k == self.k, f"expected k={self.k}, got {k}" embed_tokens = self.llm.get_input_embeddings() llm_dtype = next(self.llm.parameters()).dtype start_emb = embed_tokens(torch.tensor([[self.config.audio_start_id]], device=device)) end_emb = embed_tokens(torch.tensor([[self.config.audio_end_id]], device=device)) clamped = audio_codes.clamp(0, self.config.codebook_size - 1) audio_sum = self.audio_embeds[0](clamped[:, 0, :]) for q in range(1, self.k): audio_sum = audio_sum + self.audio_embeds[q](clamped[:, q, :]) audio_sum = audio_sum * self.embed_scale prompt_embeds = torch.cat([ start_emb.expand(B, -1, -1), audio_sum, end_emb.expand(B, -1, -1), ], dim=1).to(llm_dtype) # IMPORTANT: pass attention_mask explicitly. HF's generate cannot infer # it when pad_token_id == eos_token_id (Qwen's case), and can't infer it # at all from `inputs_embeds`. Without this, transformers ≥4.45 frequently # collapses to repeated token 0 ("!") with `inputs_embeds`-only input. attn_mask = torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=device) eos = eos_token_id if eos_token_id is not None else self.config.eos_token_id gen_ids = self.llm.generate( inputs_embeds = prompt_embeds, attention_mask = attn_mask, max_new_tokens = max_new_tokens, do_sample = do_sample, temperature = temperature if do_sample else 1.0, top_k = top_k if do_sample else 0, top_p = top_p if do_sample else 1.0, eos_token_id = eos, pad_token_id = pad_token_id if pad_token_id is not None else eos, ) # When called with `inputs_embeds`, HF generate returns ONLY the # generated ids (the prompt has no token ids to echo back). return gen_ids