import gradio as gr
from transformers import pipeline

pipe = pipeline("text-classification", model="kitrofimov/news-clf", top_k=3)
label_names = dict(zip([f"LABEL_{i}" for i in range(4)], ["World", "Sports", "Business", "Sci/Tech"]))
examples = [
    ["NASA announces new discovery on Mars water."],
    ["Stock markets rally after positive earnings report."],
    ["Lionel Messi scores a hat-trick in the Champions League."],
    ["UN summit discusses climate change and global policies."]
]

def classify(text):
    preds = pipe(text)[0]
    return {label_names[p["label"]]: float(p["score"]) for p in preds}

with gr.Blocks() as demo:
    gr.Markdown("# News Classifier")
    gr.Markdown("Paste a news article below and see which category it belongs to! (one of \"world\", \"sports\", \"business\" and \"science / technology\")")
    gr.Markdown("This model is based on a [`distilbert-base-uncased`](https://huggingface.co/distilbert/distilbert-base-uncased) architecture and was fine-tuned on the [AG News](https://huggingface.co/datasets/fancyzhx/ag_news) dataset for 3 epochs. Training code [here](https://colab.research.google.com/drive/1KTai0S1dzwIoS3Sba_jJG9ZNISRjSKGo)")
    
    with gr.Row():
        with gr.Column():
            input = gr.Textbox(lines=5, placeholder="Enter your news article...")
            gr.Examples(examples=examples, inputs=input)
            classify_btn = gr.Button("Classify")
        with gr.Column():
            output = gr.Label(num_top_classes=3)

    classify_btn.click(classify, inputs=input, outputs=output)

demo.launch()