File size: 846 Bytes
8d48081 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | 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():
# Use a light SentenceTransformer and pass to Chroma
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
|