# Copyright The Marin Authors # SPDX-License-Identifier: Apache-2.0 """HF modeling for the SODA hierarchical (backbone + depth) audio LM. This file is copied verbatim into every exported checkpoint directory and loaded via trust_remote_code, so it may import only torch/transformers. The model consumes and produces the same flat frame-interleaved token stream as the flattened arm (8 audio ids per Mimi frame: semantic then 7 acoustics), so it is a drop-in for likelihood evals that gather next-token log-probs from ``model(ids).logits`` and for ``model.generate``: - Internally, positions are grouped into backbone "steps": one text/special token, or one whole frame (its 8 codebook embeddings summed). - ``logits[t]`` is the model's true factorized conditional for token ``t+1``: the 130,308-way unified head (text/special/semantic ids — identical to flat ids 0..130307) when ``t+1`` starts a step, or the 2,048-way depth head for codebook k mapped into its flat id block when ``t+1`` is acoustic. All other vocabulary entries are -inf, so ``log_softmax`` reproduces the factorized log-prob exactly. - The depth factorization matches training exactly: codebook k is predicted from the backbone hidden of the PREVIOUS step plus codebooks 0..k-2 of the same frame (the immediately preceding codebook is not in the conditioning set for k >= 2 — a property of the trained shifted-prefix scheme, replicated verbatim). """ import torch import torch.nn as nn from transformers import PreTrainedModel from transformers.cache_utils import DynamicCache from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.qwen3.modeling_qwen3 import Qwen3Model from .configuration_soda_hier import SodaHierConfig def _group_steps(ids: torch.Tensor, audio_id_lo: int, num_codebooks: int): """Map a flat interleaved stream to steps. Returns (steps, step_of, slot_of): ``steps`` is (S, num_codebooks) with -1 padding on non-frame steps; ``step_of[t]``/``slot_of[t]`` locate position t. """ T = ids.shape[0] step_of = torch.empty(T, dtype=torch.long) slot_of = torch.empty(T, dtype=torch.long) rows: list[list[int]] = [] run = 0 for t in range(T): tok = int(ids[t]) if tok < audio_id_lo: rows.append([tok] + [-1] * (num_codebooks - 1)) run = 0 else: if run % num_codebooks == 0: rows.append([-1] * num_codebooks) rows[-1][run % num_codebooks] = tok slot_of[t] = run % num_codebooks step_of[t] = len(rows) - 1 run += 1 continue step_of[t] = len(rows) - 1 slot_of[t] = 0 steps = torch.tensor(rows, dtype=torch.long) return steps, step_of, slot_of class SodaHierForCausalLM(PreTrainedModel): config_class = SodaHierConfig _no_split_modules = ["Qwen3DecoderLayer"] main_input_name = "input_ids" _tied_weights_keys = [] def __init__(self, config: SodaHierConfig): super().__init__(config) self.backbone = Qwen3Model(config.backbone_config()) self.depth = Qwen3Model(config.depth_config()) e, e_d = config.hidden_size, config.depth_hidden_size self.unified_head = nn.Linear(e, config.unified_vocab_size, bias=False) self.bd_proj = nn.Linear(e, e_d, bias=False) self.acoustic_heads = nn.ModuleList( nn.Linear(e_d, config.codebook_size, bias=False) for _ in range(config.num_codebooks - 1) ) self.post_init() # ------------------------------------------------------------------ core def _embed_steps(self, steps: torch.Tensor) -> torch.Tensor: """(S, num_codebooks) step ids (-1 = empty slot) -> (S, E) summed embeddings.""" valid = steps >= 0 emb = self.backbone.embed_tokens(steps.clamp(min=0)) return (emb * valid.unsqueeze(-1)).sum(dim=1) def _depth_hidden_for_frames(self, cond: torch.Tensor, frames: torch.Tensor) -> torch.Tensor: """Teacher-forced depth pass. cond (F, E_d); frames (F, 8) LM ids -> (F, 8, E_d).""" cfg = self.config audio_idx = (frames - cfg.audio_id_lo).clamp(0, cfg.num_codebooks * cfg.codebook_size - 1) prefix = self.depth.embed_tokens(audio_idx) # (F, 8, E_d) shifted = torch.roll(prefix, 1, dims=1) shifted[:, 0] = 0.0 x = cond.unsqueeze(1) + shifted return self.depth(inputs_embeds=x).last_hidden_state def _forward_one(self, ids: torch.Tensor) -> torch.Tensor: """One unpadded row (T,) -> logits (T, vocab).""" cfg = self.config dev = ids.device steps, step_of, slot_of = _group_steps(ids.cpu(), cfg.audio_id_lo, cfg.num_codebooks) steps, step_of, slot_of = steps.to(dev), step_of.to(dev), slot_of.to(dev) T = ids.shape[0] emb = self._embed_steps(steps) # (S, E) h = self.backbone(inputs_embeds=emb.unsqueeze(0)).last_hidden_state[0] # (S, E) u = self.unified_head(h) # (S, unified) is_audio = ids >= cfg.audio_id_lo is_frame_step = steps[:, 1] >= 0 # frame steps have slot-1 filled frame_steps = torch.nonzero(is_frame_step, as_tuple=False)[:, 0] # depth conditions on the hidden of the step BEFORE the frame cond_frames = frame_steps[frame_steps >= 1] d = None frame_row = torch.full((steps.shape[0],), -1, dtype=torch.long, device=dev) if len(cond_frames): d = self._depth_hidden_for_frames(self.bd_proj(h[cond_frames - 1]), steps[cond_frames]) frame_row[cond_frames] = torch.arange(len(cond_frames), device=dev) logits = torch.full((T, cfg.vocab_size), float("-inf"), dtype=h.dtype, device=dev) # positions whose NEXT token starts a step: non-audio positions and slot-7 audio positions primary = (~is_audio) | (slot_of == cfg.num_codebooks - 1) logits[primary, : cfg.unified_vocab_size] = u[step_of[primary]] # positions whose next token is acoustic codebook k+1 of the SAME frame for k in range(cfg.num_codebooks - 1): pos = torch.nonzero(is_audio & (slot_of == k), as_tuple=False)[:, 0] if not len(pos): continue rows = frame_row[step_of[pos]] ok = rows >= 0 lo = cfg.audio_id_lo + (k + 1) * cfg.codebook_size if ok.any(): logits[pos[ok], lo : lo + cfg.codebook_size] = self.acoustic_heads[k](d[rows[ok], k]) if (~ok).any(): # frame at step 0: the factorization defines no conditional; keep finite logits[pos[~ok]] = 0.0 return logits def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs) -> CausalLMOutputWithPast: if input_ids is None: raise ValueError("SodaHierForCausalLM.forward requires input_ids") rows = [] for b in range(input_ids.shape[0]): ids = input_ids[b] if attention_mask is not None: length = int(attention_mask[b].sum()) logits_b = torch.zeros( (ids.shape[0], self.config.vocab_size), dtype=torch.float32, device=ids.device ) logits_b[:length] = self._forward_one(ids[:length]) rows.append(logits_b) else: rows.append(self._forward_one(ids)) logits = torch.stack(rows) loss = None if labels is not None: shift_logits = logits[:, :-1].reshape(-1, self.config.vocab_size) shift_labels = labels[:, 1:].reshape(-1) loss = nn.functional.cross_entropy(shift_logits, shift_labels, ignore_index=-100) return CausalLMOutputWithPast(loss=loss, logits=logits) # -------------------------------------------------------------- sampling @staticmethod def _sample(logits: torch.Tensor, do_sample: bool, temperature: float, top_p: float) -> int: if not do_sample or temperature <= 0: return int(logits.argmax()) probs = torch.softmax(logits / temperature, dim=-1) if top_p is not None and top_p < 1.0: sorted_probs, sorted_idx = probs.sort(descending=True) cum = sorted_probs.cumsum(-1) # drop tokens entirely beyond the nucleus; the crossing token stays remove = (cum - sorted_probs) > top_p sorted_probs[remove] = 0.0 sorted_probs /= sorted_probs.sum() return int(sorted_idx[torch.multinomial(sorted_probs, 1)]) return int(torch.multinomial(probs, 1)) @torch.no_grad() def generate( self, input_ids=None, attention_mask=None, max_new_tokens: int = 200, do_sample: bool = False, temperature: float = 1.0, top_p: float = 1.0, eos_token_id=None, pad_token_id=None, **ignored, ) -> torch.Tensor: """Two-stage autoregressive decode over the flat interleaved stream. The backbone advances once per step with a KV cache; whenever the unified head emits a semantic token, the depth transformer fills in the frame's 7 acoustic codebooks before the backbone moves on. Only whole frames are emitted: if fewer than 8 tokens of budget remain when a frame starts, generation stops early instead. """ cfg = self.config if input_ids.shape[0] != 1: raise NotImplementedError("SodaHierForCausalLM.generate supports batch size 1") dev = input_ids.device eos = set() if eos_token_id is not None: eos = {eos_token_id} if isinstance(eos_token_id, int) else set(eos_token_id) ids = input_ids[0] steps, _, _ = _group_steps(ids.cpu(), cfg.audio_id_lo, cfg.num_codebooks) steps = steps.to(dev) emb = self._embed_steps(steps) cache = DynamicCache() out = self.backbone(inputs_embeds=emb.unsqueeze(0), past_key_values=cache, use_cache=True) cache = out.past_key_values h_last = out.last_hidden_state[0, -1] step_pos = steps.shape[0] generated: list[int] = [] while len(generated) < max_new_tokens: primary = self._sample(self.unified_head(h_last), do_sample, temperature, top_p) if primary < cfg.audio_id_lo: # text or special: a one-token step generated.append(primary) next_emb = self.backbone.embed_tokens(torch.tensor([primary], device=dev))[0] if primary in eos: break else: # semantic token: emit a whole frame via the depth transformer if max_new_tokens - len(generated) < cfg.num_codebooks: break cond = self.bd_proj(h_last) frame = [primary] xs = [cond] for j in range(cfg.num_codebooks - 1): d = self.depth(inputs_embeds=torch.stack(xs).unsqueeze(0)).last_hidden_state[0, -1] idx = self._sample(self.acoustic_heads[j](d), do_sample, temperature, top_p) frame.append(cfg.audio_id_lo + (j + 1) * cfg.codebook_size + idx) xs.append(cond + self.depth.embed_tokens(torch.tensor(frame[-1] - cfg.audio_id_lo, device=dev))) generated.extend(frame) frame_t = torch.tensor(frame, device=dev) next_emb = self.backbone.embed_tokens(frame_t).sum(dim=0) if eos & set(frame): break out = self.backbone( inputs_embeds=next_emb.view(1, 1, -1), past_key_values=cache, use_cache=True, position_ids=torch.tensor([[step_pos]], device=dev), ) cache = out.past_key_values h_last = out.last_hidden_state[0, -1] step_pos += 1 return torch.cat([ids, torch.tensor(generated, dtype=ids.dtype, device=dev)]).unsqueeze(0)