| """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) # prompt + audio |
| tokens = ids[0].tolist()[inputs.input_ids.shape[1]:] # drop the prompt |
| codes = model.audio_tokens_to_codes(tokens) # [1, 8, T] for Mimi |
| wav = mimi.decode(codes).audio_values # your call |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from transformers.models.cohere2.modeling_cohere2 import Cohere2ForCausalLM |
|
|
| from .configuration_rumik_oss import RumikOSSConfig |
|
|
|
|
| class RumikOSSForCausalLM(Cohere2ForCausalLM): |
| config_class = RumikOSSConfig |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| hidden = int(config.hidden_size) |
| self.stop_predictor = nn.Sequential( |
| nn.LayerNorm(hidden), |
| nn.Linear(hidden, max(64, hidden // 4)), |
| nn.GELU(), |
| nn.Linear(max(64, hidden // 4), 1), |
| ) |
|
|
| |
| 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) -> torch.Tensor: |
| """Generated ids -> ``[1, num_quantizers, num_frames]`` codec codes. |
| |
| Shaped for ``MimiModel.decode``. Purely arithmetic -- no tokenizer, no |
| codec, no audio. Pass only the generated tail, not the prompt. |
| |
| A token that breaks the round robin does not end the audio: the |
| partial frame is dropped and the next complete frame is picked up |
| again, so one bad token costs one frame rather than the rest of the |
| clip. ``</audio>`` ends it. |
| """ |
| c = self.config |
| first, last, Q = int(c.first_unit_id), int(c.last_unit_id), int(c.num_quantizers) |
| end_id = int(c.audio_end_token_id) |
| if torch.is_tensor(token_ids): |
| token_ids = token_ids.flatten().tolist() |
|
|
| frames: list[list[int]] = [] |
| frame: list[int] = [] |
| for tid in (int(t) for t in token_ids): |
| if tid == end_id: |
| break |
| if not first <= tid <= last: |
| frame = [] |
| continue |
| code, q = divmod(tid - first, Q) |
| if q == len(frame): |
| frame.append(code) |
| if len(frame) == Q: |
| frames.append(frame) |
| frame = [] |
| else: |
| frame = [code] if q == 0 else [] |
| if not frames: |
| raise ValueError( |
| f"no complete codec frame in {len(token_ids)} tokens " |
| f"(need at least {Q})") |
| return torch.tensor(frames, dtype=torch.long).T.unsqueeze(0) |
|
|
| |
| @staticmethod |
| def _constrained_sample(scores, do_sample, temperature, top_k): |
| |
| |
| |
| if not do_sample: |
| return scores.argmax(dim=-1, keepdim=True) |
| scores = scores / max(float(temperature), 1e-5) |
| if int(top_k) > 0: |
| k = min(int(top_k), scores.shape[-1]) |
| cutoff = torch.topk(scores, k, dim=-1).values[:, -1:] |
| scores = scores.masked_fill(scores < cutoff, torch.finfo(scores.dtype).min) |
| return torch.multinomial(torch.softmax(scores.float(), dim=-1), 1) |
|
|
| @torch.inference_mode() |
| def generate_audio( |
| self, |
| input_ids, |
| attention_mask, |
| 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, |
| use_cache=True, |
| output_hidden_states=True, |
| return_dict=True, |
| ) |
| emitted = input_ids |
| cache = out.past_key_values |
| scores = out.logits[:, -1, :] |
| hidden = out.hidden_states[-1][:, -1:, :] |
| mask = attention_mask |
| allowed_ids = allowed_ids.to(scores.device) |
| audio_end_id = int(self.config.audio_end_token_id) |
| for step in range(int(max_new_tokens)): |
| if step >= int(min_new_tokens): |
| stop = torch.sigmoid(self.stop_predictor(hidden).squeeze(-1)) |
| if bool((stop > 0.5).all()): |
| eos = input_ids.new_full((input_ids.shape[0], 1), audio_end_id) |
| return torch.cat((emitted, eos), dim=1) |
| restricted = torch.full_like(scores, torch.finfo(scores.dtype).min) |
| 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._constrained_sample(restricted, do_sample, temperature, top_k) |
| emitted = torch.cat((emitted, token), dim=1) |
| if bool((token == audio_end_id).all()): |
| break |
| mask = torch.cat((mask, torch.ones_like(token)), dim=1) |
| out = self( |
| input_ids=token, |
| attention_mask=mask, |
| past_key_values=cache, |
| use_cache=True, |
| output_hidden_states=True, |
| return_dict=True, |
| ) |
| cache = out.past_key_values |
| scores = out.logits[:, -1, :] |
| hidden = out.hidden_states[-1][:, -1:, :] |
| return emitted |
|
|
