Update app.py

app.py

@@ -1,169 +1,97 @@
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
-from langchain_core.runnables.history import RunnableWithMessageHistory
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_community.chat_message_histories import ChatMessageHistory
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Model configuration
 MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
 
-# System prompt that guides the bot's behavior
-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. Start by greeting the user politely and ask 
-them to describe their health concern. For each user reply, ask only 1 or 2 follow-up questions at a time to gather more details. 
-Be structured and thorough in your questioning. Organize the information into categories: symptoms, duration, severity, 
-possible causes, past medical history, medications, allergies, habits (e.g., smoking, alcohol), and family history. 
-Always confirm and summarize what the user tells you. Respond empathetically and clearly. If unsure, ask for clarification. 
-**IMPORTANT**make a final diagnosis or suggest treatments.
-Wait for the user's answer before asking more questions.
-"""
+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.
 
-print("Loading model...")
-try:
-    # Initialize the tokenizer and model
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        torch_dtype="auto",
-        device_map="auto"
-    )
-    
-    # Create a pipeline for text generation
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.9,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    
-    llm = HuggingFacePipeline(pipeline=pipe)
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    # Fallback to a smaller model or provide an error message
-    raise
+Ask 1-2 follow-up questions at a time to gather more details about:
+- Detailed description of symptoms
+- Duration (when did it start?)
+- Severity (scale of 1-10)
+- Aggravating or alleviating factors
+- Related symptoms
+- Medical history
+- Current medications and allergies
+
+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.
 
-# Modify the prompt template with a clearer structure to prevent system prompt leakage
-prompt = ChatPromptTemplate.from_messages([
-    ("system", SYSTEM_PROMPT),
-    MessagesPlaceholder(variable_name="history"),
-    ("human", "{input}"),
-    ("system", "Remember to respond as Virtual Doctor without including system instructions in your reply.")
-])
+Respond empathetically and clearly. Always be professional and thorough."""
 
-# Memory store to maintain conversation history
-store = {}
+print("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+print("Model loaded successfully!")
 
-def get_session_history(session_id: str) -> ChatMessageHistory:
-    """Get or create a chat history for the given session ID"""
-    if session_id not in store:
-        store[session_id] = ChatMessageHistory()
-    return store[session_id]
+# Conversation state tracking
+conversation_turns = {}
 
-# Create a more robust filtering chain that will intercept the model's responses
-def filter_response(response_text):
-    """Filter out system prompts and format the response correctly"""
-    # Remove any system prompt references
-    if "system" in response_text.lower() and ("your goal is" in response_text.lower() or "professional virtual doctor" in response_text.lower()):
-        # Find the actual doctor response after any system text
-        for marker in ["Virtual Doctor:", "Virtual doctor:", "Human:"]:
-            if marker in response_text:
-                parts = response_text.split(marker)
-                if len(parts) > 1:
-                    # Get the last part after any system prompts
-                    response_text = parts[-1].strip()
-                    break
+def build_llama2_prompt(system_prompt, history, user_input):
+    """Format the conversation history and user input for Llama-2 chat models."""
+    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"
     
-    # Remove any remaining system prompt text or instructions
-    filtered_text = []
-    skip_line = False
-    for line in response_text.split('\n'):
-        lower_line = line.lower()
-        if any(phrase in lower_line for phrase in [
-            "system:", "your goal is", "start by greeting", "wait for the user", 
-            "do not make a final diagnosis", "be structured", "ask only 1 or 2"
-        ]):
-            skip_line = True
-        elif any(marker in line for marker in ["Virtual Doctor:", "Virtual doctor:", "Hello", "Thank you"]):
-            skip_line = False
-        
-        if not skip_line:
-            filtered_text.append(line)
+    # Add conversation history
+    for user_msg, assistant_msg in history:
+        prompt += f"{user_msg} [/INST] {assistant_msg} </s><s>[INST] "
     
-    clean_text = '\n'.join(filtered_text).strip()
+    # Add the current user input
+    prompt += f"{user_input} [/INST] "
     
-    # Ensure proper formatting with "Virtual Doctor:" prefix
-    if not clean_text.startswith("Virtual Doctor:") and not clean_text.startswith("Virtual doctor:"):
-        clean_text = f"Virtual Doctor: {clean_text}"
-    
-    return clean_text
+    return prompt
 
-# Chain with memory
-chain = prompt | llm
-chain_with_history = RunnableWithMessageHistory(
-    chain,
-    get_session_history,
-    input_messages_key="input",
-    history_messages_key="history"
-)
-
-# Our handler for chat interactions
-@spaces.GPU  # Request GPU for this space
-def gradio_chat(user_message, history):
-    """Process the user message and return the chatbot response"""
-    # Use a unique session ID in production
+@spaces.GPU
+def generate_response(message, history):
+    """Generate a response using the Llama-2 model with proper formatting."""
+    # Track conversation turns
     session_id = "default-session"
+    if session_id not in conversation_turns:
+        conversation_turns[session_id] = 0
+    conversation_turns[session_id] += 1
+    
+    # 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[session_id] >= 4:
+        prompt = prompt.replace("[/INST] ", "[/INST] Now summarize what you've learned and suggest when professional care may be needed. ")
+    
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
-    # Invoke the chain with history
-    try:
-        response = chain_with_history.invoke(
-            {"input": user_message},
-            config={"configurable": {"session_id": session_id}}
+    # Generate the response
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs.input_ids,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
-        
-        # Extract the text from the response
-        response_text = response.content if hasattr(response, "content") else str(response)
-        
-        # Apply our filtering function to clean up the response
-        clean_response = filter_response(response_text)
-        
-        return clean_response
-    except Exception as e:
-        print(f"Error processing message: {e}")
-        return "Virtual Doctor: I apologize, but I'm experiencing technical difficulties. Please try again."
-
-# Customize the CSS for better appearance
-css = """
-.gradio-container {
-    font-family: 'Arial', sans-serif;
-}
-.chat-bot .bot-message {
-    background-color: #f0f7ff !important;
-}
-.chat-bot .user-message {
-    background-color: #e6f7e6 !important;
-}
-"""
+    
+    # Decode and extract only the assistant's response
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    assistant_response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
+    
+    return assistant_response
 
 # Create the Gradio interface
 demo = gr.ChatInterface(
-    fn=gradio_chat,
-    title="Medbot Chatbot (Llama-2 + LangChain + Gradio)",
-    description="Medical chatbot using Llama-2-7b-chat-hf, LangChain memory, and Gradio UI.",
+    fn=generate_response,
+    title="Medical Assistant Chatbot",
+    description="Ask about your symptoms and I'll help gather relevant information.",
     examples=[
         "I have a cough and my throat hurts",
         "I've been having headaches for a week",
         "My stomach has been hurting since yesterday"
     ],
-    css=css
+    theme="soft"
 )
 
-# Launch the app
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()