Update app.py

app.py

@@ -10,7 +10,9 @@ from langchain.llms import CTransformers
 
 # === Configuration ===
 pdfs_directory = 'pdfs'
+vectorstores_directory = 'vectorstores'
 os.makedirs(pdfs_directory, exist_ok=True)
+os.makedirs(vectorstores_directory, exist_ok=True)
 
 PREDEFINED_BOOKS = [f for f in os.listdir(pdfs_directory) if f.endswith(".pdf")]
 
@@ -23,10 +25,9 @@ Context: {context}
 Answer:
 """
 
-# === Load Embeddings (CPU Friendly) ===
+# === Embeddings and LLM (CPU-friendly) ===
 embedding_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
 
-# === LLM (Quantized, CPU Efficient) ===
 llm = CTransformers(
     model='TheBloke/Mistral-7B-Instruct-v0.1-GGUF',
     model_file='mistral-7b-instruct-v0.1.Q4_K_M.gguf',
@@ -53,8 +54,24 @@ def split_text(documents):
     )
     return splitter.split_documents(documents)
 
-def create_vector_store(docs):
-    return FAISS.from_documents(docs, embedding_model)
+def get_vectorstore_path(book_filename):
+    base_name = os.path.splitext(book_filename)[0]
+    return os.path.join(vectorstores_directory, base_name)
+
+def load_or_create_vectorstore(book_filename, documents=None):
+    vs_path = get_vectorstore_path(book_filename)
+    
+    if os.path.exists(os.path.join(vs_path, "index.faiss")):
+        return FAISS.load_local(vs_path, embedding_model, allow_dangerous_deserialization=True)
+
+    if documents is None:
+        raise ValueError("Documents must be provided when creating a new vectorstore.")
+
+    os.makedirs(vs_path, exist_ok=True)
+    chunks = split_text(documents)
+    vector_store = FAISS.from_documents(chunks, embedding_model)
+    vector_store.save_local(vs_path)
+    return vector_store
 
 def retrieve_docs(vector_store, query):
     return vector_store.similarity_search(query)
@@ -66,8 +83,8 @@ def answer_question(question, documents):
     return chain.run({"question": question, "context": context})
 
 # === UI ===
-st.set_page_config(page_title="📄 PDF Q&A (CPU Version)", layout="centered")
-st.title("📚 Chat with PDF - CPU Optimized")
+st.set_page_config(page_title="📄 PDF Q&A (Cached FAISS)", layout="centered")
+st.title("📚 Chat with PDF - Cached Vector Stores")
 
 with st.sidebar:
     st.header("Select or Upload a Book")
@@ -80,17 +97,26 @@ with st.sidebar:
             st.success(f"Uploaded: {filename}")
             selected_book = filename
 
+# === Load or Create Vector Store ===
 if selected_book and selected_book != "Upload new book":
-    st.info(f"📖 You selected: {selected_book}")
     file_path = os.path.join(pdfs_directory, selected_book)
-
-    documents = load_pdf(file_path)
-    chunks = split_text(documents)
-    vector_store = create_vector_store(chunks)
-
-    question = st.chat_input("Ask a question about the book...")
-    if question:
-        st.chat_message("user").write(question)
-        related_docs = retrieve_docs(vector_store, question)
-        answer = answer_question(question, related_docs)
-        st.chat_message("assistant").write(answer)
+    vectorstore_path = get_vectorstore_path(selected_book)
+
+    try:
+        if os.path.exists(os.path.join(vectorstore_path, "index.faiss")):
+            st.info("✅ Using cached vector store.")
+            vector_store = load_or_create_vectorstore(selected_book)
+        else:
+            st.warning("⏳ Creating new vector store (first-time load)...")
+            documents = load_pdf(file_path)
+            vector_store = load_or_create_vectorstore(selected_book, documents)
+            st.success("✅ Vector store created and cached.")
+
+        question = st.chat_input("Ask a question about the book...")
+        if question:
+            st.chat_message("user").write(question)
+            related_docs = retrieve_docs(vector_store, question)
+            answer = answer_question(question, related_docs)
+            st.chat_message("assistant").write(answer)
+    except Exception as e:
+        st.error(f"❌ Error loading PDF or vector store: {e}")