Update initialize.py

initialize.py

@@ -4,6 +4,7 @@ from langchain_community.vectorstores import Chroma
 import os
 import gen_splits
 from langchain_ollama import OllamaEmbeddings
+import ollama
 
 
 
@@ -21,10 +22,20 @@ embedding_model = OpenAIEmbeddings(api_key=OPENAI_API_KEY, model="text-embedding
 persist_directory = './chroma/'
 
 def initialize():
+    # splits = gen_splits.gen_splits()
+    # vectordb = Chroma.from_documents(documents=splits, persist_directory=persist_directory, embedding=embedding_model)
+    # vectordb.persist()
+
     splits = gen_splits.gen_splits()
-    vectordb = Chroma.from_documents(documents=splits, persist_directory=persist_directory, embedding=embedding_model)
-    vectordb.persist()
+    collection = client.create_collection(name="docs")
+    # store each document in a vector embedding database
+    for i, d in enumerate(splits):
+        response = ollama.embeddings(model="mxbai-embed-large", prompt=d)
+        embedding = response["embedding"]
+        collection.add(ids=[str(i)],embeddings=[embedding], documents=[d])
+    
     return vectordb
+    
 
 if __name__ == "__main__":
     vectordb = initialize()