"""Gradio demo for Industry Classification."""

import json
from pathlib import Path

import gradio as gr
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Load model (HF Spaces: model files in root)
MODEL_PATH = Path(".")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
model.eval()

# Load label names
with open(MODEL_PATH / "taxonomy.json") as f:
    taxonomy = json.load(f)
id_to_name = {cat["id"]: cat["name"] for cat in taxonomy["categories"]}


def classify(text: str) -> dict:
    """Classify industry text."""
    if not text.strip():
        return {}

    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=64)

    with torch.no_grad():
        probs = torch.softmax(model(**inputs).logits, dim=-1)[0]

    top_probs, top_idx = torch.topk(probs, 5)

    return {
        f"{id_to_name.get(model.config.id2label[i.item()], '?')}": float(p)
        for p, i in zip(top_probs, top_idx)
    }


demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(label="Industry", placeholder="e.g. software development"),
    outputs=gr.Label(label="GICS Classification"),
    examples=["software development", "investment banking", "oil and gas", "retail stores", "pharmaceuticals"],
    title="Industry Classification",
    description="Classify text into GICS industries using fine-tuned DistilBERT.",
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()