import json, soundfile as sf, torch, collections
from transformers import AutoModelForAudioClassification, AutoFeatureExtractor
M="/root/models/lid_best"
fe=AutoFeatureExtractor.from_pretrained(M)
m=AutoModelForAudioClassification.from_pretrained(M,dtype=torch.bfloat16).cuda().eval()
id2l=m.config.id2label
ROWS=[json.loads(l) for l in open("/root/devhard/devhard_linsna.jsonl")]
pred={}; conf={}
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][:16000*20] for r in b]
        x=fe(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()}
        pr=m(**x).logits.float().softmax(-1)
        for r,p in zip(b,pr):
            k=int(p.argmax()); pred[r["id"]]=id2l[k] if not isinstance(id2l,dict) else id2l[str(k)] if str(k) in id2l else id2l[k]
            conf[r["id"]]=float(p.max())
ok=sum(1 for r in ROWS if pred[r["id"]]==r["lang"])
print(f"LID accuracy devhard: {ok}/{len(ROWS)} = {100*ok/len(ROWS):.2f}%")
cm=collections.Counter((r["lang"],pred[r["id"]]) for r in ROWS)
print("confusions:",dict(cm))
# accuracy par seuil de confiance
for t in [0.0,0.9,0.95,0.99]:
    sub=[r for r in ROWS if conf[r["id"]]>=t]
    o=sum(1 for r in sub if pred[r["id"]]==r["lang"])
    print(f"  conf>={t}: n={len(sub)} acc={100*o/max(len(sub),1):.2f}%")
json.dump({"pred":pred,"conf":conf},open("/root/lid_devhard.json","w"))
print("LID_DEVHARD_DONE")