import gradio as gr
import torch
from transformers import BertTokenizer, BertForSequenceClassification
import zipfile
import os

# ✅ Unzip the fine-tuned model if it's not already extracted
if not os.path.exists("fine_tuned_model"):
    with zipfile.ZipFile("fine_tuned_model.zip", 'r') as zip_ref:
        zip_ref.extractall("fine_tuned_model")

# ✅ Load your fine-tuned model and tokenizer
model_path = "./fine_tuned_model"
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path)
model.eval()

# ✅ Define label mapping (adjust based on your labels)
label_map = {0: "Original-like", 1: "Swapped-like"}

# ✅ Inference function
def detect_bias(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=1).squeeze()
        pred_label = torch.argmax(probs).item()
        confidence = round(probs[pred_label].item(), 2)
    
    return {
        "Predicted Class": label_map[pred_label],
        "Confidence": confidence
    }

# ✅ Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Bias Bin – Fine-Tuned BERT Version")
    gr.Markdown("This interface uses a fine-tuned BERT model to classify gender bias in narrative text.")
    
    text_input = gr.Textbox(label="Enter Narrative Text", lines=4, placeholder="Type here...")
    submit_btn = gr.Button("Detect Bias")
    output = gr.JSON(label="Output")

    submit_btn.click(fn=detect_bias, inputs=text_input, outputs=output)

demo.launch()