Update app.py

app.py

@@ -1,15 +1,16 @@
-import gradio as gr
+from langchain.chains import ConversationChain, LLMChain
+from langchain.prompts import PromptTemplate
+from langchain.llms import HuggingFacePipeline
+from langchain.memory import ConversationBufferMemory
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from typing import Annotated, List, Dict, Any
-from typing_extensions import TypedDict
-from langgraph.graph import StateGraph, START
-from langgraph.graph.message import add_messages
+import gradio as gr
 
 # Model configuration
 LLAMA_MODEL = "meta-llama/Llama-2-7b-chat-hf"
 MEDITRON_MODEL = "epfl-llm/meditron-7b"
 
+# System prompts
 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.
 Ask 1-2 follow-up questions at a time to gather more details about:
 - Detailed description of symptoms
@@ -37,213 +38,109 @@ Patient information: {patient_info}
 """
 
 print("Loading Llama-2 model...")
-tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-    
-model = AutoModelForCausalLM.from_pretrained(
+# Create LangChain wrapper for Llama-2
+llama_tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
+llama_model = AutoModelForCausalLM.from_pretrained(
     LLAMA_MODEL,
     torch_dtype=torch.float16,
     device_map="auto"
 )
+
+# Create a pipeline for LangChain
+llama_pipeline = pipeline(
+    "text-generation",
+    model=llama_model,
+    tokenizer=llama_tokenizer,
+    max_new_tokens=512,
+    temperature=0.7,
+    top_p=0.9,
+    do_sample=True
+)
+llama_llm = HuggingFacePipeline(pipeline=llama_pipeline)
 print("Llama-2 model loaded successfully!")
 
 print("Loading Meditron model...")
 meditron_tokenizer = AutoTokenizer.from_pretrained(MEDITRON_MODEL)
-if meditron_tokenizer.pad_token is None:
-    meditron_tokenizer.pad_token = meditron_tokenizer.eos_token
-    
 meditron_model = AutoModelForCausalLM.from_pretrained(
     MEDITRON_MODEL,
     torch_dtype=torch.float16,
     device_map="auto"
 )
+# Create a pipeline for Meditron
+meditron_pipeline = pipeline(
+    "text-generation",
+    model=meditron_model,
+    tokenizer=meditron_tokenizer,
+    max_new_tokens=256,
+    temperature=0.7,
+    top_p=0.9,
+    do_sample=True
+)
+meditron_llm = HuggingFacePipeline(pipeline=meditron_pipeline)
 print("Meditron model loaded successfully!")
 
-# Define the state for our LangGraph
-class ChatbotState(TypedDict):
-    messages: Annotated[List, add_messages]
-    turn_count: int
-    patient_info: List[str]
-
-# Function to build Llama-2 prompt
-def build_llama2_prompt(messages):
-    """Format the conversation history for Llama-2 chat models."""
-    prompt = f"<s>[INST] <<SYS>>\n{SYSTEM_PROMPT}\n<</SYS>>\n\n"
-    
-    # Add conversation history
-    for i, msg in enumerate(messages[:-1]):
-        if i % 2 == 0:  # User message
-            prompt += f"{msg.content} [/INST] "
-        else:  # Assistant message
-            prompt += f"{msg.content} </s><s>[INST] "
-    
-    # Add the current user input
-    prompt += f"{messages[-1].content} [/INST] "
-    
-    return prompt
-
-# Function to get Llama-2 response
-def get_llama2_response(prompt, turn_count):
-    """Generate response from Llama-2 model."""
-    # Add summarization instruction after 4 turns
-    if turn_count >= 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)
-    
-    with torch.no_grad():
-        outputs = 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=tokenizer.pad_token_id
-        )
-    
-    full_response = tokenizer.decode(outputs[0], skip_special_tokens=False)
-    response = full_response.split('[/INST]')[-1].split('</s>')[0].strip()
-    
-    return response
-
-# Function to get Meditron suggestions
-def get_meditron_suggestions(patient_info):
-    """Generate medicine and remedy suggestions from Meditron model."""
-    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,
-            pad_token_id=meditron_tokenizer.pad_token_id
-        )
-    
-    suggestion = meditron_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    return suggestion
-
-# Define LangGraph nodes
-def process_user_input(state: ChatbotState) -> ChatbotState:
-    """Process user input and update state."""
-    # Extract the latest user message
-    user_message = state["messages"][-1].content
-    
-    # Update patient info
-    return {
-        "patient_info": state["patient_info"] + [user_message],
-        "turn_count": state["turn_count"] + 1
-    }
+# Create LangChain conversation with memory
+memory = ConversationBufferMemory(return_messages=True)
+conversation = ConversationChain(
+    llm=llama_llm,
+    memory=memory,
+    verbose=True
+)
 
-def generate_llama_response(state: ChatbotState) -> ChatbotState:
-    """Generate response using Llama-2 model."""
-    prompt = build_llama2_prompt(state["messages"])
-    response = get_llama2_response(prompt, state["turn_count"])
-    
-    return {"messages": [{"role": "assistant", "content": response}]}
+# Create a template for the Meditron model
+meditron_template = PromptTemplate(
+    input_variables=["patient_info"],
+    template=MEDITRON_PROMPT
+)
+meditron_chain = LLMChain(
+    llm=meditron_llm,
+    prompt=meditron_template,
+    verbose=True
+)
 
-def check_turn_count(state: ChatbotState) -> str:
-    """Check if we need to add medicine suggestions."""
-    if state["turn_count"] >= 4:
-        return "add_suggestions"
-    return "continue"
+# Track conversation turns
+conversation_turns = 0
+patient_data = []
 
-def add_medicine_suggestions(state: ChatbotState) -> ChatbotState:
-    """Add medicine suggestions from Meditron model."""
-    # Get the last assistant response
-    last_response = state["messages"][-1].content
-    
-    # Collect full patient conversation
-    full_patient_info = "\n".join(state["patient_info"]) + "\n\nSummary: " + last_response
+def generate_response(message, history):
+    global conversation_turns, patient_data
+    conversation_turns += 1
     
-    # Get medicine suggestions
-    medicine_suggestions = get_meditron_suggestions(full_patient_info)
+    # Store patient message
+    patient_data.append(message)
     
-    # Format final response
-    final_response = (
-        f"{last_response}\n\n"
-        f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
-        f"{medicine_suggestions}"
-    )
-    
-    # Return updated message
-    return {"messages": [{"role": "assistant", "content": final_response}]}
-
-# Build the LangGraph
-def build_graph():
-    """Build and return the LangGraph for our chatbot."""
-    graph = StateGraph(ChatbotState)
-    
-    # Add nodes
-    graph.add_node("process_input", process_user_input)
-    graph.add_node("generate_response", generate_llama_response)
-    graph.add_node("add_suggestions", add_medicine_suggestions)
+    # Format the prompt with system instructions
+    if conversation_turns >= 4:
+        # Add summarization instruction after 4 turns
+        prompt = f"{SYSTEM_PROMPT}\n\nNow summarize what you've learned and suggest when professional care may be needed.\n\n{message}"
+    else:
+        prompt = f"{SYSTEM_PROMPT}\n\n{message}"
     
-    # Add edges
-    graph.add_edge(START, "process_input")
-    graph.add_edge("process_input", "generate_response")
-    graph.add_conditional_edges(
-        "generate_response",
-        check_turn_count,
-        {
-            "add_suggestions": "add_suggestions",
-            "continue": END
-        }
-    )
-    graph.add_edge("add_suggestions", END)
+    # Generate response using LangChain conversation
+    llama_response = conversation.predict(input=prompt)
     
-    return graph.compile()
-
-# Initialize the graph
-chatbot_graph = build_graph()
-
-# Function for Gradio interface
-def chat_response(message, history):
-    """Generate chatbot response using LangGraph."""
-    # Initialize state if this is the first message
-    if not history:
-        state = {
-            "messages": [{"role": "user", "content": message}],
-            "turn_count": 0,
-            "patient_info": []
-        }
-    else:
-        # Convert history to messages format
-        messages = []
-        for user_msg, bot_msg in history:
-            messages.append({"role": "user", "content": user_msg})
-            messages.append({"role": "assistant", "content": bot_msg})
+    # 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
         
-        # Add current message
-        messages.append({"role": "user", "content": message})
+        # Get medicine suggestions using LangChain
+        medicine_suggestions = meditron_chain.run(patient_info=full_patient_info)
         
-        # Get turn count from history
-        turn_count = len(history)
-        
-        # Build patient info from history
-        patient_info = [user_msg for user_msg, _ in history]
-        
-        state = {
-            "messages": messages,
-            "turn_count": turn_count,
-            "patient_info": patient_info
-        }
-    
-    # Process through LangGraph
-    result = chatbot_graph.invoke(state)
+        # Format final response
+        final_response = (
+            f"{llama_response}\n\n"
+            f"--- MEDICATION AND HOME CARE SUGGESTIONS ---\n\n"
+            f"{medicine_suggestions}"
+        )
+        return final_response
     
-    # Return the latest assistant message
-    return result["messages"][-1].content
+    return llama_response
 
 # Create the Gradio interface
 demo = gr.ChatInterface(
-    fn=chat_response,
-    title="Medical Assistant with LangGraph",
+    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.",
     examples=[
         "I have a cough and my throat hurts",
@@ -254,4 +151,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()