| """EmbeddingGemma 2 audio encoder + CTC reader with a head in Gemma's text-embedding space. |
| |
| wav -> log-mel -> Gemma audio encoder (12 conformer layers, 25 frames/s, chunked attention) |
| -> learned softmax mix of the 13 hidden states (subsampler + each layer) |
| -> reader: LayerNorm + Linear + causal depthwise conv + N pre-norm transformer blocks |
| -> to_text (d -> 512) @ unit_rows^T + bias, unit rows initialized from Gemma input embeddings |
| -> CTC |
| |
| Encoder adaptation: "frozen", "top_k" (unfreeze the top k layers) or "lora" (low-rank adapters on the linears). |
| Head: "trainable" (rows init from text embeddings and trained), "frozen" (rows fixed), |
| "residual" (fixed text rows + alpha * learned delta). |
| Streaming: reader self-attention is restricted to the current chunk and previous chunks; the encoder is already |
| chunked/left-context only and the reader conv is causal. |
| """ |
|
|
| import math |
| from dataclasses import asdict, dataclass, field, fields |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import AutoModel |
|
|
| from .features import LogMel, spec_augment |
|
|
|
|
| @dataclass |
| class ModelConfig: |
| encoder_id: str = "google/embeddinggemma-2" |
| vocab_size: int = 0 |
| d_model: int = 768 |
| n_layers: int = 6 |
| n_heads: int = 12 |
| ff_mult: int = 4 |
| conv_kernel: int = 15 |
| dropout: float = 0.1 |
| layer_mix: bool = True |
| encoder_mode: str = "top_k" |
| unfreeze_top_k: int = 2 |
| lora_r: int = 16 |
| lora_alpha: float = 32.0 |
| lora_targets: list = field(default_factory=lambda: ["q_proj", "k_proj", "v_proj", "post", "ffw_layer_1", "ffw_layer_2"]) |
| head_mode: str = "trainable" |
| residual_alpha: float = 1.0 |
| head_cosine: bool = True |
| chunk_frames: int = 12 |
| dynamic_chunk: bool = True |
| max_chunk_units: int = 4 |
| max_left_chunks: int = 8 |
| specaug: dict = field(default_factory=lambda: dict(freq_masks=2, freq_width=27, time_masks=10, time_width_frac=0.04)) |
|
|
| @classmethod |
| def from_dict(cls, d: dict) -> "ModelConfig": |
| known = {f.name for f in fields(cls)} |
| return cls(**{k: v for k, v in d.items() if k in known}) |
|
|
| def to_dict(self) -> dict: |
| return asdict(self) |
|
|
|
|
| class LoRALinear(nn.Module): |
| def __init__(self, base: nn.Linear, r: int, alpha: float): |
| super().__init__() |
| self.base = base |
| self.lora_a = nn.Parameter(torch.empty(r, base.in_features)) |
| self.lora_b = nn.Parameter(torch.zeros(base.out_features, r)) |
| nn.init.kaiming_uniform_(self.lora_a, a=math.sqrt(5)) |
| self.scale = alpha / r |
|
|
| def forward(self, x): |
| return self.base(x) + (x @ self.lora_a.t().to(x.dtype)) @ self.lora_b.t().to(x.dtype) * self.scale |
|
|
|
|
| def load_encoder_and_text_embeddings(encoder_id: str): |
| full = AutoModel.from_pretrained(encoder_id, dtype=torch.bfloat16, vision_config=None) |
| encoder = full.audio_tower.float() |
| text_emb = full.language_model.embed_tokens.weight.detach().float().clone() |
| del full |
| return encoder, text_emb |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, d, n_heads, ff_mult, dropout): |
| super().__init__() |
| self.n_heads = n_heads |
| self.norm1 = nn.LayerNorm(d) |
| self.qkv = nn.Linear(d, 3 * d) |
| self.proj = nn.Linear(d, d) |
| self.norm2 = nn.LayerNorm(d) |
| self.ff = nn.Sequential(nn.Linear(d, ff_mult * d), nn.GELU(), nn.Dropout(dropout), nn.Linear(ff_mult * d, d)) |
| self.drop = nn.Dropout(dropout) |
| self.attn_dropout = dropout |
|
|
| def forward(self, x, attn_mask): |
| B, T, D = x.shape |
| q, k, v = self.qkv(self.norm1(x)).view(B, T, 3, self.n_heads, D // self.n_heads).permute(2, 0, 3, 1, 4) |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, |
| dropout_p=self.attn_dropout if self.training else 0.0) |
| x = x + self.drop(self.proj(a.transpose(1, 2).reshape(B, T, D))) |
| return x + self.drop(self.ff(self.norm2(x))) |
|
|
|
|
| class CTCModel(nn.Module): |
| def __init__(self, cfg: ModelConfig, gemma_ids: list[int], encoder=None, text_emb=None): |
| super().__init__() |
| assert cfg.vocab_size == len(gemma_ids) + 2 |
| self.cfg = cfg |
| if encoder is None: |
| encoder, text_emb = load_encoder_and_text_embeddings(cfg.encoder_id) |
| self.mel = LogMel(cfg.encoder_id) |
| self.encoder = encoder |
| n_enc = encoder.config.num_hidden_layers |
| enc_dim = encoder.config.hidden_size |
| self._setup_encoder_training(n_enc) |
|
|
| |
| |
| self._acts = [] |
| grab = lambda mod, inp, out: self._acts.append(out[0] if isinstance(out, tuple) else out) |
| self.encoder.subsample_conv_projection.register_forward_hook(grab) |
| for layer in self.encoder.layers: |
| layer.register_forward_hook(grab) |
|
|
| d = cfg.d_model |
| self.layer_logits = nn.Parameter(torch.zeros(n_enc + 1)) if cfg.layer_mix else None |
| self.in_norm = nn.LayerNorm(enc_dim) |
| self.in_proj = nn.Linear(enc_dim, d) |
| self.pos_conv = nn.Conv1d(d, d, cfg.conv_kernel, groups=d) |
| self.blocks = nn.ModuleList([TransformerBlock(d, cfg.n_heads, cfg.ff_mult, cfg.dropout) for _ in range(cfg.n_layers)]) |
| self.out_norm = nn.LayerNorm(d) |
|
|
| rows = text_emb[torch.tensor(gemma_ids)] |
| rows = rows / rows.norm(dim=-1, keepdim=True).mean() |
| text_dim = rows.shape[1] |
| self.to_text = nn.Linear(d, text_dim) |
| if cfg.head_mode == "trainable": |
| self.unit_rows = nn.Parameter(rows.clone()) |
| elif cfg.head_mode == "frozen": |
| self.register_buffer("unit_rows", rows.clone()) |
| elif cfg.head_mode == "residual": |
| self.register_buffer("text_rows", rows.clone()) |
| self.delta_rows = nn.Parameter(torch.zeros_like(rows)) |
| else: |
| raise ValueError(cfg.head_mode) |
| self.special_rows = nn.Parameter(torch.randn(2, text_dim) * 0.02) |
| self.head_bias = nn.Parameter(torch.zeros(cfg.vocab_size)) |
| if cfg.head_cosine: |
| self.logit_scale = nn.Parameter(torch.tensor(math.log(10.0))) |
|
|
| def _setup_encoder_training(self, n_enc): |
| cfg = self.cfg |
| for p in self.encoder.parameters(): |
| p.requires_grad_(False) |
| if cfg.encoder_mode == "top_k": |
| for layer in self.encoder.layers[n_enc - cfg.unfreeze_top_k:]: |
| for p in layer.parameters(): |
| p.requires_grad_(True) |
| elif cfg.encoder_mode == "lora": |
| for layer in self.encoder.layers: |
| for name, mod in list(layer.named_modules()): |
| if name.split(".")[-1] in cfg.lora_targets and hasattr(mod, "linear"): |
| mod.linear = LoRALinear(mod.linear, cfg.lora_r, cfg.lora_alpha) |
| elif cfg.encoder_mode != "frozen": |
| raise ValueError(cfg.encoder_mode) |
|
|
| def head_rows(self): |
| if self.cfg.head_mode == "residual": |
| rows = self.text_rows + self.cfg.residual_alpha * self.delta_rows |
| else: |
| rows = self.unit_rows |
| return torch.cat([self.special_rows, rows], dim=0) |
|
|
| @staticmethod |
| def chunk_mask(T, chunk, left_chunks, device): |
| c = torch.arange(T, device=device) // chunk |
| allow = c[None, :] <= c[:, None] |
| if left_chunks is not None: |
| allow &= c[None, :] >= c[:, None] - left_chunks |
| return allow |
|
|
| def forward(self, wav, wav_lens, streaming=False, left_chunks=None): |
| """Returns float32 logits (B, T, V) and output lengths (B,).""" |
| feats, fmask = self.mel(wav, wav_lens) |
| if self.training and self.cfg.specaug: |
| feats = spec_augment(feats, fmask, **self.cfg.specaug) |
|
|
| self._acts.clear() |
| out = self.encoder(feats, fmask, return_dict=True) |
| acts, self._acts = self._acts, [] |
| mask = out.attention_mask |
| if self.cfg.layer_mix: |
| w = torch.softmax(self.layer_logits.float(), dim=0) |
| h = sum(w[i] * a.float() for i, a in enumerate(acts)) |
| else: |
| h = acts[-1].float() |
|
|
| x = self.in_proj(self.in_norm(h)) |
| xc = F.pad((x * mask[..., None]).transpose(1, 2), (self.cfg.conv_kernel - 1, 0)) |
| x = x + F.gelu(self.pos_conv(xc).transpose(1, 2)) |
|
|
| T = x.shape[1] |
| attn_mask = mask[:, None, None, :] |
| chunk = None |
| if streaming: |
| chunk = self.cfg.chunk_frames |
| elif self.training and self.cfg.dynamic_chunk and torch.rand(()) < 0.5: |
| chunk = self.cfg.chunk_frames * int(torch.randint(1, self.cfg.max_chunk_units + 1, ())) |
| left_chunks = None if torch.rand(()) < 0.5 else int(torch.randint(0, self.cfg.max_left_chunks + 1, ())) |
| if chunk is not None: |
| attn_mask = attn_mask & self.chunk_mask(T, chunk, left_chunks, x.device)[None, None] |
| |
| attn_mask = attn_mask | torch.eye(T, dtype=torch.bool, device=x.device)[None, None] |
| for blk in self.blocks: |
| x = blk(x, attn_mask) |
|
|
| t = self.to_text(self.out_norm(x)).float() |
| rows = self.head_rows().float() |
| if self.cfg.head_cosine: |
| logits = self.logit_scale.exp() * (F.normalize(t, dim=-1) @ F.normalize(rows, dim=-1).t()) |
| else: |
| logits = t @ rows.t() |
| logits = logits + self.head_bias.float() |
| return logits, mask.sum(1) |
|
|
| def trainable_state_dict(self): |
| """Only what is not recoverable from the base encoder: trained params + LoRA + head buffers.""" |
| keep = {n for n, p in self.named_parameters() if p.requires_grad} |
| sd = self.state_dict() |
| return {k: v for k, v in sd.items() |
| if k in keep or not k.startswith("encoder.") or "lora_" in k} |
|
|
|
|
| def greedy_decode(logits, lens, blank=0): |
| best = logits.argmax(-1).cpu() |
| out = [] |
| for b in range(best.shape[0]): |
| seq = best[b, : int(lens[b])] |
| keep = torch.ones_like(seq, dtype=torch.bool) |
| keep[1:] = seq[1:] != seq[:-1] |
| seq = seq[keep] |
| out.append(seq[seq != blank].tolist()) |
| return out |
|
|
