Update app.py

app.py

@@ -4,25 +4,31 @@ import re
 import time
 import torch
 import gradio as gr
+import numpy as np
 from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList
+from transformers import pipeline
 from collections import defaultdict
 from datetime import datetime, timedelta
+import tempfile
 
 # =============================
 # Configuration
 # =============================
 MODEL_PATH = r"Muhammadidrees/JayConverstionalModel"
+WHISPER_MODEL = "openai/whisper-small"  # Change to "openai/whisper-base" for faster, or "openai/whisper-medium" for better accuracy
+TTS_MODEL = "suno/bark-small"  # Alternative: "facebook/mms-tts-eng" for faster TTS
+
 MAX_NEW_TOKENS = 200
 TEMPERATURE = 0.5
 TOP_K = 50
 REPETITION_PENALTY = 1.1
-MAX_HISTORY_TURNS = 5  # Limit conversation history
+MAX_HISTORY_TURNS = 5
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"🚀 Loading model from {MODEL_PATH} on {device}...")
+print(f"🚀 Loading models on {device}...")
 
 # =============================
-# Rate Limiting (Simple IP-based)
+# Rate Limiting
 # =============================
 rate_limit_store = defaultdict(list)
 MAX_REQUESTS_PER_MINUTE = 10
@@ -42,9 +48,11 @@ def check_rate_limit(session_id):
     return True
 
 # ==========================
-# Load Model & Tokenizer
+# Load Models
 # =============================
 try:
+    # Load ChatDoctor Model
+    print("Loading ChatDoctor model...")
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_PATH,
@@ -52,9 +60,33 @@ try:
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
         low_cpu_mem_usage=True
     )
-    print("✅ ChatDoctor model loaded successfully!\n")
+    print("✅ ChatDoctor model loaded!")
+    
+    # Load Whisper (Speech-to-Text)
+    print("Loading Whisper ASR model...")
+    whisper_pipe = pipeline(
+        "automatic-speech-recognition",
+        model=WHISPER_MODEL,
+        device=0 if torch.cuda.is_available() else -1
+    )
+    print("✅ Whisper model loaded!")
+    
+    # Load TTS Model
+    print("Loading TTS model...")
+    try:
+        tts_pipe = pipeline(
+            "text-to-speech",
+            model=TTS_MODEL,
+            device=0 if torch.cuda.is_available() else -1
+        )
+        print("✅ TTS model loaded!")
+        TTS_AVAILABLE = True
+    except Exception as e:
+        print(f"⚠️ TTS model not available: {e}")
+        TTS_AVAILABLE = False
+    
 except Exception as e:
-    print(f"❌ Error loading model: {e}")
+    print(f"❌ Error loading models: {e}")
     raise
 
 # =============================
@@ -75,7 +107,6 @@ class StopOnTokens(StoppingCriteria):
                         return True
         return False
 
-
 # =============================
 # Medical Keywords and Validation
 # =============================
@@ -89,7 +120,6 @@ MEDICAL_KEYWORDS = [
     "rash", "swelling", "injury", "bruise", "cold", "sneeze", "tired", "weak"
 ]
 
-# Emergency keywords that should trigger immediate medical attention warning
 EMERGENCY_KEYWORDS = [
     "suicide", "kill myself", "end my life", "chest pain", "can't breathe",
     "severe bleeding", "overdose", "poisoning", "unconscious", "seizure",
@@ -103,23 +133,23 @@ CASUAL_PATTERNS = [
     r"^what'?s\s+up\s*[\?\!\.]*$",
 ]
 
+DANGEROUS_PATTERNS = [
+    r"take\s+\d+\s+(pills|tablets|capsules)",
+    r"inject\s+(yourself|myself)",
+    r"(don't|do not)\s+go\s+to\s+(hospital|doctor|emergency)",
+    r"ignore\s+(doctor|medical|professional)",
+]
 
 def is_emergency_query(message):
-    """Detect if query contains emergency keywords"""
     message_lower = message.lower()
     return any(keyword in message_lower for keyword in EMERGENCY_KEYWORDS)
 
-
 def is_medical_query(message):
-    """Enhanced medical query detection"""
     message_lower = message.lower()
-    
-    # Check for medical keywords
     for keyword in MEDICAL_KEYWORDS:
         if keyword in message_lower:
             return True
     
-    # Check for question patterns with sufficient length
     question_words = ["what", "how", "why", "when", "where", "can", "should", "is", "are", "do", "does", "could", "would"]
     words = message_lower.split()
     has_question = any(q in words[:4] for q in question_words)
@@ -129,42 +159,75 @@ def is_medical_query(message):
     
     return False
 
-
 def is_only_greeting(message):
-    """Improved greeting detection using regex"""
     message_clean = message.lower().strip()
-    
-    # Remove punctuation for matching
     message_clean = re.sub(r'[!?.]+$', '', message_clean)
     
-    # Check if it matches any casual pattern
     for pattern in CASUAL_PATTERNS:
         if re.match(pattern, message_clean):
             return True
     
     return False
 
-
-# =============================
-# Safety Filter
-# =============================
-DANGEROUS_PATTERNS = [
-    r"take\s+\d+\s+(pills|tablets|capsules)",
-    r"inject\s+(yourself|myself)",
-    r"(don't|do not)\s+go\s+to\s+(hospital|doctor|emergency)",
-    r"ignore\s+(doctor|medical|professional)",
-]
-
 def contains_dangerous_advice(response):
-    """Check if response contains potentially dangerous medical advice"""
     response_lower = response.lower()
-    
     for pattern in DANGEROUS_PATTERNS:
         if re.search(pattern, response_lower):
             return True
-    
     return False
 
+# =============================
+# Speech Processing Functions
+# =============================
+def transcribe_audio(audio):
+    """Convert speech to text using Whisper"""
+    if audio is None:
+        return ""
+    
+    try:
+        # Handle different audio input formats
+        if isinstance(audio, tuple):
+            sample_rate, audio_data = audio
+        else:
+            audio_data = audio
+        
+        # Ensure audio is in the right format
+        if isinstance(audio_data, np.ndarray):
+            if audio_data.dtype != np.float32:
+                audio_data = audio_data.astype(np.float32) / np.iinfo(audio_data.dtype).max
+        
+        # Transcribe
+        result = whisper_pipe(audio_data)
+        transcription = result["text"].strip()
+        
+        return transcription
+        
+    except Exception as e:
+        print(f"Error in transcription: {e}")
+        return ""
+
+def text_to_speech(text):
+    """Convert text to speech"""
+    if not TTS_AVAILABLE or not text:
+        return None
+    
+    try:
+        # Limit text length for TTS (to avoid timeout)
+        if len(text) > 500:
+            text = text[:500] + "..."
+        
+        # Generate speech
+        speech = tts_pipe(text)
+        
+        # Extract audio data
+        audio_data = speech["audio"]
+        sampling_rate = speech["sampling_rate"]
+        
+        return (sampling_rate, audio_data)
+        
+    except Exception as e:
+        print(f"Error in TTS: {e}")
+        return None
 
 # =============================
 # Get Response
@@ -172,11 +235,9 @@ def contains_dangerous_advice(response):
 def get_response(user_input, history_context, session_id="default"):
     """Generate response with enhanced safety and quality checks"""
     
-    # Rate limiting check
     if not check_rate_limit(session_id):
         return "⏰ You've made too many requests. Please wait a minute before trying again."
     
-    # Emergency detection
     if is_emergency_query(user_input):
         return (
             "🚨 **EMERGENCY DETECTED** 🚨\n\n"
@@ -187,11 +248,9 @@ def get_response(user_input, history_context, session_id="default"):
             "This AI cannot provide emergency medical care. Please seek immediate professional help."
         )
     
-    # Greeting detection
     if is_only_greeting(user_input):
         return "👋 Hello! I'm ChatDoctor — your AI medical assistant. Please tell me about any health symptoms or medical concerns you'd like to discuss."
 
-    # Non-medical query handling
     if not is_medical_query(user_input):
         return (
             "Hello! I'm ChatDoctor, an AI medical assistant specialized in health and wellness.\n\n"
@@ -202,7 +261,6 @@ def get_response(user_input, history_context, session_id="default"):
             "Please describe your health concern in detail to get started."
         )
 
-    # Build prompt with limited history
     human_prefix = "Patient:"
     doctor_prefix = "ChatDoctor:"
     system_instruction = (
@@ -212,7 +270,6 @@ def get_response(user_input, history_context, session_id="default"):
         "Never provide dosage instructions or tell patients to avoid seeking professional help.\n\n"
     )
 
-    # Limit history to prevent token overflow
     limited_history = history_context[-MAX_HISTORY_TURNS:] if len(history_context) > MAX_HISTORY_TURNS else history_context
     
     history_text = [system_instruction]
@@ -228,7 +285,6 @@ def get_response(user_input, history_context, session_id="default"):
     try:
         input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
         
-        # Stop words for cleaner output
         stop_words = ["Patient:", "\nPatient:", "Patient :", "\n\nPatient"]
         stop_ids = [tokenizer.encode(word, add_special_tokens=False) for word in stop_words]
         stopping_criteria = StoppingCriteriaList([StopOnTokens(stop_ids)])
@@ -248,7 +304,6 @@ def get_response(user_input, history_context, session_id="default"):
 
         response = tokenizer.decode(output_ids[0], skip_special_tokens=True)[len(prompt):].strip()
 
-        # Clean up response
         for stop_word in ["Patient:", "Patient :", "\nPatient", "Patient"]:
             if stop_word in response:
                 response = response.split(stop_word)[0].strip()
@@ -256,26 +311,22 @@ def get_response(user_input, history_context, session_id="default"):
 
         response = response.strip()
         
-        # Safety filter
         if contains_dangerous_advice(response):
             response = (
                 "I apologize, but I cannot provide that specific medical advice. "
                 "Please consult with a qualified healthcare professional who can properly evaluate your situation."
             )
         
-        # Filter out inappropriate content
         if any(x in response.lower() for x in ["chatbot", "api key", "error", "cloud", "sorry, i don't have"]):
             response = (
                 "I apologize for the confusion. I'm ChatDoctor, trained to assist with medical and health-related topics. "
                 "Please tell me more about your symptoms or health concerns so I can help you better."
             )
         
-        # Add disclaimer for serious conditions
         serious_conditions = ["cancer", "tumor", "heart disease", "stroke", "diabetes complications"]
         if any(condition in response.lower() for condition in serious_conditions):
             response += "\n\n⚠️ **Important:** Please consult a healthcare professional for proper diagnosis and treatment."
 
-        # Clean up memory
         del input_ids, output_ids
         gc.collect()
         if torch.cuda.is_available():
@@ -287,7 +338,6 @@ def get_response(user_input, history_context, session_id="default"):
         print(f"Error generating response: {e}")
         return "I apologize, but I encountered an error processing your request. Please try rephrasing your question or try again later."
 
-
 # =============================
 # Gradio Interface
 # =============================
@@ -320,6 +370,12 @@ custom_css = """
     margin: 15px 0;
     color: #721c24;
 }
+.voice-section {
+    background: linear-gradient(135deg, #e0c3fc 0%, #8ec5fc 100%);
+    border-radius: 10px;
+    padding: 20px;
+    margin: 15px 0;
+}
 footer {
     margin-top: 30px;
     padding: 15px;
@@ -328,13 +384,14 @@ footer {
     font-size: 0.9em;
 }
 """
+
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
-    session_state = gr.State(value=str(time.time()))  # Unique session ID
+    session_state = gr.State(value=str(time.time()))
     
     gr.HTML("""
         <div id="header">
             <h1>🩺 ChatDoctor AI Assistant</h1>
-            <p>Your AI-powered medical consultation partner</p>
+            <p>🎤 Voice-Enabled Medical Consultation Partner</p>
         </div>
     """)
     
@@ -356,32 +413,73 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         </div>
     """)
 
-    chatbot = gr.Chatbot(
-        height=500,
-        placeholder="<div style='text-align:center;padding:50px;'><h3>👋 Welcome to ChatDoctor!</h3><p style='color:#6c757d;'>Describe your symptoms or ask a health-related question to begin.</p><p style='color:#dc3545;margin-top:15px;'><strong>Remember:</strong> This is not a replacement for professional medical care.</p></div>",
-        show_label=False,
-        avatar_images=(None, "🤖"),
-    )
-
-    with gr.Row():
-        msg = gr.Textbox(
-            placeholder="Type your medical concern here... (e.g., 'I have a headache for 3 days')",
+    with gr.Tab("💬 Text Chat"):
+        chatbot = gr.Chatbot(
+            height=500,
+            placeholder="<div style='text-align:center;padding:50px;'><h3>👋 Welcome to ChatDoctor!</h3><p style='color:#6c757d;'>Describe your symptoms or ask a health-related question to begin.</p></div>",
             show_label=False,
-            scale=9,
-            container=False,
-            lines=1
+            avatar_images=(None, "🤖"),
         )
-        send_btn = gr.Button("Send 📤", scale=1, variant="primary")
 
-    with gr.Row():
-        clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
-        retry_btn = gr.Button("🔄 Retry", scale=1)
+        with gr.Row():
+            msg = gr.Textbox(
+                placeholder="Type your medical concern here...",
+                show_label=False,
+                scale=9,
+                container=False,
+                lines=1
+            )
+            send_btn = gr.Button("Send 📤", scale=1, variant="primary")
+
+        with gr.Row():
+            clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
+            retry_btn = gr.Button("🔄 Retry", scale=1)
+
+    with gr.Tab("🎤 Voice Chat"):
+        gr.HTML('<div class="voice-section"><h3>🎙️ Voice Interaction</h3><p>Record your medical question and get voice responses!</p></div>')
+        
+        voice_chatbot = gr.Chatbot(
+            height=400,
+            placeholder="<div style='text-align:center;padding:40px;'><h3>🎤 Voice Chat Mode</h3><p>Click the microphone to record your question</p></div>",
+            show_label=False,
+            avatar_images=(None, "🤖"),
+        )
+        
+        with gr.Row():
+            audio_input = gr.Audio(
+                sources=["microphone"],
+                type="numpy",
+                label="🎤 Record Your Question",
+                scale=8
+            )
+            voice_send_btn = gr.Button("Send Voice 🎙️", scale=2, variant="primary")
+        
+        audio_output = gr.Audio(
+            label="🔊 Voice Response",
+            autoplay=True,
+            visible=TTS_AVAILABLE
+        )
+        
+        transcribed_text = gr.Textbox(
+            label="📝 Transcribed Text",
+            interactive=False,
+            visible=True
+        )
+        
+        with gr.Row():
+            voice_clear_btn = gr.Button("🗑️ Clear Voice Chat", scale=1)
+        
+        if not TTS_AVAILABLE:
+            gr.Warning("⚠️ TTS model not available. Voice responses disabled. Text responses will still work.")
 
     with gr.Accordion("⚙️ Advanced Settings", open=False):
         temp_slider = gr.Slider(0.1, 1.0, TEMPERATURE, 0.1, label="Temperature (Lower = More Focused)")
         max_tok_slider = gr.Slider(50, 500, MAX_NEW_TOKENS, 50, label="Max Tokens")
         top_k_slider = gr.Slider(1, 100, TOP_K, 1, label="Top-K Sampling")
 
+    # =============================
+    # Text Chat Functions
+    # =============================
     def user_message(user_msg, history):
         if not user_msg.strip():
             return "", history
@@ -407,6 +505,36 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         history[-1][1] = bot_msg
         return history
 
+    # =============================
+    # Voice Chat Functions
+    # =============================
+    def process_voice_input(audio, history, temp, max_tok, topk, session_id):
+        """Process voice input: transcribe, get response, convert to speech"""
+        if audio is None:
+            return history, "", None
+        
+        # Transcribe audio to text
+        transcribed = transcribe_audio(audio)
+        
+        if not transcribed:
+            return history, "⚠️ Could not transcribe audio. Please try again.", None
+        
+        # Add to history
+        history = history + [[transcribed, None]]
+        
+        # Get bot response
+        global TEMPERATURE, MAX_NEW_TOKENS, TOP_K
+        TEMPERATURE, MAX_NEW_TOKENS, TOP_K = temp, int(max_tok), int(topk)
+        
+        bot_msg = get_response(transcribed, history[:-1], session_id)
+        history[-1][1] = bot_msg
+        
+        # Convert response to speech
+        audio_response = text_to_speech(bot_msg) if TTS_AVAILABLE else None
+        
+        return history, transcribed, audio_response
+
+    # Text Chat Events
     msg.submit(user_message, [msg, chatbot], [msg, chatbot], queue=False).then(
         bot_response, [chatbot, temp_slider, max_tok_slider, top_k_slider, session_state], chatbot
     )
@@ -416,10 +544,19 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     clear_btn.click(lambda: None, None, chatbot, queue=False)
     retry_btn.click(retry_last, [chatbot, temp_slider, max_tok_slider, top_k_slider, session_state], chatbot)
 
+    # Voice Chat Events
+    voice_send_btn.click(
+        process_voice_input,
+        [audio_input, voice_chatbot, temp_slider, max_tok_slider, top_k_slider, session_state],
+        [voice_chatbot, transcribed_text, audio_output]
+    )
+    voice_clear_btn.click(lambda: (None, "", None), None, [voice_chatbot, transcribed_text, audio_output], queue=False)
+
     gr.HTML(f"""
         <footer>
-            <p><strong>🧠 Powered by LLaMA-based ChatDoctor</strong></p>
+            <p><strong>🧠 Powered by LLaMA + Whisper + TTS</strong></p>
             <p>Device: {device.upper()} | Rate Limit: {MAX_REQUESTS_PER_MINUTE} requests/minute</p>
+            <p>🎤 Voice: Whisper ASR | 🔊 TTS: {"Enabled" if TTS_AVAILABLE else "Disabled"}</p>
             <p style='font-size:0.85em;margin-top:10px;'>
                 This AI provides general health information only. Always consult healthcare professionals for medical advice.
             </p>
@@ -430,11 +567,11 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 # Launch App
 # =============================
 if __name__ == "__main__":
-    print("\n💡 Launching Enhanced ChatDoctor Gradio Interface...")
+    print("\n💡 Launching ChatDoctor with Voice Support...")
     print(f"📊 Configuration:")
-    print(f"   - Max History Turns: {MAX_HISTORY_TURNS}")
-    print(f"   - Rate Limit: {MAX_REQUESTS_PER_MINUTE} requests/minute")
     print(f"   - Device: {device.upper()}")
+    print(f"   - Whisper Model: {WHISPER_MODEL}")
+    print(f"   - TTS Available: {TTS_AVAILABLE}")
+    print(f"   - Rate Limit: {MAX_REQUESTS_PER_MINUTE} requests/minute")
     demo.queue()
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
-    
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)