Update app.py

app.py

@@ -16,6 +16,7 @@ if torch.cuda.is_available():
 # Step 1: Load the PDF and create a vector store
 @st.cache_resource
 def load_pdf_to_vectorstore(pdf_path):
+    # Load and split PDF
     loader = PyPDFLoader(pdf_path)
     documents = loader.load()
     
@@ -27,6 +28,7 @@ def load_pdf_to_vectorstore(pdf_path):
     
     chunks = text_splitter.split_documents(documents)
     
+    # Create embeddings and vector store
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2"
     )
@@ -37,10 +39,11 @@ def load_pdf_to_vectorstore(pdf_path):
 # Step 2: Initialize the LaMini model
 @st.cache_resource
 def setup_model():
-    model_id = "MBZUAI/LaMini-Flan-T5-248M"
+    model_id = "MBZUAI/LaMini-Flan-T5-248M"  # Using smaller model for faster inference
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     model = AutoModelForSeq2SeqLM.from_pretrained(
         model_id,
+        # Removed low_cpu_mem_usage parameter
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
     )
     
@@ -51,16 +54,17 @@ def setup_model():
         "text2text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_length=512,  # Increased max length for better context
+        max_length=256,
         do_sample=False,
         temperature=0.3,
         top_p=0.95,
         device=0 if torch.cuda.is_available() else -1,
+        batch_size=1
     )
     return pipe
 
-# Step 3: Generate response with conversation history
-def generate_response(pipe, vectorstore, user_input, chat_history):
+# Step 3: Generate a response using the model and vector store
+def generate_response(pipe, vectorstore, user_input):
     # Get relevant context
     docs = vectorstore.similarity_search(user_input, k=2)
     context = "\n".join([
@@ -68,89 +72,83 @@ def generate_response(pipe, vectorstore, user_input, chat_history):
         for doc in docs
     ])
     
-    # Format conversation history
-    history_text = "\n".join(
-        [f"{msg['role'].capitalize()}: {msg['content']}" 
-         for msg in chat_history]
-    ) if chat_history else "No previous conversation"
-
-    # Create contextual prompt
-    prompt_template = PromptTemplate(
-        input_variables=["history", "context", "question"],
+    # Create prompt
+    prompt = PromptTemplate(
+        input_variables=["context", "question"],
         template="""
-        Previous Conversation:
-        {history}
-
-        Medical Context:
-        {context}
-
-        Current Question: {question}
-
-        Instructions:
-        1. Answer based on context and conversation history
-        2. Cite page numbers when possible
-        3. If unsure, recommend consulting a professional
-        4. Maintain natural conversation flow
-
-        Assistant Response:
-        """
-    )
-    
-    prompt = prompt_template.format(
-        history=history_text,
-        context=context,
-        question=user_input
+        Using the following medical text excerpts, answer the question.
+        If the information isn't clearly provided in the context, or if you're unsure, please say so and recommend consulting a healthcare professional.
+        
+        Context: {context}
+        
+        Question: {question}
+        
+        Answer (citing relevant page numbers when possible):"""
     )
     
-    # Generate response
-    response = pipe(prompt, max_length=512)[0]['generated_text']
+    # Generate response using the new method
+    prompt_text = prompt.format(context=context, question=user_input)
+    response = pipe(prompt_text)[0]['generated_text']
     
     return response
 
-# Streamlit UI with conversation memory
+# Cache responses for repeated questions
+@st.cache_data
+def cached_generate_response(user_input, _pipe, _vectorstore):
+    return generate_response(_pipe, _vectorstore, user_input)
+
+# Batch processing for multiple questions
+def batch_generate_responses(pipe, vectorstore, questions, batch_size=4):
+    responses = []
+    for i in range(0, len(questions), batch_size):
+        batch = questions[i:i + batch_size]
+        batch_responses = [generate_response(pipe, vectorstore, q) for q in batch]
+        responses.extend(batch_responses)
+    return responses
+
+# Streamlit UI
 def main():
     st.title("Medical Chatbot Assistant 🏥")
     
-    # Initialize session state for chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-    
-    # Load resources
+    # Use the PDF file from the root directory
     pdf_path = "Medical_book.pdf"
+    
     if os.path.exists(pdf_path):
-        with st.spinner("Initializing system..."):
+        # Initialize progress
+        progress_text = "Operation in progress. Please wait."
+        
+        # Load vector store and model with progress indication
+        with st.spinner("Loading PDF and initializing model..."):
             vectorstore = load_pdf_to_vectorstore(pdf_path)
             pipe = setup_model()
+            st.success("Ready to answer questions!")
         
-        # Display chat messages
+        # Create a chat-like interface
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
+
+        # Display chat history
         for message in st.session_state.messages:
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
 
-        # User input handling
+        # User input
         if prompt := st.chat_input("Ask your medical question:"):
-            # Add user message to history
+            # Add user message to chat history
             st.session_state.messages.append({"role": "user", "content": prompt})
-            
-            # Generate response with conversation context
-            with st.spinner("Analyzing question..."):
-                response = generate_response(
-                    pipe,
-                    vectorstore,
-                    prompt,
-                    chat_history=st.session_state.messages[:-1]  # Exclude current prompt
-                )
-            
-            # Add and display assistant response
-            st.session_state.messages.append({"role": "assistant", "content": response})
-            
-            # Display conversation
             with st.chat_message("user"):
                 st.markdown(prompt)
+
+            # Generate and display response
             with st.chat_message("assistant"):
-                st.markdown(response)
+                with st.spinner("Generating response..."):
+                    response = cached_generate_response(prompt, pipe, vectorstore)
+                    st.markdown(response)
+                    # Add assistant message to chat history
+                    st.session_state.messages.append({"role": "assistant", "content": response})
+                    
     else:
-        st.error("Medical reference book not found!")
+        st.error("The file 'Medical_book.pdf' was not found in the root directory.")
+
 
-if __name__ == "__main__":
-    main()
+main()