Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import spaces
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from langchain.memory import ConversationBufferMemory
 import re
 
 # Model configuration
@@ -26,229 +25,281 @@ SYSTEM_PROMPT = """You are a professional virtual doctor conducting a medical co
 - Build each question logically from their previous responses
 
 After 4-5 meaningful exchanges, provide assessment and recommendations.
-
 Do NOT make specific prescriptions for prescription-only drugs.
-
 Always maintain a professional, caring tone throughout the consultation."""
 
-MEDITRON_PROMPT = """<|im_start|>system
-You are a board-certified physician with extensive clinical experience. Your role is to provide evidence-based medical assessment and recommendations following standard medical practice.
+MEDITRON_PROMPT = """You are a board-certified physician providing evidence-based medical assessment. 
 
-For each patient case:
-1. Analyze presented symptoms systematically using medical terminology
-2. Create a structured differential diagnosis with most likely conditions first
+Based on the patient information provided, please:
+1. Analyze the symptoms systematically
+2. Provide a differential diagnosis with most likely conditions
 3. Recommend appropriate next steps (testing, monitoring, or treatment)
-4. Provide specific medication recommendations with precise dosing regimens
-5. Include clear red flags that would necessitate urgent medical attention
-6. Base all recommendations on current clinical guidelines and evidence-based medicine
-7. Maintain professional, clear, and compassionate communication
-
-Follow standard clinical documentation format when appropriate and prioritize patient safety at all times. Remember to include appropriate medical disclaimers.
-<|im_start|>user
-Patient information: {patient_info}
-<|im_end|>
-<|im_start|>assistant
-"""
+4. Suggest appropriate medications or remedies with dosing if applicable
+5. Include red flags that would require urgent medical attention
+6. Base recommendations on clinical guidelines
 
-print("Loading Llama-2 model...")
-tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
-model = AutoModelForCausalLM.from_pretrained(
-    LLAMA_MODEL,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
-print("Llama-2 model loaded successfully!")
-
-print("Loading Meditron model...")
-meditron_tokenizer = AutoTokenizer.from_pretrained(MEDITRON_MODEL)
-meditron_model = AutoModelForCausalLM.from_pretrained(
-    MEDITRON_MODEL,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
-print("Meditron model loaded successfully!")
-
-# Initialize LangChain memory for conversation tracking
-memory = ConversationBufferMemory(return_messages=True)
-
-# Simple state for basic info tracking
-conversation_state = {
-    'name': None,
-    'age': None,
-    'medical_turns': 0,
-    'has_name': False,
-    'has_age': False
-}
-
-def get_meditron_suggestions(patient_info):
-    """Use Meditron model to generate medicine and remedy suggestions."""
-    prompt = MEDITRON_PROMPT.format(patient_info=patient_info)
-    inputs = meditron_tokenizer(prompt, return_tensors="pt").to(meditron_model.device)
-    
-    with torch.no_grad():
-        outputs = meditron_model.generate(
-            inputs.input_ids,
-            attention_mask=inputs.attention_mask,
-            max_new_tokens=256,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True
-        )
-    
-    suggestion = meditron_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    return suggestion
-
-def build_prompt_with_memory(system_prompt, current_input):
-    """Build prompt using LangChain memory for full conversation context"""
-    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
-    
-    # Get conversation history from memory
-    messages = memory.chat_memory.messages
-    
-    # Add conversation history to prompt
-    for msg in messages:
-        if msg.type == "human":
-            prompt += f"{msg.content} [/INST] "
-        elif msg.type == "ai":
-            prompt += f"{msg.content} </s><s>[INST] "
-    
-    # Add current input
-    prompt += f"{current_input} [/INST] "
-    
-    return prompt
+Patient Information: {patient_info}
 
-@spaces.GPU
-def generate_response(message, history):
-    """Generate a response using LangChain ConversationBufferMemory."""
-    global conversation_state
-    
-    # Reset state if this is a new conversation
-    if not history:
-        conversation_state = {
-            'name': None,
-            'age': None,
-            'medical_turns': 0,
-            'has_name': False,
-            'has_age': False
-        }
-        # Clear memory for new conversation
-        memory.clear()
-    
-    # Save current user message to memory (we'll save bot response later)
-    memory.save_context({"input": message}, {"output": ""})
-    
-    # Step 1: Ask for name if not provided
-    if not conversation_state['has_name']:
-        conversation_state['has_name'] = True
-        bot_response = "Hello! Before we discuss your health concerns, could you please tell me your name?"
-        # Update memory with bot response
-        memory.save_context({"input": message}, {"output": bot_response})
-        return bot_response
-    
-    # Step 2: Store name and ask for age
-    if conversation_state['name'] is None:
-        conversation_state['name'] = message.strip()
-        bot_response = f"Nice to meet you, {conversation_state['name']}! Could you please tell me your age?"
-        # Update memory with bot response
-        memory.save_context({"input": message}, {"output": bot_response})
-        return bot_response
-    
-    # Step 3: Store age and start medical questions
-    if not conversation_state['has_age']:
-        conversation_state['age'] = message.strip()
-        conversation_state['has_age'] = True
-        bot_response = f"Thank you, {conversation_state['name']}! Now, what brings you here today? Please tell me about any symptoms or health concerns you're experiencing."
-        # Update memory with bot response
-        memory.save_context({"input": message}, {"output": bot_response})
-        return bot_response
+Please provide a structured medical assessment:"""
+
+# Load models
+print("Loading models...")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
     
-    # Step 4: Medical consultation phase using ConversationBufferMemory
-    conversation_state['medical_turns'] += 1
+    model = AutoModelForCausalLM.from_pretrained(
+        LLAMA_MODEL,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
+    print("Llama-2 model loaded successfully!")
     
-    # Build the prompt using memory for full conversation context
-    if conversation_state['medical_turns'] <= 5:
-        # Still gathering information - let LLM ask intelligent follow-up questions
-        prompt = build_prompt_with_memory(SYSTEM_PROMPT, message)
+    meditron_tokenizer = AutoTokenizer.from_pretrained(MEDITRON_MODEL)
+    if meditron_tokenizer.pad_token is None:
+        meditron_tokenizer.pad_token = meditron_tokenizer.eos_token
         
-        # Generate response with intelligent follow-up questions
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        with torch.no_grad():
-            outputs = model.generate(
-                inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_new_tokens=384,
-                temperature=0.8,
-                top_p=0.95,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
-            )
+    meditron_model = AutoModelForCausalLM.from_pretrained(
+        MEDITRON_MODEL,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
+    print("Meditron model loaded successfully!")
+except Exception as e:
+    print(f"Error loading models: {e}")
+
+class MedicalConsultationBot:
+    def __init__(self):
+        self.reset_conversation()
+    
+    def reset_conversation(self):
+        """Reset all conversation state"""
+        self.conversation_history = []
+        self.patient_name = None
+        self.patient_age = None
+        self.medical_turns = 0
+        self.stage = "greeting"  # greeting -> name -> age -> symptoms -> diagnosis
         
-        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        llama_response = full_response.split('[/INST]')[-1].strip()
+    def add_to_history(self, user_message, bot_response):
+        """Add exchange to conversation history"""
+        self.conversation_history.append({
+            "user": user_message,
+            "bot": bot_response
+        })
+    
+    def get_conversation_context(self):
+        """Get full conversation context as string"""
+        context = ""
+        if self.patient_name:
+            context += f"Patient Name: {self.patient_name}\n"
+        if self.patient_age:
+            context += f"Patient Age: {self.patient_age}\n"
         
-        # Save bot response to memory
-        memory.save_context({"input": message}, {"output": llama_response})
+        context += "\nConversation History:\n"
+        for exchange in self.conversation_history:
+            context += f"Patient: {exchange['user']}\n"
+            context += f"Doctor: {exchange['bot']}\n"
         
-        return llama_response
+        return context
     
-    else:
-        # Time for diagnosis and treatment (after 5+ turns)
-        # Get all conversation messages from memory
-        all_messages = memory.chat_memory.messages
-        
-        # Compile patient information from memory
-        patient_info = f"Patient: {conversation_state['name']}, Age: {conversation_state['age']}\n\n"
-        patient_info += "Complete Conversation History:\n"
+    def build_llama_prompt(self, current_message):
+        """Build prompt for Llama model with conversation context"""
+        prompt = f"<s>[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\n"
         
-        # Add all messages from memory
-        for msg in all_messages:
-            if msg.type == "human":
-                patient_info += f"Patient: {msg.content}\n"
-            elif msg.type == "ai":
-                patient_info += f"Doctor: {msg.content}\n"
+        # Add conversation context
+        context = self.get_conversation_context()
+        if context.strip():
+            prompt += f"Previous conversation context:\n{context}\n\n"
         
-        patient_info += f"Current: {message}\n"
+        prompt += f"Current patient message: {current_message}\n\nProvide a professional medical response with appropriate follow-up questions. [/INST]"
         
-        # Generate diagnosis with full conversation context
-        diagnosis_prompt = f"<s>[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\nBased on the complete conversation history, please provide a comprehensive medical assessment including likely diagnosis and recommendations for {conversation_state['name']}.\n\nComplete Patient Information:\n{patient_info} [/INST] "
+        return prompt
+
+# Global bot instance
+medical_bot = MedicalConsultationBot()
+
+def get_meditron_diagnosis(patient_info):
+    """Use Meditron model to generate medical assessment"""
+    try:
+        prompt = MEDITRON_PROMPT.format(patient_info=patient_info)
+        inputs = meditron_tokenizer(
+            prompt, 
+            return_tensors="pt", 
+            max_length=512, 
+            truncation=True
+        ).to(meditron_model.device)
         
-        inputs = tokenizer(diagnosis_prompt, return_tensors="pt").to(model.device)
         with torch.no_grad():
-            outputs = model.generate(
+            outputs = meditron_model.generate(
                 inputs.input_ids,
                 attention_mask=inputs.attention_mask,
-                max_new_tokens=384,
+                max_new_tokens=300,
                 temperature=0.7,
                 top_p=0.9,
                 do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
+                pad_token_id=meditron_tokenizer.pad_token_id
             )
         
-        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        diagnosis = full_response.split('[/INST]')[-1].strip()
+        response = meditron_tokenizer.decode(
+            outputs[0][inputs.input_ids.shape[1]:], 
+            skip_special_tokens=True
+        ).strip()
         
-        # Get treatment suggestions from Meditron using memory context
-        treatment_suggestions = get_meditron_suggestions(patient_info)
-        
-        # Combine responses
-        final_response = f"{diagnosis}\n\n--- TREATMENT RECOMMENDATIONS ---\n\n{treatment_suggestions}\n\n**Important:** These are general recommendations. Please consult with a healthcare professional for personalized medical advice."
+        return response
+    except Exception as e:
+        return f"Error generating medical assessment: {str(e)}"
+
+@spaces.GPU
+def medical_chat_response(message, history):
+    """Main chat response function with proper state management"""
+    global medical_bot
+    
+    # If this is a new conversation (empty history), reset the bot
+    if not history:
+        medical_bot.reset_conversation()
+    
+    user_message = message.strip()
+    
+    # Stage 1: Initial greeting and ask for name
+    if medical_bot.stage == "greeting":
+        bot_response = "Hello! I'm your AI medical assistant. Before we discuss your health concerns, could you please tell me your name?"
+        medical_bot.stage = "name"
+        medical_bot.add_to_history(user_message, bot_response)
+        return bot_response
+    
+    # Stage 2: Collect name and ask for age
+    elif medical_bot.stage == "name":
+        medical_bot.patient_name = user_message
+        bot_response = f"Nice to meet you, {medical_bot.patient_name}! Could you please tell me your age?"
+        medical_bot.stage = "age"
+        medical_bot.add_to_history(user_message, bot_response)
+        return bot_response
+    
+    # Stage 3: Collect age and start medical consultation
+    elif medical_bot.stage == "age":
+        medical_bot.patient_age = user_message
+        bot_response = f"Thank you, {medical_bot.patient_name}! Now, what brings you here today? Please tell me about any symptoms or health concerns you're experiencing."
+        medical_bot.stage = "symptoms"
+        medical_bot.add_to_history(user_message, bot_response)
+        return bot_response
+    
+    # Stage 4: Medical consultation - gather symptoms with intelligent follow-ups
+    elif medical_bot.stage == "symptoms":
+        medical_bot.medical_turns += 1
         
-        # Save final response to memory
-        memory.save_context({"input": message}, {"output": final_response})
+        # If we've had enough turns, move to diagnosis
+        if medical_bot.medical_turns >= 4:
+            medical_bot.stage = "diagnosis"
+            return generate_final_diagnosis(user_message)
         
-        return final_response
+        # Generate intelligent follow-up questions
+        try:
+            prompt = medical_bot.build_llama_prompt(user_message)
+            inputs = tokenizer(
+                prompt, 
+                return_tensors="pt", 
+                max_length=1024, 
+                truncation=True
+            ).to(model.device)
+            
+            with torch.no_grad():
+                outputs = model.generate(
+                    inputs.input_ids,
+                    attention_mask=inputs.attention_mask,
+                    max_new_tokens=200,
+                    temperature=0.8,
+                    top_p=0.95,
+                    do_sample=True,
+                    pad_token_id=tokenizer.pad_token_id
+                )
+            
+            full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            bot_response = full_response.split('[/INST]')[-1].strip()
+            
+            # Clean up the response
+            bot_response = bot_response.replace('<s>', '').replace('</s>', '').strip()
+            
+            medical_bot.add_to_history(user_message, bot_response)
+            return bot_response
+            
+        except Exception as e:
+            bot_response = f"I understand. Could you tell me more about how long you've been experiencing this and if there are any specific triggers or patterns you've noticed?"
+            medical_bot.add_to_history(user_message, bot_response)
+            return bot_response
+    
+    # Stage 5: Final diagnosis and treatment recommendations
+    elif medical_bot.stage == "diagnosis":
+        return generate_final_diagnosis(user_message)
+    
+    # Handle any questions after diagnosis
+    else:
+        # Check if they're asking about their name or previous information
+        if "name" in user_message.lower() and medical_bot.patient_name:
+            return f"Your name is {medical_bot.patient_name}."
+        elif "age" in user_message.lower() and medical_bot.patient_age:
+            return f"You told me you are {medical_bot.patient_age} years old."
+        else:
+            return "Is there anything else about your health concerns I can help you with today?"
+
+def generate_final_diagnosis(current_message):
+    """Generate final diagnosis using both models"""
+    global medical_bot
+    
+    # Add current message to history
+    medical_bot.add_to_history(current_message, "")
+    
+    # Compile complete patient information
+    patient_info = f"""
+Patient Name: {medical_bot.patient_name}
+Patient Age: {medical_bot.patient_age}
+
+Complete Consultation History:
+"""
+    
+    for exchange in medical_bot.conversation_history[:-1]:  # Exclude the empty last entry
+        patient_info += f"Doctor: {exchange['bot']}\n"
+        patient_info += f"Patient: {exchange['user']}\n"
+    
+    patient_info += f"Patient: {current_message}\n"
+    
+    # Get diagnosis from Meditron
+    meditron_assessment = get_meditron_diagnosis(patient_info)
+    
+    # Generate comprehensive response
+    final_response = f"""Thank you for providing all this information, {medical_bot.patient_name}. Based on our consultation, here is my assessment:
+
+**MEDICAL ASSESSMENT & RECOMMENDATIONS:**
+
+{meditron_assessment}
+
+**IMPORTANT DISCLAIMER:** This assessment is for informational purposes only and should not replace professional medical advice. Please consult with a healthcare provider for proper diagnosis and treatment.
+
+**NEXT STEPS:** I recommend scheduling an appointment with your primary care physician or appropriate specialist for further evaluation and personalized treatment.
+
+Is there anything specific about this assessment you'd like me to clarify?"""
+    
+    # Update conversation history with final response
+    medical_bot.conversation_history[-1]["bot"] = final_response
+    medical_bot.stage = "complete"
+    
+    return final_response
 
-# Create the Gradio interface
+# Create Gradio interface
 demo = gr.ChatInterface(
-    fn=generate_response,
-    title="🩺 AI Medical Assistant",
-    description="I'll ask for your basic information first, then gather details about your symptoms to provide medical insights.",
+    fn=medical_chat_response,
+    title="🩺 AI Medical Assistant with Memory",
+    description="I'm an AI medical assistant that will remember our conversation. I'll first ask for your basic information, then gather details about your symptoms through intelligent follow-up questions, and finally provide a medical assessment.",
     examples=[
+        "Hello, I need medical help",
         "I have a persistent cough",
         "I've been having headaches",
         "My stomach hurts"
     ],
-    theme="soft"
+    theme="soft",
+    retry_btn=None,
+    undo_btn=None,
+    clear_btn="🔄 Start New Consultation"
 )
 
 if __name__ == "__main__":