Update app.py

app.py

@@ -3,11 +3,11 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import spaces
 
-# Model configuration
-LLAMA_MODEL = "meta-llama/Llama-2-7b-chat-hf"
-MEDITRON_MODEL = "epfl-llm/meditron-7b"
+# Model configuration - Using only Me-LLaMA 13B-chat
+ME_LLAMA_MODEL = "clinicalnlplab/me-llama-13b-chat"
 
-SYSTEM_PROMPT = """You are a professional virtual doctor. Your goal is to collect detailed information about the user's health condition, symptoms, medical history, medications, lifestyle, and other relevant data.
+# System prompts for different phases
+CONSULTATION_PROMPT = """You are a professional virtual doctor. Your goal is to collect detailed information about the user's health condition, symptoms, medical history, medications, lifestyle, and other relevant data.
 Ask 1-2 follow-up questions at a time to gather more details about:
 - Detailed description of symptoms
 - Duration (when did it start?)
@@ -19,30 +19,23 @@ Ask 1-2 follow-up questions at a time to gather more details about:
 After collecting sufficient information (4-5 exchanges), summarize findings and suggest when they should seek professional care. Do NOT make specific diagnoses or recommend specific treatments.
 Respond empathetically and clearly. Always be professional and thorough."""
 
-MEDITRON_PROMPT = """<|im_start|>system
-You are a specialized medical assistant focusing ONLY on suggesting over-the-counter medicines and home remedies based on patient information.
-Based on the following patient information, provide ONLY:
+MEDICINE_PROMPT = """You are a specialized medical assistant. Based on the patient information gathered, provide:
 1. One specific over-the-counter medicine with proper adult dosing instructions
 2. One practical home remedy that might help
 3. Clear guidance on when to seek professional medical care
+
 Be concise, practical, and focus only on general symptom relief. Do not diagnose. Include a disclaimer that you are not a licensed medical professional.
-<|im_end|>
-<|im_start|>user
-Patient information: {patient_info}
-<|im_end|>
-<|im_start|>assistant
-"""
 
-# Global variables to store models (will be loaded lazily)
-llama_model = None
-llama_tokenizer = None
-meditron_model = None
-meditron_tokenizer = None
+Patient information: {patient_info}"""
+
+# Global variables
+me_llama_model = None
+me_llama_tokenizer = None
 conversation_turns = 0
 patient_data = []
 
-def build_llama2_prompt(system_prompt, history, user_input):
-    """Format the conversation history and user input for Llama-2 chat models."""
+def build_me_llama_prompt(system_prompt, history, user_input):
+    """Format the conversation for Me-LLaMA chat model."""
     prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
     
     # Add conversation history
@@ -55,126 +48,143 @@ def build_llama2_prompt(system_prompt, history, user_input):
     return prompt
 
 @spaces.GPU
-def load_models_if_needed():
-    """Load models only when GPU is available and only if not already loaded."""
-    global llama_model, llama_tokenizer, meditron_model, meditron_tokenizer
-    
-    if llama_model is None:
-        print("Loading Llama-2 model...")
-        llama_tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
-        llama_model = AutoModelForCausalLM.from_pretrained(
-            LLAMA_MODEL,
-            torch_dtype=torch.float16,
-            device_map="auto"
-        )
-        print("Llama-2 model loaded successfully!")
-    
-    if meditron_model is None:
-        print("Loading Meditron model...")
-        meditron_tokenizer = AutoTokenizer.from_pretrained(MEDITRON_MODEL)
-        meditron_model = AutoModelForCausalLM.from_pretrained(
-            MEDITRON_MODEL,
+def load_model_if_needed():
+    """Load Me-LLaMA model only when GPU is available."""
+    global me_llama_model, me_llama_tokenizer
+    
+    if me_llama_model is None:
+        print("Loading Me-LLaMA 13B-chat model...")
+        me_llama_tokenizer = AutoTokenizer.from_pretrained(ME_LLAMA_MODEL)
+        me_llama_model = AutoModelForCausalLM.from_pretrained(
+            ME_LLAMA_MODEL,
             torch_dtype=torch.float16,
-            device_map="auto"
+            device_map="auto",
+            trust_remote_code=True
         )
-        print("Meditron model loaded successfully!")
+        print("Me-LLaMA 13B-chat model loaded successfully!")
 
 @spaces.GPU
-def get_meditron_suggestions(patient_info):
-    """Use Meditron model to generate medicine and remedy suggestions."""
-    load_models_if_needed()  # Ensure models are loaded
+def generate_medicine_suggestions(patient_info):
+    """Use Me-LLaMA to generate medicine and remedy suggestions."""
+    load_model_if_needed()
     
-    prompt = MEDITRON_PROMPT.format(patient_info=patient_info)
-    inputs = meditron_tokenizer(prompt, return_tensors="pt")
+    # Create a simple prompt for medicine suggestions
+    prompt = f"<s>[INST] {MEDICINE_PROMPT.format(patient_info=patient_info)} [/INST] "
+    
+    inputs = me_llama_tokenizer(prompt, return_tensors="pt")
     
     # Move inputs to the same device as the model
     if torch.cuda.is_available():
-        inputs = {k: v.to(meditron_model.device) for k, v in inputs.items()}
+        inputs = {k: v.to(me_llama_model.device) for k, v in inputs.items()}
     
     with torch.no_grad():
-        outputs = meditron_model.generate(
+        outputs = me_llama_model.generate(
             inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
-            max_new_tokens=256,
+            max_new_tokens=300,
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
-            pad_token_id=meditron_tokenizer.eos_token_id
+            pad_token_id=me_llama_tokenizer.eos_token_id
         )
     
-    suggestion = meditron_tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    suggestion = me_llama_tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
     return suggestion
 
 @spaces.GPU
 def generate_response(message, history):
-    """Generate a response using both models."""
+    """Generate response using only Me-LLaMA for both consultation and medicine suggestions."""
     global conversation_turns, patient_data
     
-    # Load models if needed
-    load_models_if_needed()
+    # Load model if needed
+    load_model_if_needed()
     
     # Track conversation turns
     conversation_turns += 1
     
-    # Store the entire conversation for reference
+    # Store patient data
     patient_data.append(message)
     
-    # Build the prompt with proper Llama-2 formatting
-    prompt = build_llama2_prompt(SYSTEM_PROMPT, history, message)
-    
-    # Add summarization instruction after 4 turns
-    if conversation_turns >= 4:
-        prompt = prompt.replace("[/INST] ", "[/INST] Now summarize what you've learned and suggest when professional care may be needed. ")
-    
-    inputs = llama_tokenizer(prompt, return_tensors="pt")
-    
-    # Move inputs to the same device as the model
-    if torch.cuda.is_available():
-        inputs = {k: v.to(llama_model.device) for k, v in inputs.items()}
-    
-    # Generate the Llama-2 response
-    with torch.no_grad():
-        outputs = llama_model.generate(
-            inputs["input_ids"],
-            attention_mask=inputs["attention_mask"],
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=llama_tokenizer.eos_token_id
+    # Phase 1-3: Information gathering
+    if conversation_turns < 4:
+        # Build consultation prompt
+        prompt = build_me_llama_prompt(CONSULTATION_PROMPT, history, message)
+        
+        inputs = me_llama_tokenizer(prompt, return_tensors="pt")
+        
+        # Move inputs to the same device as the model
+        if torch.cuda.is_available():
+            inputs = {k: v.to(me_llama_model.device) for k, v in inputs.items()}
+        
+        # Generate consultation response
+        with torch.no_grad():
+            outputs = me_llama_model.generate(
+                inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                max_new_tokens=400,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=me_llama_tokenizer.eos_token_id
+            )
+        
+        # Decode response
+        full_response = me_llama_tokenizer.decode(outputs[0], skip_special_tokens=False)
+        response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
+        
+        return response
+    
+    # Phase 4+: Summary and medicine suggestions
+    else:
+        # First, get summary from consultation
+        summary_prompt = build_me_llama_prompt(
+            CONSULTATION_PROMPT + "\n\nNow summarize what you've learned and suggest when professional care may be needed.",
+            history, 
+            message
         )
-    
-    # Decode and extract Llama-2's response
-    full_response = llama_tokenizer.decode(outputs[0], skip_special_tokens=False)
-    llama_response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
-    
-    # After 4 turns, add medicine suggestions from Meditron
-    if conversation_turns >= 4:
-        # Collect full patient conversation
-        full_patient_info = "\n".join(patient_data) + "\n\nSummary: " + llama_response
         
-        # Get medicine suggestions
-        medicine_suggestions = get_meditron_suggestions(full_patient_info)
+        inputs = me_llama_tokenizer(summary_prompt, return_tensors="pt")
+        
+        if torch.cuda.is_available():
+            inputs = {k: v.to(me_llama_model.device) for k, v in inputs.items()}
+        
+        # Generate summary
+        with torch.no_grad():
+            outputs = me_llama_model.generate(
+                inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                max_new_tokens=400,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=me_llama_tokenizer.eos_token_id
+            )
         
-        # Format final response
+        summary_response = me_llama_tokenizer.decode(outputs[0], skip_special_tokens=False)
+        summary = summary_response.split('[/INST]')[-1].split('</s>')[0].strip()
+        
+        # Then get medicine suggestions using the same model
+        full_patient_info = "\n".join(patient_data) + f"\n\nMedical Summary: {summary}"
+        medicine_suggestions = generate_medicine_suggestions(full_patient_info)
+        
+        # Combine both responses
         final_response = (
-            f"{llama_response}\n\n"
-            f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
-            f"{medicine_suggestions}"
+            f"**MEDICAL SUMMARY:**\n{summary}\n\n"
+            f"**MEDICATION AND HOME CARE SUGGESTIONS:**\n{medicine_suggestions}\n\n"
+            f"**DISCLAIMER:** This is AI-generated advice for informational purposes only. Please consult a licensed healthcare provider for proper medical diagnosis and treatment."
         )
+        
         return final_response
-    
-    return llama_response
 
 # Create the Gradio interface
 demo = gr.ChatInterface(
     fn=generate_response,
-    title="Medical Assistant with Medicine Suggestions",
-    description="Tell me about your symptoms, and after gathering enough information, I'll suggest potential remedies.",
+    title="🏥 Complete Medical Assistant - Me-LLaMA 13B",
+    description="Comprehensive medical consultation powered by Me-LLaMA 13B-chat. One model handles both consultation and medicine suggestions. Tell me about your symptoms!",
     examples=[
-        "I have a cough and my throat hurts",
-        "I've been having headaches for a week",
-        "My stomach has been hurting since yesterday"
+        "I have a persistent cough and sore throat for 3 days",
+        "I've been having severe headaches and feel dizzy",
+        "My stomach hurts and I feel nauseous after eating"
     ],
     theme="soft"
 )