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

model_path = hf_hub_download(
    repo_id="ahmando/TextClassifier",
    filename="model.pkl"
)
learn = load_learner(model_path)

label_names = {0: 'World', 1: 'Sports', 2: 'Business', 3: 'Sci/Tech'}

def classify_news(text):
    pred_label, pred_idx, probs = learn.predict(text)
    return {label_names[i]: float(probs[i]) for i in range(len(probs))}

demo = gr.Interface(
    fn=classify_news,
    inputs=gr.Textbox(placeholder="Paste a news headline here...", lines=3, label="News text"),
    outputs=gr.Label(num_top_classes=4, label="Category"),
    title="AG News Classifier",
    description="Classifies news into World, Sports, Business or Sci/Tech using ULMFiT (FastAI).",
    examples=[
        ["NASA launches new Mars rover mission to search for signs of ancient life"],
        ["PSG wins Champions League after dramatic penalty shootout"],
        ["Federal Reserve raises interest rates amid inflation concerns"],
        ["United Nations calls emergency meeting over escalating conflict"],
    ]
)

demo.launch()