| """Streaming-capable ASR worker (runs in ~/venv_omni): Meta Omnilingual ASR CTC (Apache-2.0), language-agnostic decoding, |
| so code-switched speech needs no language hint. POST /asr body = raw float32 mono 16 kHz -> {"text", "ms"}. |
| Partial hypotheses = re-decoding the growing buffer (CTC is cheap: ~10-40 ms per call on H100).""" |
| import os, time, numpy as np, torch |
| import soundfile |
| from fastapi import FastAPI, Request |
| from omnilingual_asr.models.inference.pipeline import ASRInferencePipeline |
| |
| |
| |
| |
| |
| |
| |
| LC = {"ibo": "ibo_Latn", "yor": "yor_Latn", "hau": "hau_Latn", "pcm": "pcm_Latn", "eng": "eng_Latn"} |
| import unicodedata, types |
| import omnilingual_asr.models.inference.pipeline as PL |
| DYN = {"dynamic-asr": "real/model.pt", "dynamic-asr-cs": "codeswitch-experimental/model.pt"} |
| MODEL = os.environ.get("ASR_MODEL", "dynamic-asr") |
| CKPT = os.environ.get("ASR_CKPT", "") |
| CARD = os.environ.get("ASR_CARD", "omniASR_CTC_1B_v2" if (MODEL in DYN or CKPT) else "omniASR_CTC_1B") |
|
|
| def _load(): |
| global CARD, CKPT, MODEL |
| if MODEL in DYN: |
| CARD = "omniASR_CTC_1B_v2" |
| if MODEL in DYN and not CKPT: |
| try: |
| from huggingface_hub import hf_hub_download |
| CKPT = hf_hub_download(os.environ.get("ASR_REPO", "Modularcomputing/dynamic-asr"), DYN[MODEL]) |
| except Exception as e: |
| print("Dynamic ASR download failed, falling back to omniASR_CTC_1B:", repr(e)[:200], flush=True) |
| CARD, MODEL, CKPT = "omniASR_CTC_1B", "omni", "" |
| try: |
| pipe = ASRInferencePipeline(model_card=CARD) |
| if CKPT: |
| ck = torch.load(CKPT, map_location="cpu", weights_only=False) |
| print("ASR_CKPT", CKPT, pipe.model.load_state_dict({k: v.to(pipe.dtype) for k, v in ck["model"].items()}), flush=True) |
| del ck |
| return pipe |
| except Exception as e: |
| if CARD == "omniASR_CTC_1B": |
| raise |
| print("Dynamic ASR load failed (", repr(e)[:200], ") -> falling back to omniASR_CTC_1B", flush=True) |
| CARD, MODEL, CKPT = "omniASR_CTC_1B", "omni", "" |
| return ASRInferencePipeline(model_card=CARD) |
|
|
| p = _load() |
| LATIN_ONLY = os.environ.get("ASR_LATIN_ONLY", "1") == "1" |
|
|
| def _ok(sx): |
| for c in sx: |
| 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 |
|
|
| _BAD = {} |
| def _apply_latin(self, batch): |
| """CTC greedy decoding restricted to Latin-script tokens: all five target languages are written in Latin script, and the |
| language-agnostic CTC otherwise drifts into Arabic/other scripts on accented or synthetic speech.""" |
| bl = PL.BatchLayout(batch.source_seqs.shape, seq_lens=batch.source_seq_lens, device=batch.source_seqs.device) |
| logits, bl_out = self.model(batch.source_seqs, bl) |
| V = logits.shape[-1] |
| if V not in _BAD: |
| bad = [i for i in range(V) if not _ok(self.token_decoder(torch.tensor([i])))] |
| _BAD[V] = torch.tensor(bad, dtype=torch.long, device=logits.device); print("latin mask: blocked", len(bad), "of", V, flush=True) |
| logits[..., _BAD[V]] = -float("inf") |
| pred = torch.argmax(logits, dim=-1); out = [] |
| for i in range(pred.shape[0]): |
| seq = pred[i][: bl_out.seq_lens[i]]; m = torch.ones(seq.shape[0], dtype=torch.bool, device=seq.device); m[1:] = seq[1:] != seq[:-1] |
| out.append(self.token_decoder(seq[m])) |
| return out |
|
|
| if LATIN_ONLY and "CTC" in CARD: |
| p._apply_model_wav2vec2asr = types.MethodType(_apply_latin, p) |
| app = FastAPI() |
| w = np.zeros(16000, np.float32); p.transcribe([{"waveform": w, "sample_rate": 16000}], batch_size=1) |
|
|
| @app.post("/asr") |
| async def asr(req: Request, lang: str = ""): |
| y = np.frombuffer(await req.body(), dtype=np.float32).copy() |
| if len(y) < 1600: |
| return {"text": "", "ms": 0.0} |
| t = time.perf_counter() |
| kw = {} if "CTC" in CARD or not lang else dict(lang=[LC.get(lang, "eng_Latn")]) |
| h = p.transcribe([{"waveform": y, "sample_rate": 16000}], batch_size=1, **kw)[0] |
| return {"text": h, "ms": round(1000 * (time.perf_counter() - t), 1), "card": CARD} |
|
|
| @app.get("/health") |
| def health(): |
| return {"ok": True, "model": MODEL, "card": CARD, "ckpt": CKPT} |
|
|
