gamingflexer commited on
Commit
c498106
1 Parent(s): d012361

Chroma Fucntions added

Browse files
Files changed (1) hide show
  1. src/db/vector_fucntions.py +24 -0
src/db/vector_fucntions.py CHANGED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from config import *
2
+ import chromadb
3
+ from chromadb.utils import embedding_functions
4
+ from config import OPENAI_API_KEY
5
+ from utils import generate_uuid
6
+
7
+ emmbedding_model = "text-embedding-3-large"
8
+ openai_ef = embedding_functions.OpenAIEmbeddingFunction(model_name=emmbedding_model,api_key=OPENAI_API_KEY)
9
+ if deploy:
10
+ chroma_client = chromadb.PersistentClient(path="./data/emeddeings")
11
+ else:
12
+ chroma_client = chromadb.PersistentClient(path="/home/ubuntu/research/data/emeddeings")
13
+
14
+ collection_doc = chroma_client.get_or_create_collection(name="2024_main_document_lvl")
15
+ collection_para = chroma_client.get_or_create_collection(name="2024_main_paragraph_lvl")
16
+
17
+ def add_document_chroma_collection(collection_object, document_list, embedding_list, metadata):
18
+
19
+ metadata_list = [metadata for i in range(len(document_list))]
20
+ ids_gen = [generate_uuid() for i in range(len(document_list))]
21
+ collection_object.add(embeddings = embedding_list,documents = document_list,metadatas = metadata_list,ids = ids_gen)
22
+ if collection_object:
23
+ return True
24
+