# app.py — Beautiful web UI with Gradio

import gradio as gr
from transformers import pipeline

MODEL_PATH = "fabriceyhc/bert-base-uncased-ag_news"

LABELS = {
    "LABEL_0": "🌍 World News",
    "LABEL_1": "⚽ Sports",
    "LABEL_2": "💼 Business",
    "LABEL_3": "💻 Sci / Tech"
}

# Load model once
print("Loading model...")
classifier = pipeline(
    "text-classification",
    model=MODEL_PATH,
    return_all_scores=True
)

def classify_article(text):
    if not text.strip():
        return {}, "Please enter some text!"

    results = classifier(text)

    # Fix: handle both list-of-list and list-of-dict formats
    if isinstance(results[0], list):
        results = results[0]

    # Build confidence dict for Gradio bar chart
    scores = {LABELS[r["label"]]: round(r["score"], 4) for r in results}

    best = max(results, key=lambda x: x["score"])
    label = LABELS[best["label"]]
    conf  = best["score"] * 100
    verdict = f"**{label}** — {conf:.1f}% confidence"

    return scores, verdict

# ── Build UI ──
with gr.Blocks(title="📰 News Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📰 News Article Classifier")
    gr.Markdown("Powered by **DistilBERT** fine-tuned on AG News dataset")

    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="Paste your news article here",
                placeholder="e.g. Apple revealed its new MacBook Pro...",
                lines=6
            )
            classify_btn = gr.Button("🔍 Classify", variant="primary")

            gr.Examples(
                examples=[
                    ["NASA launched a new Mars rover to study ancient riverbeds on the planet's surface."],
                    ["The stock market fell sharply after the Federal Reserve raised interest rates."],
                    ["Lionel Messi scored a hat-trick to lead Argentina to victory in the World Cup final."],
                    ["Scientists discovered a new method to generate clean energy from seawater."],
                ],
                inputs=text_input,
                label="📋 Try these examples"
            )

        with gr.Column(scale=1):
            verdict_out = gr.Markdown(label="Result")
            scores_out  = gr.Label(label="Confidence Scores", num_top_classes=4)

    classify_btn.click(
        fn=classify_article,
        inputs=text_input,
        outputs=[scores_out, verdict_out]
    )
    text_input.submit(
        fn=classify_article,
        inputs=text_input,
        outputs=[scores_out, verdict_out]
    )

demo.launch(share=True)  # share=True gives a public link!