Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import re
 
 print("Loading MindBridge model...")
 model_name = "prats010/mindbridge-mental-health-model"
@@ -23,6 +24,29 @@ Never diagnose. Use "it sounds like..." not "you have...".
 If user mentions suicide or self harm, add [CRISIS] at the end of your response.
 Always recommend professional help for serious symptoms."""
 
+def clean_response(text):
+    # Remove URLs
+    text = re.sub(r'http\S+|www\S+', '', text)
+    # Remove email addresses
+    text = re.sub(r'\S+@\S+', '', text)
+    # Remove social media handles
+    text = re.sub(r'facebook\S*|twitter\S*|instagram\S*', '', text, flags=re.IGNORECASE)
+    # Remove lines with names/personal info patterns
+    text = re.sub(r'My name is\s+\w+', '', text, flags=re.IGNORECASE)
+    # Remove excessive whitespace
+    text = re.sub(r'\s+', ' ', text).strip()
+    # Cut at first URL-like pattern or organization name
+    bad_patterns = ['MentalHealthCanada', 'myfitnesscareer', 'facebook.com', '.ca', '.com']
+    for pattern in bad_patterns:
+        if pattern in text:
+            text = text[:text.index(pattern)].strip()
+    # Keep only first 2 sentences
+    sentences = [s.strip() for s in text.split('.') if s.strip()]
+    sentences = [s for s in sentences if len(s) > 10]
+    if len(sentences) > 2:
+        text = '. '.join(sentences[:2]) + '.'
+    return text
+
 def chat(user_input):
     input_text = f"{SYSTEM_CONTEXT}\n\nUser: {user_input}\nCounselor:"
     
@@ -38,13 +62,13 @@ def chat(user_input):
         outputs = model.generate(
             inputs,
             attention_mask=attention_mask,
-            max_new_tokens=200,
-            temperature=0.8,
-            top_p=0.85,
-            top_k=50,
+            max_new_tokens=100,
+            temperature=0.7,
+            top_p=0.8,
+            top_k=40,
             do_sample=True,
-            repetition_penalty=1.5,
-            no_repeat_ngram_size=3,
+            repetition_penalty=2.0,
+            no_repeat_ngram_size=4,
             pad_token_id=tokenizer.eos_token_id
         )
     
@@ -56,9 +80,10 @@ def chat(user_input):
     if "You are MindBridge" in response:
         response = response.split("You are MindBridge")[0].strip()
     
-    sentences = response.split(". ")
-    if len(sentences) > 4:
-        response = ". ".join(sentences[:4]) + "."
+    response = clean_response(response)
+    
+    if not response or len(response) < 10:
+        response = "I hear you. It sounds like you're going through a difficult time. Would you like to share more about what you're feeling?"
     
     return response