app.py

import gradio as gr
from transformers import pipeline

pipe = pipeline("text-classification", model="kitrofimov/news-classifier", top_k=3)

def classify(text):
    preds = pipe(text)[0]
    return {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...")
            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()