Update app.py

app.py

@@ -1,44 +1,47 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import BertTokenizer, BertForSequenceClassification
+import zipfile
+import os
 
-# Load the zero-shot classifier for bias detection using Facebook's BART MNLI.
-classifier = pipeline(
-    "zero-shot-classification",
-    model="facebook/bart-large-mnli"
-)
+# ✅ 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")
 
-def process_text(text):
-    # Define candidate labels for bias classification.
-    candidate_labels = ["biased", "neutral"]
-    
-    # Run zero-shot classification.
-    classification = classifier(text, candidate_labels)
-    detected_bias = classification["labels"][0]
-    confidence = classification["scores"][0]
+# ✅ 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 the results.
     return {
-        "Detected Bias": detected_bias,
-        "Confidence": round(confidence, 2),
+        "Predicted Class": label_map[pred_label],
+        "Confidence": confidence
     }
 
-# Build the Gradio UI.
+# ✅ Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Bias Bin")
-    gr.Markdown(
-        "Detect gender stereotypes in narrative text. "
-        "Enter a story or sentence below and click the **Submit** button."
-    )
-    
-    text_input = gr.Textbox(
-        label="Enter Story Text",
-        placeholder="Type a story or sentence here...",
-        lines=5
-    )
+    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.")
     
-    submit_btn = gr.Button("Submit")
-    result_output = gr.JSON(label="Output")
-    
-    submit_btn.click(fn=process_text, inputs=[text_input], outputs=[result_output])
-    
-demo.launch()
+    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()