| """Dynamic ASR: transcribe Nigerian English, Pidgin, Igbo, Yoruba, Hausa and code-switched speech. |
| |
| pip install "omnilingual-asr>=0.2.0" huggingface_hub soundfile scipy |
| python transcribe.py audio.wav [more.wav ...] [--model real|codeswitch] [--no-latin-mask] |
| |
| Model: Meta omniASR_CTC_1B_v2 (Apache-2.0) fine-tuned by Modular. Greedy CTC decoding, no language hint needed. |
| Output: lowercase text without punctuation (Igbo without tone marks; Yoruba with tone marks). |
| """ |
| import argparse, unicodedata |
| import numpy as np, torch, torch.nn.functional as F |
| import soundfile as sf |
| from scipy.signal import resample_poly |
| from huggingface_hub import hf_hub_download |
| from omnilingual_asr.models.inference.pipeline import ASRInferencePipeline |
| from fairseq2.models.hub import load_model |
| from fairseq2.data.tokenizers.hub import load_tokenizer |
| from fairseq2.nn.batch_layout import BatchLayout |
|
|
| REPO = "Modularcomputing/dynamic-asr" |
| FILES = {"real": "real/model.pt", "codeswitch": "codeswitch-experimental/model.pt"} |
| BASE = "omniASR_CTC_1B_v2" |
|
|
|
|
| def latin_ok(s): |
| for c in s: |
| cat = unicodedata.category(c) |
| if cat[0] == "M": |
| if not ("\u0300" <= c <= "\u036f"): |
| return False |
| elif c.isalpha() and "LATIN" not in unicodedata.name(c, ""): |
| return False |
| return True |
|
|
|
|
| class DynamicASR: |
| def __init__(self, model="real", device="cuda" if torch.cuda.is_available() else "cpu", latin_mask=True, ckpt=None): |
| self.dev = torch.device(device) |
| dtype = torch.bfloat16 if self.dev.type == "cuda" else torch.float32 |
| self.model = load_model(BASE, device=self.dev, dtype=dtype).eval() |
| ck = torch.load(ckpt or hf_hub_download(REPO, FILES[model]), map_location="cpu", weights_only=False) |
| self.model.load_state_dict({k: v.to(dtype) for k, v in ck["model"].items()}) |
| self.dtype = dtype |
| self.dec = load_tokenizer(BASE).create_decoder(skip_special_tokens=True) |
| V = self.model.final_proj.weight.shape[0] |
| self.bad = torch.tensor([i for i in range(V) if not latin_ok(str(self.dec(torch.tensor([i]))))], device=self.dev) if latin_mask else None |
|
|
| @torch.no_grad() |
| def __call__(self, wav, sr=16000): |
| y = np.asarray(wav, dtype=np.float32) |
| if y.ndim > 1: y = y.mean(1) |
| if sr != 16000: y = resample_poly(y, 16000, sr).astype(np.float32) |
| x = torch.from_numpy(y); x = F.layer_norm(x, x.shape)[None].to(self.dev, self.dtype) |
| logits, lay = self.model(x, BatchLayout(x.shape, seq_lens=[x.shape[1]], device=self.dev)) |
| if self.bad is not None: logits[..., self.bad] = -float("inf") |
| s = logits[0, : lay.seq_lens[0]].argmax(-1); m = torch.ones_like(s, dtype=torch.bool); m[1:] = s[1:] != s[:-1] |
| return str(self.dec(s[m])) |
|
|
|
|
| if __name__ == "__main__": |
| ap = argparse.ArgumentParser(); ap.add_argument("audio", nargs="+"); ap.add_argument("--model", default="real", choices=list(FILES)) |
| ap.add_argument("--no-latin-mask", action="store_true"); ap.add_argument("--ckpt", default=None, help="local model.pt instead of the Hub") |
| a = ap.parse_args() |
| asr = DynamicASR(a.model, latin_mask=not a.no_latin_mask, ckpt=a.ckpt) |
| for f in a.audio: |
| w, sr = sf.read(f, dtype="float32"); print(f"{f}\t{asr(w, sr)}") |
|
|
