import torch
import torch.nn.functional as F
from transformers import AutoConfig, Wav2Vec2FeatureExtractor
from src.models import Wav2Vec2ForSpeechClassification  #imported from https://github.com/m3hrdadfi/soxan
import gradio as gr
import librosa

device = torch.device("cpu")
model_name_or_path = "harshit345/xlsr-wav2vec-speech-emotion-recognition"
config = AutoConfig.from_pretrained(model_name_or_path)
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name_or_path)
sampling_rate = feature_extractor.sampling_rate
model = Wav2Vec2ForSpeechClassification.from_pretrained(model_name_or_path)

#load input file and resample to 16kHz
def load_data(path):
    speech, sampling_rate = librosa.load(path)
    if len(speech.shape) > 1:
        speech = speech[:,0] + speech[:,1]
    if sampling_rate != 16000:
        speech = librosa.resample(speech, sampling_rate,16000)
    return speech

#modified version of predict function from https://github.com/m3hrdadfi/soxan
def inference(path):
    speech = load_data(path)
    inputs = feature_extractor(speech, return_tensors="pt").input_values
    with torch.no_grad():
        logits = model(inputs).logits
    scores = F.softmax(logits, dim=1).detach().cpu().numpy()[0]
    outputs = {config.id2label[i]: float(round(score,2)) for i, score in enumerate(scores)}
    return outputs

inputs = gr.inputs.Audio(label="Input Audio", type="filepath", source="microphone")
outputs = gr.outputs.Label(type="confidences", label = "Output Scores")
title = "Wav2Vec2 Speech Emotion Recognition"
description = "This is a demo of the Wav2Vec2 Speech Emotion Recognition model. Record an audio file and the top emotions inferred will be displayed."
examples = ['data/heart.wav', 'data/happy26.wav', 'data/jm24.wav', 'data/newton.wav', 'data/speeding.wav']
article = "<a href = 'https://github.com/m3hrdadfi/soxan'> Wav2Vec2 Speech Classification Github Repository"


iface = gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, theme="peach", examples=examples)
iface.launch(debug=True)