Spaces:
Running on Zero
Running on Zero
| import spaces # must come before torch | |
| import json | |
| import math | |
| import time | |
| import gradio as gr | |
| import librosa | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from huggingface_hub import hf_hub_download | |
| from transformers import EmbeddingGemma2Model, Gemma4AudioFeatureExtractor | |
| # Full float32 math, as the reference (no TF32). | |
| torch.backends.cuda.matmul.allow_tf32 = False | |
| torch.backends.cudnn.allow_tf32 = False | |
| EG2_REPO = "google/embeddinggemma-2" | |
| EG2_REV = "914f7f89142e33e77833254d9c9b90c3cef7303b" | |
| READER_REPO = "mlboydaisuke/EmbeddingGemma-2-Transcribe-Reader" | |
| SR = 16000 | |
| MAX_SECONDS = 15 * 60 | |
| # Window contract (reader_config.json -> window): 30 s windows (3000 frames / 750 tokens), | |
| # 96 tokens of left overlap, 1 token dropped at the right of every window but the last. | |
| WIN_FRAMES, LEFT_OVERLAP, DROP_LAST = 3000, 96, 1 | |
| WIN_TOKENS = WIN_FRAMES // 4 | |
| STEP_TOKENS = WIN_TOKENS - LEFT_OVERLAP - DROP_LAST # 653 | |
| # ------------------------------------------------------------------ reader (reader/train/model.py, arch "a") | |
| D_IN, D_OUT = 512, 768 | |
| 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): | |
| def __init__(self, d=512, layers=6, 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 = (torch.arange(T, device=x.device)[None, :] < lengths[:, None])[:, None, None, :] | |
| 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): | |
| def __init__(self, layers=6): | |
| 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))) | |
| def scale(self): | |
| return self.log_scale.exp() | |
| def forward(self, x, lengths): | |
| v = self.body(x, lengths) | |
| return F.normalize(v, dim=-1), self.blank(v).squeeze(-1) | |
| # ------------------------------------------------------------------ load everything at module scope | |
| cfg = json.load(open(hf_hub_download(READER_REPO, "reader_config.json"))) | |
| with open(hf_hub_download(READER_REPO, "vocab.txt"), encoding="utf-8") as f: | |
| VOCAB = [line.rstrip("\n") for line in f if line.strip()] | |
| BLANK_ID = len(VOCAB) | |
| reader = Reader(layers=cfg["architecture"]["layers"]) | |
| reader.load_state_dict(torch.load(hf_hub_download(READER_REPO, "reader.pt"), map_location="cpu", weights_only=True)) | |
| reader = reader.to("cuda").eval() | |
| table = torch.from_numpy(np.load(hf_hub_download(READER_REPO, "word_table.npy")).astype(np.float32)) | |
| table = F.normalize(table, dim=-1).to("cuda") | |
| feature_extractor = Gemma4AudioFeatureExtractor.from_pretrained(EG2_REPO, revision=EG2_REV) | |
| encoder = EmbeddingGemma2Model.from_pretrained(EG2_REPO, revision=EG2_REV, dtype=torch.float32, vision_config=None) | |
| encoder = encoder.to("cuda").eval() | |
| print("audio attention:", getattr(encoder.config.audio_config, "_attn_implementation", None)) | |
| def n_tokens_for(n_frames): | |
| return ((n_frames + 1) // 2 + 1) // 2 | |
| def plan_windows(n_frames): | |
| """Windowing.plan of the eg2-transcribe app: (token0, frame0, frames, keep_start, keep_end).""" | |
| total = n_tokens_for(n_frames) | |
| plan, i = [], 0 | |
| while total > 0: | |
| a = i * STEP_TOKENS | |
| nf = min(WIN_FRAMES, n_frames - 4 * a) | |
| m = n_tokens_for(nf) | |
| final = a + m >= total | |
| plan.append((a, 4 * a, nf, 0 if i == 0 else LEFT_OVERLAP, m if final else WIN_TOKENS - DROP_LAST)) | |
| if final: | |
| break | |
| i += 1 | |
| return plan | |
| def load_audio(path): | |
| wav, _ = librosa.load(path, sr=SR, mono=True) | |
| return wav.astype(np.float32) | |
| def get_duration(audio_path, *args, **kwargs): | |
| try: | |
| secs = librosa.get_duration(path=audio_path) | |
| except Exception: | |
| secs = 60 | |
| return int(min(60, 8 + secs * 0.01)) | |
| def transcribe(audio_path: str) -> tuple[str, str]: | |
| """Transcribe English speech with the frozen EmbeddingGemma 2 audio encoder + the trained word reader. | |
| Args: | |
| audio_path: path to an audio file (any format librosa can read; resampled to 16 kHz mono). | |
| Returns: | |
| The lower-case transcript (no punctuation) and a short stats line. | |
| """ | |
| if not audio_path: | |
| raise gr.Error("Please upload or record some audio first.") | |
| t0 = time.time() | |
| wav = load_audio(audio_path) | |
| secs = len(wav) / SR | |
| if secs > MAX_SECONDS: | |
| raise gr.Error(f"Audio is {secs / 60:.1f} min long; this demo accepts up to {MAX_SECONDS // 60} min.") | |
| n_frames = max(0, (len(wav) - 161) // 160 + 1) | |
| if n_frames < 4: | |
| raise gr.Error("Audio is too short.") | |
| # log-mel of the whole recording (no 30 s cut); each window is a slice of its frames | |
| out = feature_extractor([wav], return_tensors="np", max_length=None, truncation=False) | |
| feats = out["input_features"][0][:n_frames].astype(np.float32) | |
| words_all = [] | |
| plan = plan_windows(n_frames) | |
| for token0, frame0, nf, k0, k1 in plan: | |
| x = torch.from_numpy(np.ascontiguousarray(feats[frame0:frame0 + nf]))[None].to("cuda") | |
| m = torch.ones(1, nf, dtype=torch.bool, device="cuda") | |
| res = encoder.get_audio_features(x, m, return_dict=True) | |
| soft = res.pooler_output[0][res.attention_mask[0]].float() # [tokens, 512] | |
| lengths = torch.tensor([soft.shape[0]], device="cuda") | |
| h, blank = reader(soft[None], lengths) | |
| best, ids = (h[0] @ table.T).max(dim=1) | |
| labels = torch.where(blank[0] > best * reader.scale, torch.full_like(ids, BLANK_ID), ids) | |
| words_all.extend(labels[k0:k1].tolist()) | |
| # greedy CTC read-out over the kept tokens of all windows | |
| text, prev = [], None | |
| for c in words_all: | |
| if c != prev and c != BLANK_ID: | |
| text.append(VOCAB[c]) | |
| prev = c | |
| dt = time.time() - t0 | |
| stats = f"{secs:.1f} s of audio · {len(plan)} window(s) of 30 s · {len(words_all)} tokens · {dt:.2f} s" | |
| return " ".join(text), stats | |
| CSS = """ | |
| #col-container { max-width: 1000px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# EmbeddingGemma 2 Transcribe\n" | |
| "English speech-to-text from the **frozen audio encoder of " | |
| "[EmbeddingGemma 2](https://huggingface.co/google/embeddinggemma-2)** plus a small trained reader " | |
| "that maps every 40 ms token to one of 200,000 English words (or blank) — greedy CTC, no language model. " | |
| "Runs the PyTorch reference path of " | |
| "[EmbeddingGemma-2-Transcribe-LiteRT](https://huggingface.co/mlboydaisuke/EmbeddingGemma-2-Transcribe-LiteRT) " | |
| "/ [Reader v2](https://huggingface.co/mlboydaisuke/EmbeddingGemma-2-Transcribe-Reader), " | |
| "with the same 30 s sliding-window contract as the on-device app." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio = gr.Audio(label="English speech", type="filepath", sources=["upload", "microphone"]) | |
| run = gr.Button("Transcribe", variant="primary") | |
| with gr.Column(): | |
| transcript = gr.Textbox(label="Transcript", lines=8, buttons=["copy"]) | |
| stats = gr.Markdown() | |
| gr.Markdown( | |
| "Output is lower case without punctuation. Trained on read audiobook speech (LibriSpeech), " | |
| f"so clean read English works best. Up to {MAX_SECONDS // 60} min of audio." | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/1284-1180-0000.flac"], | |
| ["examples/1089-134686-0000.flac"], | |
| ["examples/2300-131720-0000.flac"], | |
| ["examples/4507-16021-0047.flac"], | |
| ], | |
| inputs=[audio], | |
| outputs=[transcript, stats], | |
| fn=transcribe, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| gr.Markdown( | |
| "Examples: LibriSpeech test-clean (Panayotov et al., ICASSP 2015, CC BY 4.0). " | |
| "Use must follow the [Gemma Prohibited Use Policy](https://ai.google.dev/gemma/prohibited_use_policy)." | |
| ) | |
| run.click(transcribe, inputs=[audio], outputs=[transcript, stats]) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |