Delete createVectorDB.py
Browse files- createVectorDB.py +0 -49
createVectorDB.py
DELETED
@@ -1,49 +0,0 @@
|
|
1 |
-
import chromadb
|
2 |
-
from chromadb.utils import embedding_functions
|
3 |
-
from typing import Optional, Dict
|
4 |
-
|
5 |
-
def createVectorDB(
|
6 |
-
collection_name: Optional[str],
|
7 |
-
chroma_data_path: Optional[str] = None,
|
8 |
-
embed_model: Optional[str] = "all-MiniLM-L6-v2",
|
9 |
-
metadata: Optional[Dict[str, str]] = None
|
10 |
-
) -> chromadb.Collection:
|
11 |
-
"""Creates the vector database to store embeddings.
|
12 |
-
|
13 |
-
Args:
|
14 |
-
collection_name (str): The name of the collection.
|
15 |
-
chroma_data_path (Optional[str]): Path for chroma embeddings.
|
16 |
-
embed_model (Optional[str]): Model name for embeddings.
|
17 |
-
metadata (Optional[Dict[str, str]]): Metadata for the collection.
|
18 |
-
|
19 |
-
Returns:
|
20 |
-
chromadb.Collection: The created collection object.
|
21 |
-
"""
|
22 |
-
if chroma_data_path is None:
|
23 |
-
chroma_data_path = r"CHROMA_EMBEDDINGS_PATH" # Default path if not provided
|
24 |
-
|
25 |
-
client = chromadb.PersistentClient(path=chroma_data_path)
|
26 |
-
|
27 |
-
embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
|
28 |
-
model_name=embed_model
|
29 |
-
)
|
30 |
-
|
31 |
-
# Use provided metadata or default to empty dictionary
|
32 |
-
if metadata is None:
|
33 |
-
metadata = {"hnsw:space": "cosine"}
|
34 |
-
|
35 |
-
collection = client.create_collection(
|
36 |
-
collection_name=collection_name,
|
37 |
-
embedding_function=embedding_func,
|
38 |
-
metadata=metadata,
|
39 |
-
)
|
40 |
-
|
41 |
-
return collection
|
42 |
-
|
43 |
-
#unsure how to create unittest
|
44 |
-
|
45 |
-
#collection = createVectorDB(
|
46 |
-
#COLLECTION_NAME="123456789",
|
47 |
-
#C#HROMA_DATA_PATH=r"C:\Users\navan\Downloads\BioModelsRAG\CHROMA_EMBEDDINGS_PATH",
|
48 |
-
#EMBED_MODEL="all-MiniLM-L6-v2",
|
49 |
-
#metadata={"hnsw:space": "cosine"}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|