| import os |
| from typing import List |
| from sentence_transformers import SentenceTransformer |
| import chromadb |
| from chromadb.utils import embedding_functions |
|
|
| DEFAULT_MODEL = os.getenv("EMB_MODEL", "sentence-transformers/all-MiniLM-L6-v2") |
|
|
| def get_embedder(): |
| |
| model = SentenceTransformer(DEFAULT_MODEL) |
| ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=DEFAULT_MODEL) |
| return model, ef |
|
|
| def get_chroma(collection_name: str = "devices", persist_dir: str = ".chroma"): |
| _, ef = get_embedder() |
| client = chromadb.PersistentClient(path=persist_dir) |
| try: |
| col = client.get_collection(collection_name, embedding_function=ef) |
| except Exception: |
| col = client.create_collection(collection_name, embedding_function=ef) |
| return client, col |
|
|
|
|
