Update app.py

app.py

@@ -446,7 +446,7 @@ def load_model():
     checkpoint = "MBZUAI/LaMini-T5-738M"
     tokenizer = AutoTokenizer.from_pretrained(checkpoint)
     model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
-    pipe = pipeline('text2text-generation', model=model, tokenizer=tokenizer, max_length=512)
+    pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, max_length=512)
     return HuggingFacePipeline(pipeline=pipe)
 
 # --- Extract PDF Text ---
@@ -461,44 +461,46 @@ def read_pdf(file):
         logging.error(f"Failed to extract text: {e}")
         return ""
 
-# --- Process Answer ---
-def process_answer(question, full_text):
-    # Save the full_text to a temporary file
+# --- Build Retriever (cached per session) ---
+@st.cache_resource
+def build_retriever(full_text):
+    # Save text to temp file
     with open("temp_text.txt", "w") as f:
         f.write(full_text)
 
     loader = TextLoader("temp_text.txt")
     docs = loader.load()
 
-    # Chunk the documents
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
+    # Chunking
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=300)
     splits = text_splitter.split_documents(docs)
 
-    # Load embeddings
+    # Embeddings
     embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
 
+    # Safe temporary directory for Chroma
+    chroma_dir = os.path.join(tempfile.gettempdir(), "chroma_db_rag")
+    if os.path.exists(chroma_dir):
+        shutil.rmtree(chroma_dir)
+    os.makedirs(chroma_dir, exist_ok=True)
 
-    from langchain_community.vectorstores import Chroma
-    
-    # Create Chroma in-memory vector store
-    db = Chroma.from_documents(splits, embeddings)
-    retriever = db.as_retriever()
-
+    db = Chroma.from_documents(splits, embeddings, persist_directory=chroma_dir)
+    db.persist()
+    return db.as_retriever(search_kwargs={"k": 6})
 
-
-    # Set up the model
+# --- Process Answer ---
+def process_answer(question, full_text, retriever):
     llm = load_model()
 
-    # RAG-style retrieval QA
-    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
-
-    # Smart prompting
-    if "summarize" in question.lower() or "summary" in question.lower() or "tl;dr" in question.lower():
+    # Special handling for summary-type queries
+    if any(x in question.lower() for x in ["summarize", "summary", "tl;dr"]):
         prompt = f"Summarize the following document:\n\n{full_text[:3000]}"
-        summary = llm(prompt)
+        summary = llm(prompt)  # Uses the LLM to generate a summary
         return summary
-    else:
-        return qa_chain.run(question)
+
+    # Use RetrievalQA for general queries
+    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
+    return qa_chain.run(question)  # This is the main answer generation with retrieval
 
 # --- UI Layout ---
 with st.sidebar:
@@ -518,9 +520,12 @@ if uploaded_file:
         st.subheader("💬 Ask a Question")
         user_question = st.text_input("Type your question about the PDF content")
 
+        # Build retriever once per session
+        retriever = build_retriever(full_text)
+
         if user_question:
             with st.spinner("Thinking..."):
-                answer = process_answer(user_question, full_text)
+                answer = process_answer(user_question, full_text, retriever)
                 st.markdown("### 🤖 Answer")
                 st.write(answer)
 
@@ -540,3 +545,4 @@ if uploaded_file:
         st.error("⚠️ No text could be extracted from the PDF. Try another file.")
 else:
     st.info("Upload a PDF to begin.")
+