Update app.py

app.py

@@ -18,57 +18,56 @@ from langchain_community.document_loaders.directory import DirectoryLoader
 from HTML_templates import css, bot_template, user_template
 
 
+data_path = "data"
 
-
-def retriever_from_chroma(docs, search_type, k):
+def create_retriever_from_chroma(data_path, vectorstore_path="docs/chroma/", search_type='mmr', k=7, chunk_size=250, chunk_overlap=20):
     model_name = "sentence-transformers/all-mpnet-base-v2"
     model_kwargs = {'device': 'cpu'}
     encode_kwargs = {'normalize_embeddings': True}
+
+    # Initialize embeddings
     embeddings = HuggingFaceEmbeddings(
         model_name=model_name,
         model_kwargs=model_kwargs,
         encode_kwargs=encode_kwargs
     )
-    vectorstore_path = "docs/chroma/"
-    if not os.path.exists(vectorstore_path):
-        os.makedirs(vectorstore_path)
-    vectorstore = Chroma.from_documents(
-        documents=docs, embedding=embeddings, persist_directory="docs/chroma/")
+
+    # Check if vectorstore exists
+    if os.path.exists(vectorstore_path) and os.listdir(vectorstore_path):
+        # Load the existing vectorstore
+        vectorstore = Chroma(persist_directory=vectorstore_path, embedding=embeddings)
+    else:
+        # Load documents from the specified data path
+        documents = []
+        for filename in os.listdir(data_path):
+            if filename.endswith('.txt'):
+                file_path = os.path.join(data_path, filename)
+                loaded_docs = TextLoader(file_path).load()
+                documents.extend(loaded_docs)
+
+        # Split documents into chunks
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
+        split_docs = text_splitter.split_documents(documents)
+
+        # Ensure the directory for storing vectorstore exists
+        if not os.path.exists(vectorstore_path):
+            os.makedirs(vectorstore_path)
+
+        # Create the vectorstore
+        vectorstore = Chroma.from_documents(
+            documents=split_docs, embedding=embeddings, persist_directory=vectorstore_path
+        )
+
+    # Create and return the retriever
     retriever = vectorstore.as_retriever(search_type=search_type, search_kwargs={"k": k})
     return retriever
-        
-
-
-
-data_path = "data"
-
-
-    
-documents = []
-    
-for filename in os.listdir(data_path):
-
-    if filename.endswith('.txt'):
-
-        file_path = os.path.join(data_path, filename)
-
-        documents = TextLoader(file_path).load()
-
-        documents.extend(documents)
-    
-
-
-
-docs = split_docs(documents, 250, 20)
-
-retriever = retriever_from_chroma(docs,'mmr',7)
 
 
 
 
 
 
-def main(retriever):
+def main():
 
     st.set_page_config(page_title="Chat with multiple PDFs",
                        page_icon=":books:")
@@ -90,7 +89,8 @@ def main(retriever):
         handle_userinput(user_question,vectorstore)
 
 
-def handle_userinput(user_question,retriever):
+def handle_userinput(user_question):
+    retriever = create_retriever_from_chroma(data_path, vectorstore_path="docs/chroma/", search_type='mmr', k=7, chunk_size=250, chunk_overlap=20)
     docs = retriever.invoke(question)
     
     doc_txt = [doc.page_content for doc in docs]
@@ -135,4 +135,4 @@ def create_conversational_rag_chain(retriever):
  
 
 if __name__ == "__main__":
-    main(vectorstore)
+    main()