import torchaudio
import soundfile as sf
import torch
from transformers import AutoProcessor, AutoModelForAudioClassification
from transformers import AutoFeatureExtractor

# Load model directly
feature_extractor = AutoFeatureExtractor.from_pretrained("ThomasR/facebook_wav2vec2-large_October_03_2023_05h34PM")
model = AutoModelForAudioClassification.from_pretrained("ThomasR/facebook_wav2vec2-large_October_03_2023_05h34PM")

# Label dict
label2id={'fake':0, 'real':1}
id2label = {v:k for k,v in label2id.items()}

#Inference function
def predict(audio_path):
    wavform, sample_rate = sf.read(audio_path)
    
    inputs = feature_extractor(
            wavform, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt", max_length=16000, truncation=True, padding=True
        )
    
    with torch.no_grad():
        logits = model(**inputs).logits
    
    probabilities = torch.sigmoid(logits[0])
    # labels is a one-hot array of shape (num_frames, num_speakers)
    labels = (probabilities > 0.5).long()
    pred_probs = list(probabilities.tolist())
    LABELS=list(id2label.values())
    pred_labels_and_probs = {LABELS[i]: round(float(pred_probs[i]),4) for i in range(len(LABELS))}
    
    return pred_labels_and_probs


audio_path = 'fake_Elevenlabs_common_voice_en_36808626_Harry.wav'
predict(audio_path)