Spaces:
Running on Zero
Running on Zero
| """Backbone wrapper: a HF causal LM consumed as a 12.5 Hz sequence model. | |
| Both candidate backbones load through the same wrapper so comparing them is a | |
| config switch, not a code fork: | |
| gemma: google/gemma-3-4b-pt (262k vocab, 5:1 sliding window — small KV) | |
| qwen: Qwen/Qwen3-4B-Base (Apache-2.0, full attention) | |
| The wrapper owns three things the raw LM doesn't have: | |
| - vocab extension for the token contract (resize + mean-init of new rows); | |
| - an input-composition layer: text positions use the LM's embedding table, | |
| audio positions use the depth decoder's summed-code feedback vector, the | |
| speaker slot uses the conditioning projector's output — all merged into one | |
| (B, T, d) inputs_embeds tensor (audio never enters the text vocab); | |
| - the cb0 head: n_codes+1 classes (class 2048 = <aeos>), separate from lm_head. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from ..tokens import CONTRACT, N_CODES, AEOS_CLASS, extend_tokenizer | |
| BACKBONES = { | |
| "gemma": "google/gemma-3-4b-pt", | |
| "qwen": "Qwen/Qwen3-4B-Base", | |
| "llama": "meta-llama/Llama-3.2-3B", | |
| } | |
| def _from_config(path): | |
| import glob, os | |
| if os.environ.get("MAYA3_BACKBONE_FROM_CONFIG") == "1": | |
| return True | |
| return os.path.isdir(path) and not glob.glob(os.path.join(path, "*.safetensors")) and not glob.glob(os.path.join(path, "*.bin")) | |
| def _lm_from_config(path, dtype, attn): | |
| from transformers import AutoConfig | |
| cfg = AutoConfig.from_pretrained(path) | |
| with torch.device("meta"): | |
| lm = AutoModelForCausalLM.from_config(cfg, dtype=dtype, attn_implementation=attn) | |
| lm = lm.to_empty(device="cpu") | |
| for m in lm.modules(): # non-persistent buffers (RoPE inv_freq) were never materialised: recompute them | |
| if hasattr(m, "inv_freq") and hasattr(m, "rope_type"): | |
| if getattr(m, "rope_init_fn", None) is not None: | |
| fn = m.rope_init_fn | |
| elif m.rope_type == "default": | |
| fn = m.compute_default_rope_parameters | |
| else: | |
| from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS | |
| fn = ROPE_INIT_FUNCTIONS[m.rope_type] | |
| inv, scale = fn(m.config, torch.device("cpu")) | |
| m.inv_freq = inv | |
| if hasattr(m, "original_inv_freq"): | |
| m.original_inv_freq = inv | |
| if hasattr(m, "attention_scaling"): | |
| m.attention_scaling = scale | |
| lm.tie_weights() | |
| return lm | |
| class Backbone(nn.Module): | |
| def __init__(self, name="qwen", dtype=torch.float32, attn="flash_attention_2"): | |
| # fp32 parameters by default: FSDP2's MixedPrecisionPolicy casts to bf16 | |
| # for compute but the SHARDED source params are what AdamW updates — load | |
| # them bf16 and there are no fp32 masters. | |
| super().__init__() | |
| path = BACKBONES.get(name, name) | |
| if _from_config(path): | |
| # inference: every backbone weight comes from the Maya3 checkpoint loaded strictly afterwards, so build the | |
| # architecture without reading (or even needing) the multi-GB backbone safetensors | |
| self.lm = _lm_from_config(path, dtype, attn) | |
| else: | |
| self.lm = AutoModelForCausalLM.from_pretrained(path, dtype=dtype, | |
| attn_implementation=attn) | |
| self.tokenizer = AutoTokenizer.from_pretrained(path) | |
| self.tokenizer, self.token_ids = extend_tokenizer(self.tokenizer) | |
| old = self.lm.get_input_embeddings().weight.shape[0] | |
| if len(self.tokenizer) > old: | |
| self.lm.resize_token_embeddings(len(self.tokenizer)) | |
| with torch.no_grad(): | |
| emb = self.lm.get_input_embeddings().weight | |
| emb[old:] = emb[:old].mean(0, keepdim=True) + 0.02 * torch.randn_like(emb[old:]) | |
| self.d_model = self.lm.config.hidden_size | |
| # cb0 + <aeos>: the only audio-side output the backbone owns | |
| self.cb0_head = nn.Linear(self.d_model, N_CODES + 1, bias=False) | |
| def embed_text(self, ids): | |
| return self.lm.get_input_embeddings()(ids) | |
| def forward(self, inputs_embeds, attention_mask=None, past_key_values=None, | |
| use_cache=False): | |
| out = self.lm.model(inputs_embeds=inputs_embeds, attention_mask=attention_mask, | |
| past_key_values=past_key_values, use_cache=use_cache) | |
| return out.last_hidden_state, getattr(out, "past_key_values", None) | |
| def text_logits(self, hidden): | |
| return self.lm.lm_head(hidden) | |
| def cb0_logits(self, hidden): | |
| return self.cb0_head(hidden) | |