| |
| """Standalone transcription demo for the eg2-voice CTC reader (13.3 M parameters). |
| |
| Transcribes one 16 kHz mono WAV/FLAC file end to end, 100% locally (CPU by default): |
| |
| 1. EmbeddingGemma 2 audio encoder -> soft tokens [T, 512], 25 tokens per second |
| 2. this reader -> per-frame scores against 200,000 word vectors + a learned blank logit |
| 3. greedy CTC decode (argmax, merge repeats, drop blanks) -> words |
| |
| Requirements (no repo needed — this file is self-contained): |
| pip install torch numpy soundfile safetensors "transformers>=5.19" |
| (scipy only if your input is not 16 kHz; ffmpeg only for non-libsndfile formats) |
| |
| The EmbeddingGemma 2 base model is loaded from a local snapshot directory (recommended) |
| or downloaded from the hub on first use (several GB): |
| python example.py clip.flac --eg2 /path/to/embeddinggemma-2 |
| python example.py clip.flac --eg2 google/embeddinggemma-2 |
| """ |
| import argparse |
| import math |
| import os |
| import time |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| D_IN, D_OUT = 512, 768 |
| SR = 16000 |
|
|
|
|
| |
|
|
| def n_frames_for(n_samples): |
| """Valid mel frames: frame i needs samples up to 160 * i + 160.""" |
| return max(0, (n_samples - 161) // 160 + 1) |
|
|
|
|
| def n_tokens_for(n_frames): |
| """ceil(ceil(n_frames / 2) / 2): two stride-2 convolutions.""" |
| return ((n_frames + 1) // 2 + 1) // 2 |
|
|
|
|
| |
| |
|
|
| def key_mask(lengths, T): |
| """[B, 1, 1, T] bool, True = a real token (attend), for scaled_dot_product_attention.""" |
| return (torch.arange(T, device=lengths.device)[None, :] < lengths[:, None])[:, None, None, :] |
|
|
|
|
| def rope(T, dim, base, device): |
| inv = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)) |
| f = torch.arange(T, device=device, dtype=torch.float32)[:, None] * inv[None, :] |
| emb = torch.cat([f, f], dim=-1) |
| return emb.cos(), emb.sin() |
|
|
|
|
| def rotate_half(x): |
| x1, x2 = x.chunk(2, dim=-1) |
| return torch.cat([-x2, x1], dim=-1) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d, heads, ffn): |
| super().__init__() |
| self.heads = heads |
| self.ln1 = nn.LayerNorm(d) |
| self.qkv = nn.Linear(d, 3 * d) |
| self.proj = nn.Linear(d, d) |
| self.ln2 = nn.LayerNorm(d) |
| self.ff1 = nn.Linear(d, ffn) |
| self.ff2 = nn.Linear(ffn, d) |
|
|
| def forward(self, x, mask, cos, sin): |
| B, T, D = x.shape |
| q, k, v = self.qkv(self.ln1(x)).view(B, T, 3, self.heads, D // self.heads).permute(2, 0, 3, 1, 4) |
| q = q * cos + rotate_half(q) * sin |
| k = k * cos + rotate_half(k) * sin |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) |
| x = x + self.proj(a.transpose(1, 2).reshape(B, T, D)) |
| return x + self.ff2(F.gelu(self.ff1(self.ln2(x)))) |
|
|
|
|
| class ReaderA(nn.Module): |
| """Pre-LN bidirectional transformer: LN -> Linear 512->512 -> N blocks -> LN -> Linear 512->768.""" |
|
|
| def __init__(self, d=512, layers=4, heads=8, ffn=2048, rope_base=10000.0): |
| super().__init__() |
| self.heads, self.rope_base = heads, rope_base |
| self.ln_in = nn.LayerNorm(D_IN) |
| self.inp = nn.Linear(D_IN, d) |
| self.blocks = nn.ModuleList([Block(d, heads, ffn) for _ in range(layers)]) |
| self.ln_out = nn.LayerNorm(d) |
| self.out = nn.Linear(d, D_OUT) |
|
|
| def forward(self, x, lengths): |
| T = x.shape[1] |
| mask = key_mask(lengths, T) |
| cos, sin = rope(T, self.inp.out_features // self.heads, self.rope_base, x.device) |
| h = self.inp(self.ln_in(x)) |
| for blk in self.blocks: |
| h = blk(h, mask, cos, sin) |
| return self.out(self.ln_out(h)) |
|
|
|
|
| class Reader(nn.Module): |
| """Body + linear blank head + learned log_scale (scale = exp(log_scale), starts at 20).""" |
|
|
| def __init__(self, layers=4): |
| super().__init__() |
| self.body = ReaderA(layers=layers) |
| self.blank = nn.Linear(D_OUT, 1) |
| self.log_scale = nn.Parameter(torch.tensor(math.log(20.0))) |
|
|
| @property |
| def scale(self): |
| return self.log_scale.exp() |
|
|
| def forward(self, x, lengths): |
| """-> unit vectors [B, T, 768], blank logits [B, T].""" |
| v = self.body(x, lengths) |
| return F.normalize(v, dim=-1), self.blank(v).squeeze(-1) |
|
|
|
|
| @torch.no_grad() |
| def greedy_labels(h, blank, table, scale, chunk=1024): |
| """h [M, D] unit vectors, blank [M] -> per-frame class [M] (blank = index N). |
| |
| A frame goes to the best word only when scale * cos(h, word) beats the frame's blank logit.""" |
| M = h.shape[0] |
| if M == 0: |
| return torch.zeros(0, dtype=torch.long, device=h.device) |
| out = [] |
| for a in range(0, M, chunk): |
| v, i = (h[a:a + chunk] @ table.T).max(dim=1) |
| out.append(torch.where(blank[a:a + chunk] > v * scale, torch.full_like(i, table.shape[0]), i)) |
| return torch.cat(out) |
|
|
|
|
| def collapse(labels, blank_id): |
| """Greedy CTC read-out: merge repeats, drop blanks. labels: 1-D sequence of ints.""" |
| out, prev = [], None |
| for c in labels: |
| if c != prev and c != blank_id: |
| out.append(c) |
| prev = c |
| return out |
|
|
|
|
| |
|
|
| def load_eg2(model_dir, device): |
| """-> (feature extractor, EmbeddingGemma 2 audio+text model in float32 on device). |
| |
| transformers 5.19.0 contract: float32 (never float16), default/sdpa attention (never eager) — |
| the eager audio attention reads the wrong positions in this version.""" |
| from transformers import EmbeddingGemma2Model, Gemma4AudioFeatureExtractor |
| fe = Gemma4AudioFeatureExtractor.from_pretrained(model_dir) |
| model = EmbeddingGemma2Model.from_pretrained(model_dir, dtype=torch.float32, vision_config=None) |
| model = model.to(device).eval() |
| impl = getattr(model.config.audio_config, "_attn_implementation", None) |
| if impl == "eager": |
| raise SystemExit("audio attention is eager; transformers 5.19.0 eager reads the wrong positions") |
| return fe, model |
|
|
|
|
| @torch.no_grad() |
| def soft_tokens(model, fe, wav, device): |
| """-> soft tokens [T, 512] float32 numpy for one waveform (no 30 s cut). |
| |
| A one-item list, as the processor expects: a bare 1-D array fails in transformers 5.19.0.""" |
| out = fe([np.asarray(wav, dtype=np.float32)], return_tensors="np", max_length=None, truncation=False) |
| feats = torch.from_numpy(out["input_features"][0].astype(np.float32))[None].to(device) |
| mask = torch.from_numpy(out["input_features_mask"][0].astype(bool))[None].to(device) |
| n_frames = n_frames_for(len(wav)) |
| if int(mask.sum()) != n_frames: |
| raise AssertionError(f"extractor mask {int(mask.sum())} != {n_frames} frames") |
| res = model.get_audio_features(feats, mask, return_dict=True) |
| soft = res.pooler_output[0][res.attention_mask[0]].float().cpu().numpy() |
| if soft.shape[0] != n_tokens_for(n_frames): |
| raise AssertionError(f"{soft.shape[0]} tokens != {n_tokens_for(n_frames)}") |
| return soft |
|
|
|
|
| |
|
|
| def load_audio(path): |
| """-> float32 mono 16 kHz in [-1, 1].""" |
| import soundfile as sf |
| wav, sr = sf.read(path, dtype="float32", always_2d=False) |
| if wav.ndim == 2: |
| wav = wav.mean(axis=1) |
| wav = np.asarray(wav, dtype=np.float32) |
| np.clip(wav, -1.0, 1.0, out=wav) |
| if sr != SR: |
| try: |
| from scipy.signal import resample_poly |
| from math import gcd |
| g = gcd(SR, int(sr)) or 1 |
| wav = np.asarray(resample_poly(wav, SR // g, sr // g), dtype=np.float32) |
| except ImportError: |
| raise SystemExit(f"input is {sr} Hz; install scipy to resample, or provide 16 kHz audio") |
| return wav |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("audio", help="WAV/FLAC file (16 kHz mono recommended)") |
| ap.add_argument("--eg2", default=os.environ.get("EG2_MODEL", "google/embeddinggemma-2"), |
| help="EmbeddingGemma 2 snapshot dir or hub id (default: env EG2_MODEL or the hub)") |
| ap.add_argument("--reader", default=os.path.dirname(os.path.abspath(__file__)), |
| help="dir with reader.safetensors, vocab.txt, word_table.npy (default: beside this file)") |
| ap.add_argument("--device", default="cpu") |
| args = ap.parse_args() |
|
|
| t0 = time.perf_counter() |
| print(f"loading EmbeddingGemma 2 from {args.eg2} ...", flush=True) |
| fe, eg2 = load_eg2(args.eg2, args.device) |
|
|
| reader_dir = args.reader |
| import json |
| with open(os.path.join(reader_dir, "config.json"), encoding="utf-8") as f: |
| cfg = json.load(f) |
| with open(os.path.join(reader_dir, cfg["vocab"]), encoding="utf-8") as f: |
| vocab = [line.rstrip("\n") for line in f if line.strip()] |
| table = torch.from_numpy(np.load(os.path.join(reader_dir, cfg["table"]["file"])).astype(np.float32)) |
| table = F.normalize(table.to(args.device), dim=-1) |
| model = Reader(layers=cfg["architecture"]["layers"]) |
| from safetensors.torch import load_file |
| model.load_state_dict(load_file(os.path.join(reader_dir, "reader.safetensors")), strict=True) |
| model = model.to(args.device).eval() |
| n = sum(p.numel() for p in model.parameters()) |
| print(f"reader ready: {n:,} params, {len(vocab):,} words, scale {model.scale.item():.2f} " |
| f"({time.perf_counter() - t0:.0f} s)", flush=True) |
|
|
| wav = load_audio(args.audio) |
| if len(wav) / SR < 0.15: |
| raise SystemExit("audio too short (< 0.15 s)") |
| t0 = time.perf_counter() |
| tokens = soft_tokens(eg2, fe, wav, args.device) |
| enc_s = time.perf_counter() - t0 |
| t0 = time.perf_counter() |
| x = torch.from_numpy(tokens[None]).to(args.device) |
| lengths = torch.tensor([x.shape[1]], dtype=torch.long, device=args.device) |
| with torch.no_grad(): |
| h, blank = model(x, lengths) |
| labels = greedy_labels(h[0], blank[0], table, model.scale).cpu().tolist() |
| words = collapse(labels, len(vocab)) |
| text = " ".join(vocab[i] for i in words) |
| dec_s = time.perf_counter() - t0 |
| print(f"encode {enc_s:.2f} s | decode {dec_s:.2f} s | {tokens.shape[0]} tokens " |
| f"({len(wav) / SR:.1f} s of audio)") |
| print(text) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
