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

MODEL_PATH = "./model"

def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=False)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
    model.eval()

    with open(f"{MODEL_PATH}/label_mapping.json") as f:
        label_mapping = json.load(f)

    return tokenizer, model, label_mapping

tokenizer, model, label_mapping = load_model()

def predict(text: str):
    if not text.strip():
        return {}

    inputs = tokenizer(
        text,
        return_tensors="pt",
        padding="max_length",
        truncation=True,
        max_length=128,
    )
    with torch.no_grad():
        logits = model(**inputs).logits

    probs = torch.softmax(logits, dim=-1).squeeze()
    scores = {label_mapping[str(i)]: float(probs[i]) for i in range(len(probs))}
    return scores

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Write something",
        label="Example Tweet"
    ),
    outputs=gr.Label(
        num_top_classes=len(label_mapping),
        label="Predicted Label"
    ),
    title="Radical Speech Detector",
    description="Enter a text and the model will classify it into one of the radical speech categories.\n\n\nBy Sofia Maldonado, Viviana Toledo & Oscar Rocha",
)

demo.launch()