Update app.py

app.py

@@ -5,49 +5,47 @@ import easyocr
 from PIL import Image
 import numpy as np
 
-# Hate Speech model (example uses base CardiffNLP + extended labels for demonstration)
-MODEL_NAME = "cardiffnlp/twitter-roberta-base-hate-multiclass-latest"
-LABELS = [
+# Hate Speech Model: Only 7 valid categories!
+HATE_MODEL_NAME = "cardiffnlp/twitter-roberta-base-hate-multiclass-latest"
+HATE_LABELS = [
     "sexism",
     "racism",
     "disability",
     "sexual_orientation",
     "religion",
-    "abusive_words",  # added label - simulation only
-    "threat",         # added label - simulation only
-    "harassment",     # added label - simulation only
-    "sarcastic",      # added label - simulation only; we'll do actual sarcasm detection via separate model
+    "other",
     "not_hate"
 ]
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+hate_tokenizer = AutoTokenizer.from_pretrained(HATE_MODEL_NAME)
+hate_model = AutoModelForSequenceClassification.from_pretrained(HATE_MODEL_NAME)
 
-# Sarcasm Detection model (example pretrained; replace with your actual sarcasm model)
-SARCASM_MODEL_NAME = "microsoft/deberta-base-sarcasm"  # example, replace if unavailable
+# Sarcasm Model: REAL public model
+SARCASM_MODEL_NAME = "mrm8488/bert-tiny-finetuned-sarcasm-detection"
+SARCASM_LABELS = ["Not Sarcastic", "Sarcastic"]
 sarcasm_tokenizer = AutoTokenizer.from_pretrained(SARCASM_MODEL_NAME)
 sarcasm_model = AutoModelForSequenceClassification.from_pretrained(SARCASM_MODEL_NAME)
 
-reader = easyocr.Reader(['en'])
+reader = easyocr.Reader(['en'], gpu=False)
 
 def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    inputs = hate_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = hate_model(**inputs)
         logits = outputs.logits
         probs = torch.nn.functional.softmax(logits, dim=-1)
         pred = torch.argmax(probs).item()
         confidence = float(probs[0][pred])
-        return LABELS[pred], confidence
+        return HATE_LABELS[pred], confidence
 
 def is_sarcastic(text):
     inputs = sarcasm_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = sarcasm_model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        # assuming label 1 means sarcastic; adjust if needed
-        sarcasm_prob = probs[0][1].item()
-        return sarcasm_prob > 0.5, sarcasm_prob
+        pred = torch.argmax(probs).item()
+        sarcasm_prob = float(probs[0][pred])
+        return SARCASM_LABELS[pred], sarcasm_prob
 
 def ocr_extract(image):
     if isinstance(image, Image.Image):
@@ -56,27 +54,24 @@ def ocr_extract(image):
     return ' '.join(result)
 
 def chatbot(image=None, text=None):
-    # Priority: image with OCR, else text box
     if image is not None:
         extracted = ocr_extract(image)
         if not extracted.strip():
             return "No text found in image.", None, None
         label, confidence = classify_text(extracted)
-        sarcastic, sarcasm_prob = is_sarcastic(extracted)
-        sarcasm_text = "Yes" if sarcastic else "No"
+        sarcasm, sarcasm_prob = is_sarcastic(extracted)
         return (
-            f"OCR Extracted: {extracted}\nPrediction: {label} (Confidence: {confidence:.2f})\nSarcasm: {sarcasm_text} (Prob: {sarcasm_prob:.2f})",
+            f"OCR Extracted: {extracted}\nPrediction: {label} (Confidence: {confidence:.2f})\nSarcasm: {sarcasm} (Prob: {sarcasm_prob:.2f})",
             label,
-            sarcasm_text
+            sarcasm
         )
     elif text and text.strip():
         label, confidence = classify_text(text)
-        sarcastic, sarcasm_prob = is_sarcastic(text)
-        sarcasm_text = "Yes" if sarcastic else "No"
+        sarcasm, sarcasm_prob = is_sarcastic(text)
         return (
-            f"Text: {text}\nPrediction: {label} (Confidence: {confidence:.2f})\nSarcasm: {sarcasm_text} (Prob: {sarcasm_prob:.2f})",
+            f"Text: {text}\nPrediction: {label} (Confidence: {confidence:.2f})\nSarcasm: {sarcasm} (Prob: {sarcasm_prob:.2f})",
             label,
-            sarcasm_text
+            sarcasm
         )
     else:
         return "Please provide an image or some text.", None, None
@@ -89,14 +84,13 @@ iface = gr.Interface(
     ],
     outputs=[
         gr.Textbox(label="Prediction & Sarcasm Detection"),
-        gr.Label(num_top_classes=len(LABELS), label="Hate Speech Class"),
+        gr.Label(num_top_classes=len(HATE_LABELS), label="Hate Speech Class"),
         gr.Label(num_top_classes=2, label="Sarcasm")
     ],
-    title="Cyberbully detection system Chatbot",
+    title="Cyberbully Detection System Chatbot",
     description="""
-    Classifies text (or text extracted from image) into hate speech categories including abusive words,
-    threat, harassment, and detects sarcasm separately. Enter text or upload an image screenshot.
-    """
+Classifies text (or text extracted from image) into hate speech categories (sexism, racism, disability, sexual_orientation, religion, other, not_hate) and detects sarcasm (separately). Enter text or upload an image screenshot.
+"""
 )
 
 if __name__ == "__main__":