Blame for example.py

#!/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()