Update app.py

app.py

@@ -6,54 +6,66 @@ from langchain_core.runnables.history import RunnableWithMessageHistory
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_community.chat_message_histories import ChatMessageHistory
 
+# Model configuration
 MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
 
-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. "
-    "Do NOT make a final diagnosis or suggest treatments. You are only here to collect and organize medical data to support a licensed physician. "
-    "Wait for the user's answer before asking more questions."
-)
-
+# 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. 
+Do NOT make a final diagnosis or suggest treatments. You are only here to collect and organize medical data to support a licensed physician. 
+Wait for the user's answer before asking more questions.
+"""
 
 print("Loading model...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype="auto",
-    device_map="auto"
-)
-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!")
+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
 
-# LangChain prompt
+# LangChain prompt template
 prompt = ChatPromptTemplate.from_messages([
     ("system", SYSTEM_PROMPT),
     MessagesPlaceholder(variable_name="history"),
     ("human", "{input}")
 ])
 
-# Memory store
+# Memory store to maintain conversation history
 store = {}
 
 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]
 
-# Chain with memory
+# Create a chain with memory
 chain = prompt | llm
 chain_with_history = RunnableWithMessageHistory(
     chain,
@@ -62,22 +74,57 @@ chain_with_history = RunnableWithMessageHistory(
     history_messages_key="history"
 )
 
-@spaces.GPU
+# Our handler for chat interactions
+@spaces.GPU  # Request GPU for this space
 def gradio_chat(user_message, history):
-    session_id = "default-session"  # For demo; can be made unique per user
-    response = chain_with_history.invoke(
-        {"input": user_message},
-        config={"configurable": {"session_id": session_id}}
-    )
-    # LangChain returns a "AIMessage" object; get text
-    return response.content if hasattr(response, "content") else str(response)
+    """Process the user message and return the chatbot response"""
+    # Use a unique session ID in production
+    session_id = "default-session"
+    
+    # Invoke the chain with history
+    try:
+        response = chain_with_history.invoke(
+            {"input": user_message},
+            config={"configurable": {"session_id": session_id}}
+        )
+        
+        # Extract the text from the response
+        response_text = response.content if hasattr(response, "content") else str(response)
+        
+        # Format as "Virtual doctor: " response to match the expected format
+        formatted_response = f"Virtual doctor: {response_text}"
+        
+        return formatted_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;
+}
+"""
 
-# Gradio UI
+# 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."
+    description="Medical chatbot using Llama-2-7b-chat-hf, LangChain memory, and Gradio UI.",
+    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
 )
 
+# Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=False)