You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
51 lines
1.5 KiB
51 lines
1.5 KiB
import chromadb as db |
|
from chromadb.utils import embedding_functions |
|
from chromadb.config import Settings |
|
from chromadb.api.client import Client |
|
from chromadb.api.models.Collection import Collection |
|
|
|
|
|
|
|
class ChromaDB: |
|
""" |
|
A class representing a Chroma database. |
|
""" |
|
|
|
def __init__(self, host: str = "192.168.1.10", port: int = 8001) -> None: |
|
""" |
|
Initializes a ChromaDB object running on specified port. |
|
|
|
Args: |
|
host (str, optional): The host address of the Chroma database. Defaults to "192.168.1.10". |
|
port (int, optional): The port number of the Chroma database. Defaults to 8001. |
|
""" |
|
self.client: Client = db.HttpClient( |
|
settings=Settings(anonymized_telemetry=False), |
|
host=host, |
|
port=port, |
|
) |
|
huggingface_ef = embedding_functions.HuggingFaceEmbeddingFunction( |
|
api_key='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', |
|
model_name="KBLab/sentence-bert-swedish-cased" |
|
) |
|
self.embedding_function: embedding_functions = huggingface_ef |
|
|
|
def print_collections(self): |
|
""" |
|
Prints all collections in the database. |
|
""" |
|
collections = self.client.list_collections() |
|
for collection in collections: |
|
print(collection.name) |
|
|
|
# Initialize the ChromaDB object |
|
chroma = ChromaDB() |
|
|
|
if __name__ == '__main__': |
|
|
|
chroma.client.delete_collection('mala_persons') |
|
col = chroma.client.get_or_create_collection('mala_persons') |
|
|
|
print(col.count()) |
|
|
|
|
|
|