refactor: rename TinyAya to RumikOSS

Renames the model and config classes, the modules that hold them, and
model_type. Weight tensor names are unchanged, so the safetensors shards
and their index are untouched.

Also folds in three behaviour changes to the model class:

- _sample is renamed to _constrained_sample. It was shadowing
GenerationMixin._sample, so every call to model.generate() failed with
a TypeError on logits_processor.
- allowed_ids is now optional and derived from the config, removing the
need to regex-scan the tokenizer vocabulary for unit tokens.
- audio_tokens_to_codes() regroups generated ids into [1, 8, T] codec
frames without pulling in the codec.

config.json gains the audio vocabulary layout (first_unit_id 261008,
last_unit_id 277391, text_start_token_id 277392, frame_rate_hz 12.5),
verified against this repository's own tokenizer.

---

Mirrors the changes made to rumik-oss 1, so the two repositories
present the same API.

Renames the model and config classes to RumikOSS, moves the audio token
layout into `config.json`, unblocks `model.generate()`, adds a chat
template matching the trained prompt layout, and adds LICENSE.

Weight files are not touched: tensor names are unchanged, so the
safetensors shards and their index are identical.

Verified against the published weights on an RTX PRO 6000:

- loads on transformers 5.16.1 as RumikOSSForCausalLM, vocab 277,395
- the chat template is token-exact vs the prompt `inference.py` built
- audio generates and decodes for English and Hindi
- `inference.py` writes a wav and rejects an unknown speaker

config.json

@@ -1,21 +1,54 @@
 {
-  "vocab_size": 277395,
-  "max_position_embeddings": 500000,
-  "hidden_size": 2048,
-  "logit_scale": 1.0,
-  "intermediate_size": 11008,
-  "num_hidden_layers": 36,
-  "num_attention_heads": 16,
-  "num_key_value_heads": 4,
+  "_sliding_window_pattern": 4,
+  "add_cross_attention": false,
+  "architectures": [
+    "RumikOSSForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "audio_end_token_id": 277394,
+  "audio_start_token_id": 277393,
+  "auto_map": {
+    "AutoConfig": "configuration_rumik_oss.RumikOSSConfig",
+    "AutoModelForCausalLM": "modeling_rumik_oss.RumikOSSForCausalLM"
+  },
+  "bad_words_ids": null,
+  "begin_suppress_tokens": null,
+  "bos_token_id": 2,
+  "chunk_size_feed_forward": 0,
+  "codebook_size": 2048,
+  "cross_attention_hidden_size": null,
+  "decoder_start_token_id": null,
+  "diversity_penalty": 0.0,
+  "do_sample": false,
+  "dtype": "bfloat16",
+  "early_stopping": false,
+  "encoder_no_repeat_ngram_size": 0,
+  "eos_token_id": 3,
+  "exponential_decay_length_penalty": null,
+  "finetuning_task": null,
+  "first_unit_id": 261008,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "frame_rate_hz": 12.5,
+  "head_dim": 128,
   "hidden_act": "silu",
+  "hidden_size": 2048,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
   "initializer_range": 0.02,
+  "intermediate_size": 11008,
+  "is_decoder": false,
+  "is_encoder_decoder": false,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "last_unit_id": 277391,
   "layer_norm_eps": 1e-05,
-  "use_cache": true,
-  "rope_theta": 50000,
-  "rope_scaling": null,
-  "attention_bias": false,
-  "attention_dropout": 0.0,
-  "sliding_window": 4096,
+  "layer_switch": 4,
   "layer_types": [
     "sliding_attention",
     "sliding_attention",
@@ -54,94 +87,64 @@
     "sliding_attention",
     "full_attention"
   ],
-  "head_dim": 128,
-  "return_dict": true,
-  "output_hidden_states": false,
-  "torchscript": false,
-  "dtype": "bfloat16",
-  "pruned_heads": {},
-  "tie_word_embeddings": true,
-  "chunk_size_feed_forward": 0,
-  "is_encoder_decoder": false,
-  "is_decoder": false,
-  "cross_attention_hidden_size": null,
-  "add_cross_attention": false,
-  "tie_encoder_decoder": false,
-  "architectures": [
-    "TinyAyaForCausalLM"
-  ],
-  "finetuning_task": null,
-  "id2label": {
-    "0": "LABEL_0",
-    "1": "LABEL_1"
-  },
-  "label2id": {
-    "LABEL_0": 0,
-    "LABEL_1": 1
-  },
-  "task_specific_params": null,
-  "problem_type": null,
-  "tokenizer_class": null,
-  "prefix": null,
-  "bos_token_id": 2,
-  "pad_token_id": 0,
-  "eos_token_id": 3,
-  "sep_token_id": null,
-  "decoder_start_token_id": null,
+  "length_penalty": 1.0,
+  "logit_scale": 1.0,
   "max_length": 20,
+  "max_position_embeddings": 500000,
   "min_length": 0,
-  "do_sample": false,
-  "early_stopping": false,
-  "num_beams": 1,
-  "temperature": 1.0,
-  "top_k": 50,
-  "top_p": 1.0,
-  "typical_p": 1.0,
-  "repetition_penalty": 1.0,
-  "length_penalty": 1.0,
+  "model_type": "rumik_oss",
   "no_repeat_ngram_size": 0,
-  "encoder_no_repeat_ngram_size": 0,
-  "bad_words_ids": null,
-  "num_return_sequences": 1,
-  "output_scores": false,
-  "return_dict_in_generate": false,
-  "forced_bos_token_id": null,
-  "forced_eos_token_id": null,
-  "remove_invalid_values": false,
-  "exponential_decay_length_penalty": null,
-  "suppress_tokens": null,
-  "begin_suppress_tokens": null,
+  "num_attention_heads": 16,
   "num_beam_groups": 1,
-  "diversity_penalty": 0.0,
-  "_sliding_window_pattern": 4,
-  "cache_implementation": "hybrid",
-  "layer_switch": 4,
-  "model_type": "tinyaya_mimi",
+  "num_beams": 1,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 4,
+  "num_quantizers": 8,
+  "num_return_sequences": 1,
   "order_of_interleaved_layers": "local_attn_first",
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "output_scores": false,
+  "pad_token_id": 0,
   "position_embedding_type": "rope_gptj",
+  "prefix": null,
+  "problem_type": null,
+  "pruned_heads": {},
+  "remove_invalid_values": false,
+  "repetition_penalty": 1.0,
+  "return_dict": true,
+  "return_dict_in_generate": false,
+  "rope_scaling": null,
+  "rope_theta": 50000,
   "rotary_pct": 1.0,
+  "sep_token_id": null,
+  "sliding_window": 4096,
   "sliding_window_pattern": 4,
-  "use_embedding_sharing": true,
-  "use_gated_activation": true,
-  "use_parallel_block": true,
-  "use_parallel_embedding": false,
-  "use_qk_norm": false,
-  "tf_legacy_loss": false,
-  "use_bfloat16": false,
-  "output_attentions": false,
-  "num_quantizers": 8,
-  "codebook_size": 2048,
-  "audio_start_token_id": 277393,
-  "audio_end_token_id": 277394,
   "speakers": [
     "Ira",
     "Aisha",
     "Siya",
     "Zoya"
   ],
-  "auto_map": {
-    "AutoConfig": "configuration_tinyaya.TinyAyaConfig",
-    "AutoModelForCausalLM": "modeling_tinyaya.TinyAyaForCausalLM"
-  },
-  "torch_dtype": "bfloat16"
+  "suppress_tokens": null,
+  "task_specific_params": null,
+  "temperature": 1.0,
+  "text_start_token_id": 277392,
+  "tf_legacy_loss": false,
+  "tie_encoder_decoder": false,
+  "tie_word_embeddings": true,
+  "tokenizer_class": null,
+  "top_k": 50,
+  "top_p": 1.0,
+  "torch_dtype": "bfloat16",
+  "torchscript": false,
+  "typical_p": 1.0,
+  "use_bfloat16": false,
+  "use_cache": true,
+  "use_embedding_sharing": true,
+  "use_gated_activation": true,
+  "use_parallel_block": true,
+  "use_parallel_embedding": false,
+  "use_qk_norm": false,
+  "vocab_size": 277395
 }

configuration_rumik_oss.py

@@ -0,0 +1,55 @@
+"""Configuration for rumik-oss 1."""
+
+from transformers import Cohere2Config
+
+
+class RumikOSSConfig(Cohere2Config):
+    """rumik-oss 1: tiny aya fire extended with flattened Mimi codec tokens.
+
+    Beyond the base Cohere2 fields, this records where the audio vocabulary
+    lives so that callers can map token ids to codec (code, quantizer) pairs
+    arithmetically, without loading the tokenizer or scanning its vocabulary.
+
+    The unit tokens are laid out code-major, quantizer-minor::
+
+        <0_0> <0_1> ... <0_7> <1_0> ... <2047_7>
+
+    so for any id in ``[first_unit_id, last_unit_id]``::
+
+        code      = (token_id - first_unit_id) // num_quantizers
+        quantizer = (token_id - first_unit_id) %  num_quantizers
+    """
+
+    model_type = "rumik_oss"
+
+    def __init__(
+        self,
+        num_quantizers=8,
+        codebook_size=2048,
+        audio_start_token_id=None,
+        audio_end_token_id=None,
+        text_start_token_id=None,
+        first_unit_id=None,
+        last_unit_id=None,
+        frame_rate_hz=12.5,
+        speakers=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.num_quantizers = int(num_quantizers)
+        self.codebook_size = int(codebook_size)
+        self.audio_start_token_id = audio_start_token_id
+        self.audio_end_token_id = audio_end_token_id
+        self.text_start_token_id = text_start_token_id
+        self.first_unit_id = first_unit_id
+        # Derivable, but stored so a hand-edited config stays self-consistent.
+        if last_unit_id is None and first_unit_id is not None:
+            last_unit_id = int(first_unit_id) + self.codebook_size * self.num_quantizers - 1
+        self.last_unit_id = last_unit_id
+        self.frame_rate_hz = float(frame_rate_hz)
+        self.speakers = speakers or ["Ira", "Aisha", "Siya", "Zoya"]
+
+    @property
+    def audio_tokens_per_second(self) -> float:
+        """Tokens the model emits per second of audio (8 x 12.5 = 100)."""
+        return self.frame_rate_hz * self.num_quantizers

configuration_tinyaya.py

@@ -1,21 +0,0 @@
-from transformers import Cohere2Config
-
-
-class TinyAyaConfig(Cohere2Config):
-    model_type = "tinyaya_mimi"
-
-    def __init__(
-        self,
-        num_quantizers=8,
-        codebook_size=2048,
-        audio_start_token_id=None,
-        audio_end_token_id=None,
-        speakers=None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        self.num_quantizers = int(num_quantizers)
-        self.codebook_size = int(codebook_size)
-        self.audio_start_token_id = audio_start_token_id
-        self.audio_end_token_id = audio_end_token_id
-        self.speakers = speakers or ["Ira", "Aisha", "Siya", "Zoya"]

modeling_tinyaya.py → modeling_rumik_oss.py

@@ -1,14 +1,23 @@
+"""rumik-oss 1: text -> flattened Mimi codec tokens.
+
+This model produces *audio tokens*, not waveforms. Decoding them to audio is
+the caller's job and needs the Mimi codec (shipped in ``codec/``)::
+
+    ids   = model.generate_audio(**inputs)          # token ids
+    codes = model.audio_tokens_to_codes(ids)        # [1, 8, T], ready for Mimi
+    wav   = mimi.decode(codes).audio_values         # your call, your decoder
+"""
 from __future__ import annotations
 
 import torch
 from torch import nn
 from transformers.models.cohere2.modeling_cohere2 import Cohere2ForCausalLM
 
-from .configuration_tinyaya import TinyAyaConfig
+from .configuration_rumik_oss import RumikOSSConfig
 
 
-class TinyAyaForCausalLM(Cohere2ForCausalLM):
-    config_class = TinyAyaConfig
+class RumikOSSForCausalLM(Cohere2ForCausalLM):
+    config_class = RumikOSSConfig
 
     def __init__(self, config):
         super().__init__(config)
@@ -20,8 +29,55 @@ class TinyAyaForCausalLM(Cohere2ForCausalLM):
             nn.Linear(max(64, hidden // 4), 1),
         )
 
+    # ---- audio vocabulary -------------------------------------------------
+    def audio_token_ids(self, device=None) -> torch.Tensor:
+        """Every id the model may legally emit inside an <audio> span.
+
+        Derived from the config, so callers no longer have to scan the
+        tokenizer vocabulary for ``<code>_<quantizer>`` tokens.
+        """
+        c = self.config
+        if c.first_unit_id is None:
+            raise ValueError(
+                "config.first_unit_id is unset; pass allowed_ids explicitly or "
+                "add first_unit_id/last_unit_id to config.json.")
+        ids = torch.arange(int(c.first_unit_id), int(c.last_unit_id) + 1, dtype=torch.long)
+        end = torch.tensor([int(c.audio_end_token_id)], dtype=torch.long)
+        return torch.cat((ids, end)).to(device if device is not None else self.device)
+
+    def audio_tokens_to_codes(self, token_ids, *, drop_partial_frame=True) -> torch.Tensor:
+        """Generated ids -> ``[1, num_quantizers, num_frames]`` codec codes.
+
+        Shaped for ``MimiModel.decode``. Purely arithmetic -- no tokenizer, no
+        codec, no audio. Anything outside the unit range (delimiters, stray
+        text tokens) ends the audio span.
+        """
+        c = self.config
+        first, last, Q = int(c.first_unit_id), int(c.last_unit_id), int(c.num_quantizers)
+        if torch.is_tensor(token_ids):
+            token_ids = token_ids.flatten().tolist()
+
+        codes: list[int] = []
+        for pos, tid in enumerate(int(t) for t in token_ids):
+            if not first <= tid <= last:
+                break                       # </audio> or anything else: done
+            if (tid - first) % Q != pos % Q:
+                break                       # off the round robin: stop cleanly
+            codes.append((tid - first) // Q)
+
+        usable = len(codes) - (len(codes) % Q) if drop_partial_frame else len(codes)
+        if usable < Q:
+            raise ValueError(
+                f"no complete codec frame in {len(token_ids)} tokens "
+                f"(need at least {Q}, got {usable})")
+        return torch.tensor(codes[:usable], dtype=torch.long).reshape(-1, Q).T.unsqueeze(0)
+
+    # ---- generation -------------------------------------------------------
     @staticmethod
-    def _sample(scores, do_sample, temperature, top_k):
+    def _constrained_sample(scores, do_sample, temperature, top_k):
+        # NB: deliberately not named `_sample`. `GenerationMixin._sample` is the
+        # method `generate()` dispatches to, and shadowing it makes every call
+        # to `generate()` fail with a TypeError on `logits_processor`.
         if not do_sample:
             return scores.argmax(dim=-1, keepdim=True)
         scores = scores / max(float(temperature), 1e-5)
@@ -36,13 +92,19 @@ class TinyAyaForCausalLM(Cohere2ForCausalLM):
         self,
         input_ids,
         attention_mask,
-        allowed_ids,
+        allowed_ids=None,
         max_new_tokens=2048,
         min_new_tokens=8,
         temperature=0.8,
         top_k=30,
         do_sample=True,
     ):
+        """Autoregressively emit audio tokens, ending at ``</audio>``.
+
+        ``allowed_ids`` defaults to the audio vocabulary from the config.
+        """
+        if allowed_ids is None:
+            allowed_ids = self.audio_token_ids(device=input_ids.device)
         out = self(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -67,7 +129,7 @@ class TinyAyaForCausalLM(Cohere2ForCausalLM):
             restricted.index_copy_(1, allowed_ids, scores.index_select(1, allowed_ids))
             if step < int(min_new_tokens):
                 restricted[:, audio_end_id] = torch.finfo(scores.dtype).min
-            token = self._sample(restricted, do_sample, temperature, top_k)
+            token = self._constrained_sample(restricted, do_sample, temperature, top_k)
             emitted = torch.cat((emitted, token), dim=1)
             if bool((token == audio_end_id).all()):
                 break