import json, soundfile as sf, torch
from transformers import AutoModelForCTC, AutoProcessor
M="/root/models/joint_cont_best"
proc=AutoProcessor.from_pretrained(M)
model=AutoModelForCTC.from_pretrained(M,torch_dtype=torch.bfloat16).cuda().eval()
rows=[json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl")]
def norm(s): return " ".join(str(s).replace("|"," ").split())
out=[]
with torch.inference_mode():
    for i in range(0,len(rows),8):
        b=rows[i:i+8]
        au=[sf.read(r["audio"],dtype="float32")[0] for r in b]
        x=proc(au,sampling_rate=16000,return_tensors="pt",padding=True)
        x={k:v.to("cuda",dtype=torch.bfloat16 if v.dtype==torch.float32 else v.dtype) for k,v in x.items()}
        pid=model(**x).logits.float().argmax(-1).cpu().numpy()
        for r,h in zip(b,proc.batch_decode(pid)):
            out.append({"id":r["id"],"lang":r.get("lang"),"ref":r["text"],"hyp":norm(h)})
json.dump(out,open("/root/devhard_joint_hyps.json","w",encoding="utf-8"),ensure_ascii=False)
print("DECODE_DEVHARD_DONE",len(out))