from reader import get_article
import gradio as gr
from transformers import pipeline

info = get_article()


#model_name2id = {"Model A": "hackathon-pln-es/wav2vec2-base-finetuned-sentiment-classification-MESD", "Model B": "hackathon-pln-es/wav2vec2-base-finetuned-sentiment-mesd"}

def classify_sentiment(audio):
  pipe = pipeline("audio-classification", model="hackathon-pln-es/wav2vec2-base-finetuned-sentiment-classification-MESD")
  pred = pipe(audio)
  return {dic["label"]: dic["score"] for dic in pred}

input_audio = [gr.Audio(sources=["microphone"], type="filepath", label="Record/ Drop audio")]
label = gr.Label(num_top_classes=5)

################### Gradio Web APP ################################

description = """ Gradio demo for Sentiment Classification of Spanish audios using Wav2Vec2. We have fine-tuned two models for the aforementioned task i) [Model A](https://huggingface.co/hackathon-pln-es/wav2vec2-base-finetuned-sentiment-classification-MESD) ii) [Model B](https://huggingface.co/hackathon-pln-es/wav2vec2-base-finetuned-sentiment-mesd)

Link to Processed MESD dataset: [Please Click Here](https://huggingface.co/datasets/hackathon-pln-es/MESD)

Note: The Audio examples provided for testing this app were randomly picked from the **TEST set**.
"""


# generate and launch interface
interface = gr.Interface(fn=classify_sentiment, inputs=input_audio, outputs=label, examples=[["basta_neutral.wav"], ["detras_disgust.wav"], ["mortal_sadness.wav"], ["respiracion_happiness.wav"], ["robo_fear.wav"]], article=info['article'], css=info['css'], theme='huggingface', title=info['title'], allow_flagging='never', description=description)
interface.launch() 

#info['description']