Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import torch
+import torch.backends.cudnn as cudnn
 import streamlit as st
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
@@ -8,6 +9,10 @@ from langchain_community.vectorstores import FAISS
 from langchain.prompts import PromptTemplate
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
+# Enable CUDA optimizations if available
+if torch.cuda.is_available():
+    cudnn.benchmark = True
+
 # Step 1: Load the PDF and create a vector store
 @st.cache_resource
 def load_pdf_to_vectorstore(pdf_path):
@@ -16,8 +21,8 @@ def load_pdf_to_vectorstore(pdf_path):
     documents = loader.load()
     
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=50,
+        chunk_size=1000,
+        chunk_overlap=20,
         separators=["\n\n", "\n", ".", " ", ""]
     )
     
@@ -34,33 +39,40 @@ 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-783M"
+    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)
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        low_cpu_mem_usage=True
+    )
+    
+    if torch.cuda.is_available():
+        model = model.cuda()
     
     pipe = pipeline(
         "text2text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_length=512,
-        do_sample=True,
+        max_length=256,
+        do_sample=False,
         temperature=0.3,
         top_p=0.95,
-        device=0 if torch.cuda.is_available() else -1
+        device=0 if torch.cuda.is_available() else -1,
+        batch_size=1
     )
     return pipe
 
 # Step 3: Generate a response using the model and vector store
 def generate_response(pipe, vectorstore, user_input):
     # Get relevant context
-
-     # Increase k for more context if needed
-    docs = vectorstore.similarity_search(user_input, k=4)  # increased from 3
-    
-    # Add document metadata (like page numbers) to help track sources
-    context = "\n".join([f"Page {doc.metadata.get('page', 'unknown')}: {doc.page_content}" for doc in docs])
+    docs = vectorstore.similarity_search(user_input, k=2)
+    context = "\n".join([
+        f"Page {doc.metadata.get('page', 'unknown')}: {doc.page_content}"
+        for doc in docs
+    ])
     
-    # Enhanced prompt template
+    # Create prompt
     prompt = PromptTemplate(
         input_variables=["context", "question"],
         template="""
@@ -80,6 +92,20 @@ def generate_response(pipe, vectorstore, user_input):
     
     return response
 
+# 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 🏥")
@@ -88,17 +114,39 @@ def main():
     pdf_path = "Medical_book.pdf"
     
     if os.path.exists(pdf_path):
-        # Load vector store and model
-        vectorstore = load_pdf_to_vectorstore(pdf_path)
-        pipe = setup_model()
+        # Initialize progress
+        progress_text = "Operation in progress. Please wait."
         
-        # User input
-        user_input = st.text_input("Ask your medical question:")
+        # 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!")
         
-        if user_input:
-            with st.spinner("Generating response..."):
-                response = generate_response(pipe, vectorstore, user_input)
-                st.write(response)
+        # 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
+        if prompt := st.chat_input("Ask your medical question:"):
+            # Add user message to chat history
+            st.session_state.messages.append({"role": "user", "content": prompt})
+            with st.chat_message("user"):
+                st.markdown(prompt)
+
+            # Generate and display response
+            with st.chat_message("assistant"):
+                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("The file 'Medical_book.pdf' was not found in the root directory.")