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

MODEL_NAME = "s-nlp/roberta_toxicity_classifier"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# FIX 1: Dynamically pull the exact labels the model was trained on
labels = model.config.id2label 

def classify(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    
    with torch.no_grad():
        outputs = model(**inputs)
        
    # FIX 2: Use softmax for binary/multi-class classification to get probabilities
    scores = torch.softmax(outputs.logits, dim=-1)[0].tolist()
    
    # FIX 3: Safely map the scores to the dynamic labels
    result = {labels[i]: float(scores[i]) for i in range(len(scores))}
    
    return result

demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(label="Enter English text"),
    outputs=gr.JSON(label="Toxicity scores"),
    title="English Toxicity Detection",
    description="Model: s-nlp/roberta_toxicity_classifier"
)

demo.launch()