from huggingface_hub import from_pretrained_fastai
import gradio as gr
# from fastai.vision.all import *
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch


# repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
repo_id = "islasher/clasificador-dair-emotion"

labels = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']

# Definimos una función que se encarga de llevar a cabo las predicciones


# Cargar el modelo y el tokenizador
nombre_modelo = "islasher/clasificador-dair-emotion"
model = AutoModelForSequenceClassification.from_pretrained(nombre_modelo)
tokenizer = AutoTokenizer.from_pretrained(nombre_modelo)


def predict(frase):
    #img = PILImage.create(img)
    inputs = tokenizer(frase, return_tensors="pt")
    outputs = model(**inputs)
    predicted_class = torch.argmax(outputs.logits, dim=1).item()
    return labels[predicted_class]
    
# Creamos la interfaz y la lanzamos. 
gr.Interface(fn=predict, inputs="text", outputs="text").launch(share=False)