#!/usr/bin/env python3 """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 # ------------------------------------------------------------------ the audio contract 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 # ------------------------------------------------------------------ the reader architecture # (exact copy of the trained model definition; the checkpoint loads with strict=True) 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 # ------------------------------------------------------------------ EmbeddingGemma 2 soft tokens 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 # ------------------------------------------------------------------ the demo 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()