TAPASxHF2 / test_chroma_db_indexer.py
jskinner215's picture
Upload 17 files
25fc3a2
raw
history blame contribute delete
434 Bytes
from chroma_db_indexer import ChromaDBIndexer
def test_chroma_db_indexer():
chroma_db_indexer = ChromaDBIndexer()
# Test creating index
index = chroma_db_indexer.create_index()
assert index.metric_type == "cosine"
assert index.dim == 768
# Test uploading embeddings
embeddings = [0]*768
chroma_db_indexer.upload_embeddings(embeddings, index)
assert len(index.search(embeddings, 1)) == 1