# Importación de las bibliotecas necesarias
from huggingface_hub import from_pretrained_fastai
import gradio as gr

# Identificador del repositorio en Hugging Face donde está almacenado el modelo
repo_id = "luis56125/news"

# Cargar el modelo preentrenado desde Hugging Face
learner = from_pretrained_fastai(repo_id)

# Definir las etiquetas de clasificación disponibles
labels = ['Mundo', 'Deportes', 'Negocios', 'Ciencia/Tecnología']

# Definir una función para predecir la categoría de un texto dado
def predict(text):
    # Obtener las probabilidades de las etiquetas desde el modelo
    probs = learner.predict(text)[-1]
    # Devolver un diccionario que mapea cada etiqueta a su probabilidad correspondiente
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# Crear una interfaz de usuario para el modelo utilizando Gradio
gr.Interface(fn=predict, inputs="text", outputs=gr.components.Label(num_top_classes=4)).launch(share=False)