Update generate_embeddings.py

generate_embeddings.py

@@ -1,18 +1,12 @@
 import os
 import time
 from langchain_community.vectorstores import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.document_loaders import PyPDFLoader
 
-# ==============================
-# Folder path (IMPORTANT)
-# ==============================
 folder_path = "documents"
 
-# ==============================
-# Load PDFs
-# ==============================
 def load_pdfs(folder_path):
     docs = []
     for filename in os.listdir(folder_path):
@@ -29,9 +23,7 @@ def load_pdfs(folder_path):
     print(f"Total documents loaded: {len(docs)}")
     return docs
 
-# ==============================
-# Split documents
-# ==============================
+
 def split_documents(docs):
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=500,
@@ -41,9 +33,7 @@ def split_documents(docs):
     print(f"Total chunks: {len(chunks)}")
     return chunks
 
-# ==============================
-# Create vector DB
-# ==============================
+
 def create_vectorstore(docs):
     embedding = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2"
@@ -63,9 +53,7 @@ def create_vectorstore(docs):
     print(f"Done in {time.time() - start} sec")
     return vector_db
 
-# ==============================
-# Main
-# ==============================
+
 def main():
     docs = load_pdfs(folder_path)
     chunks = split_documents(docs)
@@ -73,5 +61,6 @@ def main():
 
     print(f"Stored documents: {vector_db._collection.count()}")
 
+
 if __name__ == "__main__":
     main()