Update app.py

app.py

@@ -5,14 +5,16 @@ import easyocr
 from PIL import Image
 import numpy as np
 
-# RoBERTa Multiclass Model
+# Set up model and labels
 MODEL_NAME = "cardiffnlp/twitter-roberta-base-hate-multiclass-latest"
 LABELS = [
-    "no hate",             # 0
-    "racist",              # 1
-    "religious hate",      # 2
-    "sexual abuse",        # 3
-    "sarcastic"            # 4
+    "sexism",             # 0
+    "racism",             # 1
+    "disability",         # 2
+    "sexual_orientation", # 3
+    "religion",           # 4
+    "other",              # 5
+    "not_hate"            # 6
 ]
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -23,28 +25,32 @@ def classify_text(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
     with torch.no_grad():
         outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        logits = outputs.logits
+        probs = torch.nn.functional.softmax(logits, dim=-1)
         pred = torch.argmax(probs).item()
-        return LABELS[pred], float(probs[0][pred])
+        confidence = float(probs[0][pred])
+        return LABELS[pred], confidence
 
 def ocr_extract(image):
+    # Convert to numpy if Image
     if isinstance(image, Image.Image):
         image = np.array(image)
     result = reader.readtext(image, detail=0)
     return ' '.join(result)
 
 def chatbot(image=None, text=None):
+    # Prioritize image
     if image is not None:
         extracted = ocr_extract(image)
         if not extracted.strip():
             return "No text found in image.", None
-        hate_class, prob = classify_text(extracted)
-        return f"OCR: {extracted}\n\nClass: {hate_class} (Prob: {prob:.2f})", hate_class
+        label, confidence = classify_text(extracted)
+        return f"OCR Extracted: {extracted}\nPrediction: {label} (Confidence: {confidence:.2f})", label
     elif text and text.strip():
-        hate_class, prob = classify_text(text)
-        return f"Text: {text}\nClass: {hate_class} (Prob: {prob:.2f})", hate_class
+        label, confidence = classify_text(text)
+        return f"Text: {text}\nPrediction: {label} (Confidence: {confidence:.2f})", label
     else:
-        return "Please provide a screenshot or text input.", None
+        return "Please provide an image or some text.", None
 
 iface = gr.Interface(
     fn=chatbot,
@@ -54,10 +60,10 @@ iface = gr.Interface(
     ],
     outputs=[
         gr.Textbox(label="Prediction & OCR"),
-        gr.Label(num_top_classes=5)
+        gr.Label(num_top_classes=7)
     ],
-    title="Multiclass Hate Speech Detector Chatbot (RoBERTa, with OCR)",
-    description="Detects religious hate, sexual abuse, racism, sarcasm or no hate. Upload a screenshot or enter text."
+    title="RoBERTa Multiclass Hate Speech Classifier (with OCR)",
+    description="Detects: sexism, racism, disability, sexual_orientation, religion, other, not_hate. Enter text or upload screenshot."
 )
 
 if __name__ == "__main__":