File size: 3,364 Bytes
e4c97d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """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 # noqa: F401 (registers omniASR cards)
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() # base weights, then overwritten
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)}")
|