import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
from peft import PeftModel

# ✅ Load model and tokenizer
base_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=9)
model = PeftModel.from_pretrained(base_model, "NightPrince/peft-distilbert-toxic-classifier")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

# ✅ Label mapping
id2label = {
    0: "Child Sexual Exploitation",
    1: "Elections",
    2: "Non-Violent Crimes",
    3: "Safe",
    4: "Sex-Related Crimes",
    5: "Suicide & Self-Harm",
    6: "Unknown S-Type",
    7: "Violent Crimes",
    8: "unsafe"
}
# ✅ Pipeline for easy inference
pipe = pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    return_all_scores=True
)

# ✅ Define prediction function
def classify_toxicity(query, image_description):
    combined_text = query + " [SEP] " + image_description
    preds = pipe(combined_text)[0]  # Get scores for all classes
    preds_sorted = sorted(preds, key=lambda x: x['score'], reverse=True)
    
    top_label = preds_sorted[0]['label']
    top_score = preds_sorted[0]['score']
    
    # Map label ID back to human-readable label
    label_id = int(top_label.split("_")[-1]) if "_" in top_label else int(top_label)
    final_label = id2label.get(label_id, "Unknown")
    
    # Display all class scores (optional)
    scores_table = "\n".join(
        [f"{id2label[int(item['label'].split('_')[-1])]}: {round(item['score']*100, 2)}%" for item in preds]
    )
    
    return f"Top Prediction: {final_label} ({round(top_score*100, 2)}%)\n\nFull Class Scores:\n{scores_table}"

# ✅ Gradio UI
iface = gr.Interface(
    fn=classify_toxicity,
    inputs=[
        gr.Textbox(label="User Query"),
        gr.Textbox(label="Image Description"),
    ],
    outputs=gr.Textbox(label="Toxicity Prediction"),
    title="Toxic Category Classifier (DistilBERT + LoRA)",
    description="Enter a user query and image description. The model will classify into one of the 9 toxic categories."
)

iface.launch()