Update app.py

app.py

@@ -5,102 +5,97 @@ import easyocr
 from PIL import Image
 import numpy as np
 
-# -------------------- MODEL SETUP --------------------
-MODEL_NAME = "microsoft/deberta-v3-base"  # Context-rich NLP model
+# -------------------------------
+# MODEL: CardiffNLP RoBERTa Hate Classifier
+# -------------------------------
+MODEL_NAME = "cardiffnlp/twitter-roberta-base-hate-multiclass-latest"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
-# Classes (Your categories)
-LABELS = [
-    "abusive_language",
-    "harassment",
-    "threat",
-    "racism",
-    "sexism",
-    "religious_hate",
-    "not_hate"
-]
-
-# Initialize OCR reader
-reader = easyocr.Reader(['en'], gpu=False)
-
-# -------------------- FUNCTIONS --------------------
-def extract_text_from_image(image):
-    """Extracts text from uploaded image using EasyOCR."""
+LABELS = ["sexism", "racism", "disability", "sexual_orientation", "religion", "other", "not_hate"]
+
+# -------------------------------
+# OCR Setup
+# -------------------------------
+reader = easyocr.Reader(["en"], gpu=False)
+
+def extract_text(image):
+    """Extract text from uploaded image using EasyOCR"""
     if image is None:
         return ""
     if isinstance(image, Image.Image):
         image = np.array(image)
-    extracted = reader.readtext(image, detail=0)
-    return " ".join(extracted)
-
-
-def classify_text_with_deberta(text):
-    """Runs text through DeBERTa model for classification."""
+    result = reader.readtext(image, detail=0)
+    return " ".join(result)
+
+# -------------------------------
+# CLASSIFICATION LOGIC
+# -------------------------------
+def classify_text(text):
+    """Classify input text using RoBERTa model"""
     if not text.strip():
         return "No text found for analysis.", None
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True)
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+        logits = model(**inputs).logits
         probs = torch.nn.functional.softmax(logits, dim=-1)
         pred = torch.argmax(probs).item()
         confidence = float(probs[0][pred])
-    return f"Prediction: {LABELS[pred]} (Confidence: {confidence:.2f})", LABELS[pred]
-
-
-def respond(chat_history, user_text, user_image):
-    """Chatbot pipeline: OCR → DeBERTa classification → Chat output"""
-    # OCR extraction if image uploaded
-    if user_image is not None:
-        extracted_text = extract_text_from_image(user_image)
-        if extracted_text.strip():
-            text = extracted_text
-            display_text = f"[Extracted from OCR] {extracted_text}"
-        elif user_text and user_text.strip():
-            text = user_text
-            display_text = text
-        else:
-            chat_history.append(("User", ""))
-            chat_history.append(("Bot", "Please enter text or upload a readable image."))
-            return chat_history, "", None
+        label = LABELS[pred]
+    return f"Category: {label} (Confidence: {confidence:.2f})", label
+
+
+# -------------------------------
+# CHATBOT FUNCTION
+# -------------------------------
+def cyberbully_chat(messages, user_message, image=None):
+    """Chat-like conversational function"""
+    history = messages or []
+
+    # Extract text from image (if any)
+    if image:
+        text = extract_text(image)
+        content = f"[Extracted from image] {text}" if text else "[No readable text found]"
     else:
-        text = user_text or ""
-        display_text = text.strip()
+        text = user_message
+        content = text.strip()
 
-    if not display_text:
-        chat_history.append(("User", ""))
-        chat_history.append(("Bot", "Empty input provided."))
-        return chat_history, "", None
+    if not content or content == "[No readable text found]":
+        history.append({"role": "assistant", "content": "Please provide valid text or an image with text."})
+        return history
 
-    # Run DeBERTa classification
-    result, label = classify_text_with_deberta(text)
-    chat_history.append(("User", display_text))
-    chat_history.append(("Cyber Bully Bot", result))
-    return chat_history, "", None
+    # Classify with model
+    classification, label = classify_text(text)
+    
+    # Append to chat
+    history.append({"role": "user", "content": content})
+    history.append({"role": "assistant", "content": classification})
+    return history
 
 
-# -------------------- GRADIO INTERFACE --------------------
+# -------------------------------
+# GRADIO ChatGPT-like UI
+# -------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 💬 Cyber Bully Detection System (OCR + DeBERTa Context Analysis)")
-    chat_history = gr.State([])
+    gr.Markdown("# 🤖 Cyber Bully Detection System")
+    gr.Markdown("Upload an image or type text. The system will analyze hate-speech categories using a RoBERTa model specialized for social media context.")
 
-    chatbot = gr.Chatbot(label="Chat History")
+    chatbot = gr.Chatbot(type="messages", label="CyberBully Chat")
     with gr.Row():
-        text_input = gr.Textbox(show_label=False, placeholder="Type a message or paste text here")
+        text_input = gr.Textbox(show_label=False, placeholder="Type a message here...")
         image_input = gr.Image(source="upload", type="pil", label="Upload Screenshot (optional)")
     with gr.Row():
         submit_btn = gr.Button("Analyze")
         clear_btn = gr.Button("Clear Chat")
 
-    submit_btn.click(respond, [chatbot, text_input, image_input], [chatbot, text_input, image_input])
-    clear_btn.click(lambda: ([], "", None), None, [chatbot, text_input, image_input])
-
-    gr.Markdown(
-        "🔍 **How it works:** Upload a screenshot or type text. "
-        "The system extracts text via OCR and uses DeBERTa to understand contextual meaning and classify it."
+    submit_btn.click(
+        cyberbully_chat,
+        [chatbot, text_input, image_input],
+        [chatbot],
+        queue=True
     )
 
-# -------------------- LAUNCH --------------------
+    clear_btn.click(lambda: [], None, chatbot, queue=False)
+
 if __name__ == "__main__":
     demo.launch()