Update app.py

app.py

@@ -5,14 +5,13 @@ import easyocr
 from PIL import Image
 import numpy as np
 
-# Sarcasm Model
-SARCASM_MODEL_NAME = "j-hartmann/emotion-english-distilroberta-base"
+# Models and labels (same as before)
+SARCASM_MODEL = "j-hartmann/emotion-english-distilroberta-base"
 sarcasm_labels = ["not sarcastic", "sarcastic"]
-sarcasm_tokenizer = AutoTokenizer.from_pretrained(SARCASM_MODEL_NAME)
-sarcasm_model = AutoModelForSequenceClassification.from_pretrained(SARCASM_MODEL_NAME)
+sarcasm_tokenizer = AutoTokenizer.from_pretrained(SARCASM_MODEL)
+sarcasm_model = AutoModelForSequenceClassification.from_pretrained(SARCASM_MODEL)
 
-# Hate Speech Model (RoBERTa multiclass)
-HATE_MODEL_NAME = "cardiffnlp/twitter-roberta-base-hate-multiclass-latest"
+HATE_MODEL = "cardiffnlp/twitter-roberta-base-hate-multiclass-latest"
 hate_labels = [
     "sexism",
     "racism",
@@ -22,17 +21,18 @@ hate_labels = [
     "other",
     "not_hate"
 ]
-hate_tokenizer = AutoTokenizer.from_pretrained(HATE_MODEL_NAME)
-hate_model = AutoModelForSequenceClassification.from_pretrained(HATE_MODEL_NAME)
+hate_tokenizer = AutoTokenizer.from_pretrained(HATE_MODEL)
+hate_model = AutoModelForSequenceClassification.from_pretrained(HATE_MODEL)
 
-# OCR Reader
 reader = easyocr.Reader(['en'], gpu=False)
 
 def extract_text(image):
+    if image is None:
+        return ""
     if isinstance(image, Image.Image):
         image = np.array(image)
-    result = reader.readtext(image, detail=0)
-    return ' '.join(result)
+    texts = reader.readtext(image, detail=0)
+    return ' '.join(texts)
 
 def detect_sarcasm(text):
     inputs = sarcasm_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -40,8 +40,8 @@ def detect_sarcasm(text):
         outputs = sarcasm_model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
         pred = torch.argmax(probs).item()
-        confidence = float(probs[0][pred])
-        return sarcasm_labels[pred], confidence
+        conf = float(probs[0][pred])
+        return sarcasm_labels[pred], conf
 
 def classify_hate(text):
     inputs = hate_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -49,56 +49,54 @@ def classify_hate(text):
         outputs = hate_model(**inputs)
         probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
         pred = torch.argmax(probs).item()
-        confidence = float(probs[0][pred])
-        return hate_labels[pred], confidence
+        conf = float(probs[0][pred])
+        return hate_labels[pred], conf
 
-
-def chatbot(conversation, user_message, image=None):
-    # Process image first if available
-    if image is not None:
-        extracted_text = extract_text(image)
-        if not extracted_text.strip():
-            response = "No text found in the uploaded image."
-            conversation.append(("User", user_message))
-            conversation.append(("Cyber Bully Bot", response))
-            return conversation, None, None
-        text = extracted_text
-        display_input = f"[Extracted from image] {text}"
+def respond(chat_history, user_text, user_image):
+    # Combine OCR and text input
+    if user_image is not None:
+        extracted = extract_text(user_image)
+        if extracted.strip():
+            input_text = extracted
+        elif user_text.strip():
+            input_text = user_text.strip()
+        else:
+            chat_history.append(("User", ""))
+            chat_history.append(("Bot", "Please provide text or an image with text."))
+            return chat_history, None, None
     else:
-        text = user_message.strip()
-        display_input = text
+        input_text = user_text.strip()
 
-    # Sarcasm detection
-    sarcasm_label, sarcasm_conf = detect_sarcasm(text)
+    sarcasm_label, sarcasm_conf = detect_sarcasm(input_text)
     if sarcasm_label == "sarcastic":
-        response = f"Text detected as SARCASTIC (Confidence: {sarcasm_conf:.2f}). Hate speech classification skipped."
-        conversation.append(("User", display_input))
-        conversation.append(("Cyber Bully Bot", response))
-        return conversation, None, None
-
-    # Hate speech classification
-    hate_label, hate_conf = classify_hate(text)
-    response = (
-        f"Hate Speech Category: {hate_label} (Confidence: {hate_conf:.2f})\n"
-        f"Sarcasm: {sarcasm_label} (Confidence: {sarcasm_conf:.2f})"
-    )
-    conversation.append(("User", display_input))
-    conversation.append(("Cyber Bully Bot", response))
-    return conversation, None, None
-
-default_conversation = []
+        response_text = f"Sarcasm detected (Confidence: {sarcasm_conf:.2f}). Hate speech detection skipped."
+        hate_label = None
+    else:
+        hate_label, hate_conf = classify_hate(input_text)
+        response_text = (
+            f"Hate Speech Category: {hate_label} (Confidence: {hate_conf:.2f})\n"
+            f"Text analyzed: \"{input_text}\""
+        )
+    chat_history.append(("User", input_text))
+    chat_history.append(("Bot", response_text))
+    return chat_history, None, None
 
-iface = gr.ChatInterface(
-    fn=chatbot,
-    title="Cyber Bully Detection System",
-    description="Upload images or enter text. Bot detects sarcasm first, then classifies hate speech categories.",
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# Cyber Bully Detection System (Chat Interface)")
+    
+    chat_history = gr.State([])
+    
+    with gr.Row():
+        chatbot = gr.Chatbot()
+    with gr.Row():
+        txt = gr.Textbox(show_label=False, placeholder="Type your message here and press enter")
+        img = gr.Image(label="Upload screenshot (optional)", type="pil")
+    with gr.Row():
+        clear = gr.Button("Clear Chat")
 
-# Add an image upload component beside text input
-iface.add_component(
-    gr.Image(source="upload", label="Upload Screenshot (optional)", interactive=True),
-    insert_before=iface.input_components[0]
-)
+    txt.submit(respond, [chatbot, txt, img], [chatbot, txt, img])
+    img.submit(respond, [chatbot, txt, img], [chatbot, txt, img])
+    clear.click(lambda: ([], None, None), None, [chatbot, txt, img])
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()