""" Kinyarwanda ASR demo — Conformer-CTC trained from scratch (/ LanguageBuild). Space: https://huggingface.co/spaces/nsarob/kinyasr2026 Transcribes microphone / uploaded audio with greedy CTC decoding. Model, tokenizer, and the feature front-end are byte-identical to training so the checkpoint loads cleanly. Honest scope: trained on ONE in-domain corpus (image-description speech, single 2025 campaign). Spontaneous / real-mic / health-domain speech is out-of-distribution and may transcribe poorly. This is a research PoC, not a product. """ import os import math import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torchaudio import torchaudio.transforms as T import gradio as gr # ----- config (must match training) ----- PRESET = dict(d_model=256, n_heads=4, n_layers=12, ff_mult=4, conv_kernel=31) SR, N_MELS, N_FFT, HOP = 16000, 80, 400, 160 MAX_SECS = 30 # cap input length for responsiveness on CPU # Checkpoint resolution: # 1) a LOCAL file named CKPT_FILE next to app.py <-- simplest: just drop the .ckpt here # (works the same locally and inside the Space repo). # 2) else, if CKPT_REPO is set, download that file from a HF *model* repo. CKPT_FILE = os.environ.get("CKPT_FILE", "conformer-ctc-15-0.240.ckpt") CKPT_REPO = os.environ.get("CKPT_REPO", "") # optional HF model-repo id, e.g. "nsarob/kinyarwanda-asr" def resolve_ckpt(): if os.path.exists(CKPT_FILE): return CKPT_FILE if CKPT_REPO: from huggingface_hub import hf_hub_download return hf_hub_download(repo_id=CKPT_REPO, filename=CKPT_FILE) raise FileNotFoundError( f"Checkpoint '{CKPT_FILE}' not found. Put the .ckpt next to app.py " f"(drop it in this folder / the Space repo), or set the CKPT_REPO env var to a " f"Hugging Face model repo that hosts it." ) # ----- tokenizer (identical to training) ----- class CTCTokenizer: BLANK = "" UNK = "" def __init__(self): chars = list("abcdefghijklmnopqrstuvwxyz '") + list("0123456789") + list('.,?!-"') self.vocab = [self.BLANK, self.UNK] + chars self.c2i = {c: i for i, c in enumerate(self.vocab)} self.i2c = {i: c for c, i in self.c2i.items()} self.blank_id = self.c2i[self.BLANK] self.unk_id = self.c2i[self.UNK] def __len__(self): return len(self.vocab) def decode_greedy(self, ids): out, prev = [], None for i in ids: i = int(i) if i != prev and i != self.blank_id and i != self.unk_id: out.append(self.i2c[i]) prev = i return "".join(out) # ----- model (identical to training) ----- class Subsample4x(nn.Module): def __init__(self, in_feats, d_model): super().__init__() self.conv = nn.Sequential( nn.Conv2d(1, d_model, 3, 2), nn.ReLU(), nn.Conv2d(d_model, d_model, 3, 2), nn.ReLU(), ) f = in_feats f = (f - 3) // 2 + 1 f = (f - 3) // 2 + 1 self.proj = nn.Linear(d_model * f, d_model) def forward(self, x, lengths): x = self.conv(x.unsqueeze(1)) B, d, Tn, Fn = x.shape x = x.permute(0, 2, 1, 3).contiguous().view(B, Tn, d * Fn) x = self.proj(x) new_lens = lengths for _ in range(2): new_lens = ((new_lens - 3) // 2 + 1).clamp(min=1) return x, new_lens class SinusoidalPE(nn.Module): def __init__(self, d_model, max_len=8000, dropout=0.1): super().__init__() pe = torch.zeros(max_len, d_model) pos = torch.arange(max_len, dtype=torch.float32).unsqueeze(1) div = torch.exp(torch.arange(0, d_model, 2, dtype=torch.float32) * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(pos * div) pe[:, 1::2] = torch.cos(pos * div) self.register_buffer("pe", pe, persistent=False) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.dropout(x + self.pe[: x.size(1)]) class FeedForward(nn.Module): def __init__(self, d_model, ff_mult, dropout=0.1): super().__init__() self.net = nn.Sequential( nn.LayerNorm(d_model), nn.Linear(d_model, d_model * ff_mult), nn.SiLU(), nn.Dropout(dropout), nn.Linear(d_model * ff_mult, d_model), nn.Dropout(dropout), ) def forward(self, x): return self.net(x) class ConvModule(nn.Module): def __init__(self, d_model, kernel, dropout=0.1): super().__init__() assert kernel % 2 == 1 self.ln = nn.LayerNorm(d_model) self.pw1 = nn.Conv1d(d_model, 2 * d_model, 1) self.glu = nn.GLU(dim=1) self.dw = nn.Conv1d(d_model, d_model, kernel, padding=(kernel - 1) // 2, groups=d_model) self.bn = nn.BatchNorm1d(d_model) self.act = nn.SiLU() self.pw2 = nn.Conv1d(d_model, d_model, 1) self.drop = nn.Dropout(dropout) def forward(self, x): y = self.ln(x).transpose(1, 2) y = self.pw2(self.act(self.bn(self.dw(self.glu(self.pw1(y)))))) return self.drop(y).transpose(1, 2) class MHSA(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() self.ln = nn.LayerNorm(d_model) self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.drop = nn.Dropout(dropout) def forward(self, x, key_padding_mask=None): h = self.ln(x) h, _ = self.attn(h, h, h, key_padding_mask=key_padding_mask, need_weights=False) return self.drop(h) class ConformerBlock(nn.Module): def __init__(self, d_model, n_heads, ff_mult, conv_kernel, dropout=0.1): super().__init__() self.ff1 = FeedForward(d_model, ff_mult, dropout) self.mhsa = MHSA(d_model, n_heads, dropout) self.conv = ConvModule(d_model, conv_kernel, dropout) self.ff2 = FeedForward(d_model, ff_mult, dropout) self.ln_out = nn.LayerNorm(d_model) def forward(self, x, key_padding_mask=None): x = x + 0.5 * self.ff1(x) x = x + self.mhsa(x, key_padding_mask) x = x + self.conv(x) x = x + 0.5 * self.ff2(x) return self.ln_out(x) class ConformerCTC(nn.Module): def __init__(self, n_mels, vocab_size, d_model, n_heads, n_layers, ff_mult, conv_kernel, dropout=0.1, **_): super().__init__() self.subsample = Subsample4x(n_mels, d_model) self.pe = SinusoidalPE(d_model, dropout=dropout) self.blocks = nn.ModuleList([ ConformerBlock(d_model, n_heads, ff_mult, conv_kernel, dropout) for _ in range(n_layers) ]) self.head = nn.Linear(d_model, vocab_size) def forward(self, mels, input_lens): x = mels.transpose(1, 2) x, lens = self.subsample(x, input_lens) x = self.pe(x) mask = torch.arange(x.size(1), device=x.device)[None, :] >= lens[:, None] for blk in self.blocks: x = blk(x, key_padding_mask=mask) return self.head(x), lens # ----- load model + features ----- tok = CTCTokenizer() _ckpt_path = resolve_ckpt() print(f"loading checkpoint: {_ckpt_path}") _ckpt = torch.load(_ckpt_path, map_location="cpu", weights_only=False) _state = _ckpt["state_dict"] if "state_dict" in _ckpt else _ckpt _model_sd = {k[len("model."):]: v for k, v in _state.items() if k.startswith("model.")} model = ConformerCTC(N_MELS, len(tok), **PRESET) model.load_state_dict(_model_sd, strict=False) model.eval() melspec = T.MelSpectrogram(sample_rate=SR, n_fft=N_FFT, win_length=N_FFT, hop_length=HOP, n_mels=N_MELS, power=2.0) def transcribe(audio): if audio is None: return "(no audio)" try: # Gradio (type='numpy') hands us (sample_rate, np.ndarray) already decoded, # so we avoid torchaudio.load (which needs the torchcodec backend on torch 2.8+). sr, data = audio data = np.asarray(data) is_int = np.issubdtype(data.dtype, np.integer) maxv = float(np.iinfo(data.dtype).max) if is_int else 1.0 if data.ndim == 2: # (samples, channels) -> mono data = data.mean(axis=1) wav = torch.from_numpy(data.astype("float32")) if is_int: wav = wav / maxv # PCM int -> [-1, 1] if sr != SR: wav = torchaudio.functional.resample(wav, sr, SR) wav = wav[: SR * MAX_SECS] # cap length mel = torch.log(melspec(wav) + 1e-6) mel = (mel - mel.mean(dim=1, keepdim=True)) / (mel.std(dim=1, keepdim=True) + 1e-5) x = mel.unsqueeze(0) lens = torch.tensor([mel.shape[1]], dtype=torch.long) with torch.no_grad(): logits, out_lens = model(x, lens) ids = logits.argmax(dim=-1)[0, : int(out_lens[0])].tolist() text = tok.decode_greedy(ids).strip() return text or "(silence / nothing recognized)" except Exception as e: return f"error: {type(e).__name__}: {e}" DESCRIPTION = """ # 🎙️ Kinyarwanda ASR — trained from scratch A ~20M-parameter Conformer-CTC speech recognizer, trained **from scratch** (no pretrained weights) on Kinyarwanda speech. Record or upload audio to get a transcript. ⚠️ **Research PoC.** Trained on a single in-domain corpus (image-description speech, one 2025 campaign). Spontaneous, noisy, or health-domain speech is out-of-distribution and may transcribe poorly — that gap is expected and is the focus of ongoing work. """ demo = gr.Interface( fn=transcribe, inputs=gr.Audio(sources=["microphone", "upload"], type="numpy", label="Kinyarwanda audio"), outputs=gr.Textbox(label="Transcript", lines=3), title="Kinyarwanda ASR (from scratch)", description=DESCRIPTION, flagging_mode="never", ) if __name__ == "__main__": demo.launch()