from huggingface_hub import from_pretrained_fastai
import gradio as gr
from fastai.text.all import *

# repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
repo_id = "aribanez/ag-news-classifier"
learner = from_pretrained_fastai(repo_id)
labels = learner.dls.vocab[1] # ['World', 'Sports', 'Business', 'Sci/Tech']

mapping = {
    0: "World",
    1: "Sports",
    2: "Business",
    3: "Sci/Tech"
}

def predict(text):
    pred, pred_idx, probs = learner.predict(text)
    return {mapping[i]: float(probs[i]) for i in range(len(mapping))}

with open('examples.txt', 'r', encoding='utf-8') as f:
    examples = f.readlines()

gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=4, placeholder="Escribe la cabecera de una noticia o artículo..."),
    outputs=gr.Label(num_top_classes=4),
    title="Clasificador AG-News",
    description="Clasifica noticias entre World, Sports, Business o Sci/Tech utilizando ULMFiT + FastAI.",
    examples=examples,
    examples_per_page=len(examples),
    cache_examples=False
).launch()