from speechbrain.pretrained.interfaces import foreign_class
import gradio as gr

import warnings
warnings.filterwarnings("ignore")

# Loading the speechbrain emotion detection model
learner = foreign_class(
    source="speechbrain/emotion-recognition-wav2vec2-IEMOCAP",
    # savedir="/home/harish3110/SeaWord/emotion/nbs/pretrained_models/CustomEncoderWav2vec2Classifier--8353113631630090076",
    pymodule_file="custom_interface.py", 
    classname="CustomEncoderWav2vec2Classifier"
)

# Building prediction function for gradio
emotion_dict = {
    'sad': 'Sad', 
    'hap': 'Happy',
    'ang': 'Anger',
    'neu': 'Neutral'
}

def predict_emotion(audio):
    out_prob, score, index, text_lab = learner.classify_file(audio.name)
    return emotion_dict[text_lab[0]]

# Loading gradio interface
inputs = gr.inputs.Audio(label="Input Audio", type="file")
outputs = "text"
title = "Emotion Detection"
description = "Gradio demo for Emotion Detection. To use it, simply upload your audio, or click one of the examples to load them. Read more at the links below."
gr.Interface(predict_emotion, inputs, outputs, title=title, description=description).launch()