Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1,14 @@
+{{ bos_token }}
+{%- for message in messages -%}
+{%- if message['role'] == 'assistant' -%}
+    <|start_header_id|>{{ message['role'] }}<|end_header_id|>
+{% generation %}{{- message['content'] | trim }}<|eot_id|>{% endgeneration %}
+
+{% else %}
+<|start_header_id|>{{ message['role'] }}<|end_header_id|>
+{{ message['content'] | trim }}<|eot_id|>
+{% endif %}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+<|start_header_id|>assistant<|end_header_id|>
+{% endif -%}

config.json

@@ -0,0 +1,56 @@
+{
+  "_name_or_path": "",
+  "architectures": [
+    "SodaHierForCausalLM"
+  ],
+  "audio_id_lo": 128260,
+  "auto_map": {
+    "AutoConfig": "configuration_soda_hier.SodaHierConfig",
+    "AutoModelForCausalLM": "modeling_soda_hier.SodaHierForCausalLM"
+  },
+  "bos_token_id": 128000,
+  "chunk_size_feed_forward": 0,
+  "codebook_size": 2048,
+  "depth_head_dim": 128,
+  "depth_hidden_size": 512,
+  "depth_intermediate_size": 2048,
+  "depth_num_heads": 4,
+  "depth_num_kv_heads": 4,
+  "depth_num_layers": 6,
+  "dtype": null,
+  "eos_token_id": 128001,
+  "hidden_size": 768,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "intermediate_size": 3072,
+  "is_encoder_decoder": false,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "max_position_embeddings": 1024,
+  "model_type": "soda_hier",
+  "num_attention_heads": 6,
+  "num_codebooks": 8,
+  "num_hidden_layers": 8,
+  "num_key_value_heads": 6,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "problem_type": null,
+  "return_dict": true,
+  "rms_norm_eps": 1e-05,
+  "rope_scaling": {
+    "factor": 8.0,
+    "high_freq_factor": 4.0,
+    "low_freq_factor": 1.0,
+    "original_max_position_embeddings": 8192,
+    "rope_type": "llama3"
+  },
+  "rope_theta": 500000.0,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.12.1",
+  "unified_vocab_size": 130308,
+  "vocab_size": 144644
+}

configuration_soda_hier.py

@@ -0,0 +1,128 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""HF configuration 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 —
+never marin/levanter code.
+"""
+
+from transformers import PretrainedConfig
+from transformers.models.qwen3.configuration_qwen3 import Qwen3Config
+
+# Llama-3-style long-rope parameters shared by backbone and depth (the values
+# the runs were trained with). Written as legacy keys so transformers 4.x and
+# 5.x both read them.
+_DEFAULT_ROPE_THETA = 500000.0
+_DEFAULT_ROPE_SCALING = {
+    "rope_type": "llama3",
+    "factor": 8.0,
+    "low_freq_factor": 1.0,
+    "high_freq_factor": 4.0,
+    "original_max_position_embeddings": 8192,
+}
+
+
+class SodaHierConfig(PretrainedConfig):
+    """Backbone-over-steps + depth-over-codebooks factorization of Mimi audio.
+
+    One backbone position ("step") is a text/special token or one whole audio
+    frame (8 Mimi codebooks summed at the input). The unified head predicts
+    the next step's text/special/semantic id over ids 0..unified_vocab_size-1
+    (identical to the flat id space); a small depth transformer predicts the
+    7 acoustic codebooks within each frame.
+    """
+
+    model_type = "soda_hier"
+
+    def __init__(
+        self,
+        # id space
+        vocab_size: int = 144644,
+        unified_vocab_size: int = 130308,
+        num_codebooks: int = 8,
+        codebook_size: int = 2048,
+        audio_id_lo: int = 128260,
+        # backbone
+        hidden_size: int = 768,
+        intermediate_size: int = 3072,
+        num_hidden_layers: int = 8,
+        num_attention_heads: int = 6,
+        num_key_value_heads: int = 6,
+        max_position_embeddings: int = 1024,
+        # depth transformer
+        depth_hidden_size: int = 384,
+        depth_intermediate_size: int = 1536,
+        depth_num_layers: int = 4,
+        depth_num_heads: int = 3,
+        depth_num_kv_heads: int = 3,
+        depth_head_dim: int = 128,
+        # shared
+        rope_theta: float = _DEFAULT_ROPE_THETA,
+        rope_scaling: dict | None = None,
+        rms_norm_eps: float = 1e-5,
+        bos_token_id: int = 128000,
+        eos_token_id: int = 128001,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.unified_vocab_size = unified_vocab_size
+        self.num_codebooks = num_codebooks
+        self.codebook_size = codebook_size
+        self.audio_id_lo = audio_id_lo
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.max_position_embeddings = max_position_embeddings
+        self.depth_hidden_size = depth_hidden_size
+        self.depth_intermediate_size = depth_intermediate_size
+        self.depth_num_layers = depth_num_layers
+        self.depth_num_heads = depth_num_heads
+        self.depth_num_kv_heads = depth_num_kv_heads
+        self.depth_head_dim = depth_head_dim
+        self.rope_theta = rope_theta
+        self.rope_scaling = dict(rope_scaling) if rope_scaling else dict(_DEFAULT_ROPE_SCALING)
+        self.rms_norm_eps = rms_norm_eps
+        kwargs.setdefault("tie_word_embeddings", False)
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+
+    def backbone_config(self) -> Qwen3Config:
+        return Qwen3Config(
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            intermediate_size=self.intermediate_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            num_key_value_heads=self.num_key_value_heads,
+            head_dim=self.hidden_size // self.num_attention_heads,
+            max_position_embeddings=self.max_position_embeddings,
+            rope_theta=self.rope_theta,
+            rope_scaling=dict(self.rope_scaling),
+            rms_norm_eps=self.rms_norm_eps,
+            attention_bias=False,
+            tie_word_embeddings=False,
+            use_sliding_window=False,
+            use_cache=True,
+        )
+
+    def depth_config(self) -> Qwen3Config:
+        return Qwen3Config(
+            vocab_size=self.num_codebooks * self.codebook_size,
+            hidden_size=self.depth_hidden_size,
+            intermediate_size=self.depth_intermediate_size,
+            num_hidden_layers=self.depth_num_layers,
+            num_attention_heads=self.depth_num_heads,
+            num_key_value_heads=self.depth_num_kv_heads,
+            head_dim=self.depth_head_dim,
+            max_position_embeddings=self.num_codebooks,
+            rope_theta=self.rope_theta,
+            rope_scaling=dict(self.rope_scaling),
+            rms_norm_eps=self.rms_norm_eps,
+            attention_bias=False,
+            tie_word_embeddings=False,
+            use_sliding_window=False,
+            use_cache=False,
+        )

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c2218c5170b7d449b81f5a1352bff6165353dc1e0b389fd56d1ae6940960713
+size 1311904880

modeling_soda_hier.py

@@ -0,0 +1,261 @@
+# 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)

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7c2f8c6c93d178640aabccd2c24b0a45fd2b0e31eb2e550ca6d76767e9dc49c
+size 20167966

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|begin_of_text|>",
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "<|end_of_text|>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 131072,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}