"""Whisper candidates (MIT) via transformers: ~/venv/bin/python asr_whisper.py <hf_id> <tag>"""
import json, sys, time, torch, soundfile as sf
from transformers import pipeline
mid, tag = sys.argv[1], sys.argv[2]; D = "/scratch/duplex/asr_eval"
WL = {"eng": "en", "yor": "yo", "hau": "ha", "pcm": "en", "ibo": None, "cs": None}
asr = pipeline("automatic-speech-recognition", model=mid, torch_dtype=torch.float16, device="cuda")
R = [json.loads(l) for l in open(f"{D}/list.jsonl")]
out = open(f"{D}/hyp_{tag}.jsonl", "w")
for i, r in enumerate(R):
    w = sf.read(r["wav"], dtype="float32")[0]
    gk = {"task": "transcribe"}
    if WL[r["lang"]]: gk["language"] = WL[r["lang"]]
    torch.cuda.synchronize(); t = time.perf_counter()
    h = asr({"raw": w, "sampling_rate": 16000}, generate_kwargs=gk)["text"].strip()
    torch.cuda.synchronize(); dt = time.perf_counter() - t
    out.write(json.dumps(dict(id=r["id"], hyp=h, lat=round(dt, 4), dur=r["dur"]), ensure_ascii=False) + "\n"); out.flush()
print("WHISPER_DONE", tag)