Update app.py

app.py

@@ -1,116 +1,81 @@
 import gradio as gr
 import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer
+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
 
-# Define the medical assistant system prompt
-SYSTEM_PROMPT = """
-You are a knowledgeable medical assistant. Follow these steps in order:
+MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
 
-1) INITIAL ASSESSMENT: First, warmly greet the user and ask about their primary concern.
+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. Based on their input, ask follow-up questions to gather as much relevant information as possible. "
+    "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. "
+    "Ask one or two questions at a time, and wait for user input."
+)
 
-2) ASK FOLLOW-UP QUESTIONS: For any health concern mentioned, systematically gather information by asking 1-2 specific follow-up questions at a time 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
-   - Family history of similar conditions
+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!")
 
-3) SUMMARIZE FINDINGS: Once you have gathered sufficient information (at least 4-5 exchanges with the user), organize what you've learned into clear categories:
-   - Symptoms
-   - Duration
-   - Severity
-   - Possible Causes
-   - Medications/Allergies
-   - Family History
+# LangChain prompt
+prompt = ChatPromptTemplate.from_messages([
+    ("system", SYSTEM_PROMPT),
+    MessagesPlaceholder(variable_name="history"),
+    ("human", "{input}")
+])
 
-4) PROVIDE RECOMMENDATIONS: Only after gathering comprehensive information, suggest:
-   - One specific OTC medicine with proper adult dosing
-   - One practical home remedy
-   - When they should seek professional medical care
+# Memory store
+store = {}
 
-5) END WITH DISCLAIMER: Always end with a clear medical disclaimer that you are not a licensed medical professional and your suggestions are not a substitute for professional medical advice.
+def get_session_history(session_id: str) -> ChatMessageHistory:
+    if session_id not in store:
+        store[session_id] = ChatMessageHistory()
+    return store[session_id]
 
-IMPORTANT: Do not skip ahead to recommendations without gathering comprehensive information through multiple exchanges. Your primary goal is information gathering through thoughtful questions.
-"""
+# Chain with memory
+chain = prompt | llm
+chain_with_history = RunnableWithMessageHistory(
+    chain,
+    get_session_history,
+    input_messages_key="input",
+    history_messages_key="history"
+)
 
-# Define model options
-MODELS = {
-    "TinyLlama-1.1B": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-    "Llama-2-7b": "meta-llama/Llama-2-7b-chat-hf"
-}
-
-# Global variables to store loaded models and tokenizers
-loaded_models = {}
-loaded_tokenizers = {}
-
-def load_model(model_name):
-    """Load model and tokenizer if not already loaded"""
-    if model_name not in loaded_models:
-        print(f"Loading {model_name}...")
-        model_path = MODELS[model_name]
-        tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            torch_dtype="auto",
-            device_map="auto"  # Use GPU if available
-        )
-        loaded_models[model_name] = model
-        loaded_tokenizers[model_name] = tokenizer
-        print(f"{model_name} loaded successfully!")
-    return loaded_models[model_name], loaded_tokenizers[model_name]
-
-# Pre-load the smaller model to start with
-print("Pre-loading TinyLlama model...")
-load_model("TinyLlama-1.1B")
-
-@spaces.GPU  # Required by ZeroGPU!
-def generate_response(message, history, model_choice):
-    """Generate a response from the selected model"""
-    # Load the selected model if not already loaded
-    model, tokenizer = load_model(model_choice)
-    
-    # Format the prompt based on the history and system prompt
-    formatted_prompt = SYSTEM_PROMPT + "\n\n"
-    
-    # Add conversation history
-    for human, assistant in history:
-        formatted_prompt += f"User: {human}\nAssistant: {assistant}\n"
-    
-    # Add the current message
-    formatted_prompt += f"User: {message}\nAssistant:"
-    
-    # Generate the response
-    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        inputs["input_ids"],
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
+@spaces.GPU
+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}}
     )
-    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-    
-    return response.strip()
+    # LangChain returns a "AIMessage" object; get text
+    return response.content if hasattr(response, "content") else str(response)
 
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Medical Assistant Chatbot")
-    gr.Markdown("This chatbot uses LLM models to provide medical information and assistance. Please note that this is not a substitute for professional medical advice.")
-    
-    with gr.Row():
-        model_dropdown = gr.Dropdown(
-            choices=list(MODELS.keys()),
-            value="TinyLlama-1.1B",
-            label="Select Model"
-        )
-    
-    chatbot = gr.ChatInterface(
-        fn=lambda message, history, model_choice: generate_response(message, history, model_choice),
-        additional_inputs=[model_dropdown],
-    )
+# Gradio UI
+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."
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()