import glob,json,os,collections,soundfile as sf,torch
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()
i2l=m.config.id2label
files=sorted(glob.glob("/scratch/p2_16k/*.wav"))
gpu={}; conf={}
with torch.inference_mode():
    for i in range(0,len(files),8):
        b=files[i:i+8]
        au=[sf.read(f,dtype="float32")[0][:16000*20] for f 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 f,p in zip(b,pr):
            k=int(p.argmax()); lab=i2l[k] if k in i2l else i2l[str(k)]
            if lab=="lug": lab="lin"
            i2=os.path.splitext(os.path.basename(f))[0]
            gpu[i2]=lab; conf[i2]=float(p.max())
lex=json.load(open("/root/test_lang.json"))
print("GPU LID :",dict(collections.Counter(gpu.values())))
print("lexical :",dict(collections.Counter(lex.values())))
dis=[i for i in gpu if gpu[i]!=lex.get(i)]
print("DESACCORDS :",len(dis))
for i in dis: print("   ",i,"gpu=",gpu[i],"(conf %.3f)"%conf[i],"lex=",lex.get(i))
json.dump({"gpu":gpu,"conf":conf,"disagree":dis},open("/root/lid_test.json","w"))
print("LID_TEST_DONE")