Update app.py

app.py

@@ -4,48 +4,45 @@ import random
 import re
 import torch
 from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
-from transformers import Wav2Vec2FeatureExtractor
 import librosa
 from gtts import gTTS
-import numpy as np
 import tempfile
 import os
 
-# ------------------------------
-# Allowed mental health keywords
+# ===== Mental health keywords (EN + AR + transliterated AR)
 MENTAL_KEYWORDS = [
     "depression", "depressed", "anxiety", "anxious", "panic", "stress", "sad", "lonely",
-    "trauma", "mental", "therapy", "therapist", "counselor", "mood", "overwhelmed", "anger",
-    "fear", "worry", "self-esteem", "confidence", "motivation", "relationship", "cope", "coping",
-    "relax", "calm", "sleep", "emotion", "feeling", "feel", "thoughts", "help", "life", "advice",
-    "unmotivated", "lost", "hopeless", "tired", "burnout", "cry", "hurt", "love", "breakup",
-    "friend", "family", "alone", "heartbroken", "scared", "fearful",
-    "ana", "zahqan", "daye2", "ha2t", "mota3ab", "mota3eb", "za3lan", "malo", "khalni", "mash3or",
-    "bakhaf", "w7ed", "msh 3aref", "mash fahem", "malish", "3ayez", "ayez", "7azeen", "mdaye2",
+    "trauma", "mental", "therapy", "therapist", "counselor", "mood", "overwhelmed",
+    "anger", "fear", "worry", "self-esteem", "confidence", "motivation", "relationship",
+    "cope", "coping", "relax", "calm", "sleep", "emotion", "feeling", "feel", "thoughts",
+    "help", "life", "advice", "unmotivated", "lost", "hopeless", "tired", "burnout",
+    "cry", "hurt", "love", "breakup", "friend", "family", "alone", "heartbroken",
+    "scared", "fearful",
+    # Transliterated Arabic
+    "ana", "zahqan", "daye2", "ha2t", "mota3ab", "mota3eb", "za3lan", "malo", "khalni",
+    "mash3or", "bakhaf", "w7ed", "msh 3aref", "mash fahem", "malish", "3ayez", "ayez",
+    "7azeen", "mdaye2",
+    # Arabic
     "حزين", "تعبان", "قلق", "خايف", "وحدة", "ضيق", "توتر", "زعلان", "اكتئاب", "علاج",
     "مشاعر", "مضغوط", "قلقان", "وحدي", "مش مبسوط", "زهقان", "ضايق", "تعب", "مش مرتاح",
 ]
 
-# Off-topic keywords
 OFF_TOPIC = [
-    "recipe", "song", "music", "lyrics", "joke", "funny", "laugh", "code", "python", "program",
-    "game", "food", "cook", "movie", "film", "series", "sport", "football", "instagram",
-    "tiktok", "money", "business", "crypto", "ai", "computer",
+    "recipe", "song", "music", "lyrics", "joke", "funny", "laugh", "code", "python",
+    "program", "game", "food", "cook", "movie", "film", "series", "sport", "football",
+    "instagram", "tiktok", "money", "business", "crypto", "ai", "computer",
+    # Arabic
     "نكتة", "ضحك", "اغنية", "اغاني", "طبخ", "اكل", "فيلم", "مسلسل", "كورة", "رياضة",
     "بيزنس", "فلوس", "العاب", "لعبة", "كود", "برمجة", "ذكاء اصطناعي"
 ]
 
-# Random off-topic responses
 OFF_TOPIC_RESPONSES = [
-    "I'm here to help with emotional and mental well-being. Let's focus on how you're feeling today.",
-    "I specialize in mental and emotional health conversations. Tell me what’s on your mind.",
-    "Let’s bring it back to your feelings — I’m here to help you talk through emotions or stress.",
-    "My goal is to support your mental health. How have things been emotionally lately?",
-    "I’m here for emotional support only. What’s been bothering you recently?",
+    "I'm here to help with emotional and mental well-being. Let's focus on your feelings.",
+    "I specialize in mental and emotional health conversations. How have you been feeling?",
+    "Let's bring it back to your emotions — I'm here to help process stress or challenges.",
 ]
 
-# ------------------------------
-# Arabic detection
+# Detect Arabic text
 def contains_arabic(text: str) -> bool:
     return bool(re.search(r"[\u0600-\u06FF]", text))
 
@@ -59,11 +56,10 @@ def is_mental_health_related(text: str) -> bool:
         return True
     return False
 
-# ------------------------------
-# Load voice emotion model
+# ===== Voice emotion detection
 voice_model_name = "Hatman/audio-emotion-detection"
 voice_model = Wav2Vec2ForSequenceClassification.from_pretrained(voice_model_name)
-voice_processor = Wav2Vec2FeatureExtractor.from_pretrained(voice_model_name)
+voice_processor = Wav2Vec2Processor.from_pretrained(voice_model_name)
 
 def detect_voice_emotion(audio_file):
     audio, sr = librosa.load(audio_file, sr=16000)
@@ -73,53 +69,55 @@ def detect_voice_emotion(audio_file):
     predicted_id = torch.argmax(logits, dim=-1).item()
     return voice_model.config.id2label[predicted_id]
 
-# ------------------------------
-# Chat + voice response with TTS
-def respond(message, history: list[dict[str, str]], system_message, max_tokens, temperature, top_p, hf_token: gr.OAuthToken, audio=None):
-    transcript = {"user": "", "bot": "", "tts": None}
+# ===== Chat function with mood, TTS, transcript
+def respond(message, history, system_message, max_tokens, temperature, top_p, hf_token: gr.OAuthToken, audio=None):
+    transcript = []
     response_text = ""
-    
-    # Detect audio emotion
+
     if audio:
-        emotion = detect_voice_emotion(audio)
-        response_text += f"[Detected mood: {emotion}] "
-    
-    transcript["user"] = message
+        mood = detect_voice_emotion(audio)
+        response_text += f"[Detected mood: {mood}] "
+
     if not is_mental_health_related(message):
         response_text += random.choice(OFF_TOPIC_RESPONSES)
-        transcript["bot"] = response_text
-    else:
-        locked_system_message = (
-            "You are a licensed mental health therapy assistant. "
-            "You respond with empathy, emotional intelligence, and a therapeutic tone. "
-            "Never answer questions unrelated to emotional or mental wellness, even if they are in another language."
-        )
-        
-        client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
-        messages = [{"role": "system", "content": locked_system_message}]
-        messages.extend(history)
-        messages.append({"role": "user", "content": message})
-        
-        # Stream response
-        for msg in client.chat_completion(messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p):
-            choices = msg.choices
-            token = ""
-            if len(choices) and choices[0].delta.content:
-                token = choices[0].delta.content
-            response_text += token
-        transcript["bot"] = response_text
-    
-    # Convert response to speech (TTS)
-    tts = gTTS(text=transcript["bot"], lang='ar' if contains_arabic(transcript["bot"]) else 'en')
-    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
-    tts.save(tmp_file.name)
-    transcript["tts"] = tmp_file.name
-    
-    return transcript
-
-# ------------------------------
-# Gradio UI
+        transcript.append(("User", message))
+        transcript.append(("Bot", response_text))
+        tts_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
+        tts = gTTS(response_text)
+        tts.save(tts_file)
+        return response_text, tts_file, transcript
+
+    locked_system_message = (
+        "You are a licensed mental health therapy assistant. "
+        "You respond with empathy, emotional intelligence, and a therapeutic tone. "
+        "Never answer questions unrelated to emotional or mental wellness."
+    )
+
+    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+    messages = [{"role": "system", "content": locked_system_message}]
+    messages.extend(history)
+    messages.append({"role": "user", "content": message})
+
+    for msg in client.chat_completion(messages, max_tokens=max_tokens, stream=True,
+                                      temperature=temperature, top_p=top_p):
+        choices = msg.choices
+        token = ""
+        if len(choices) and choices[0].delta.content:
+            token = choices[0].delta.content
+        response_text += token
+
+    transcript.append(("User", message))
+    transcript.append(("Bot", response_text))
+
+    tts_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
+    tts = gTTS(response_text)
+    tts.save(tts_file)
+
+    return response_text, tts_file, transcript
+
+# ===== Gradio UI
 with gr.Blocks() as demo:
+    gr.Markdown("## Mental Health Chatbot with Voice Mood Detection")
     with gr.Row():
         with gr.Column():
             chatbot = gr.ChatInterface(
@@ -130,8 +128,12 @@ with gr.Blocks() as demo:
                     gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
                     gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
                     gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-                    gr.Audio(source="microphone", type="filepath", label="Record your voice (optional)"),
+                    gr.Audio(label="Record your voice (optional)", type="filepath"),
                     gr.OAuthToken(label="Hugging Face Token"),
                 ],
             )
-    demo.launch()
+
+    # Output area for transcript
+    transcript_box = gr.Textbox(label="Transcript (User & Bot)", interactive=False)
+
+demo.launch()