| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
| from transformers.models.cohere2.modeling_cohere2 import Cohere2ForCausalLM |
|
|
| from .configuration_tinyaya import TinyAyaConfig |
|
|
|
|
| class TinyAyaForCausalLM(Cohere2ForCausalLM): |
| config_class = TinyAyaConfig |
|
|
| 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), |
| ) |
|
|
| @staticmethod |
| def _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, |
| max_new_tokens=2048, |
| min_new_tokens=8, |
| temperature=0.8, |
| top_k=30, |
| do_sample=True, |
| ): |
| 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._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 |
|
|
