Spaces:
Build error
Build error
File size: 2,711 Bytes
60e3a80 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
# This tests a very minimal of test_add in test_add.py as a example based test
# instead of a property based test. We can use the delta to get the property
# test working and then enable
import random
import time
from chromadb.api import ClientAPI
from chromadb.test.conftest import (
COMPACTION_SLEEP,
reset,
skip_if_not_cluster,
)
from chromadb.test.property import invariants
import numpy as np
@skip_if_not_cluster()
def test_add(
client: ClientAPI,
) -> None:
seed = time.time()
random.seed(seed)
print("Generating data with seed ", seed)
reset(client)
collection = client.create_collection(
name="test",
metadata={"hnsw:construction_ef": 128, "hnsw:search_ef": 128, "hnsw:M": 128},
)
# Add 1000 records, where each embedding has 3 dimensions randomly generated
# between 0 and 1
ids = []
embeddings = []
for i in range(1000):
ids.append(str(i))
embeddings.append(np.random.rand(1, 3)[0])
collection.add(
ids=[str(i)],
embeddings=[embeddings[-1]],
)
random_query = np.random.rand(1, 3)[0]
print("Generated data with seed ", seed)
invariants.ann_accuracy(
collection,
{
"ids": ids,
"embeddings": embeddings,
"metadatas": None,
"documents": None,
},
10,
query_embeddings=[random_query],
)
@skip_if_not_cluster()
def test_add_include_all_with_compaction_delay(client: ClientAPI) -> None:
seed = time.time()
random.seed(seed)
print("Generating data with seed ", seed)
reset(client)
collection = client.create_collection(
name="test_add_include_all_with_compaction_delay",
metadata={"hnsw:construction_ef": 128, "hnsw:search_ef": 128, "hnsw:M": 128},
)
ids = []
embeddings = []
documents = []
for i in range(1000):
ids.append(str(i))
embeddings.append(np.random.rand(1, 3)[0])
documents.append(f"document_{i}")
collection.add(
ids=[str(i)],
embeddings=[embeddings[-1]],
documents=[documents[-1]],
)
time.sleep(COMPACTION_SLEEP) # Wait for the documents to be compacted
random_query_1 = np.random.rand(1, 3)[0]
random_query_2 = np.random.rand(1, 3)[0]
print("Generated data with seed ", seed)
# Query the collection with a random query
invariants.ann_accuracy(
collection,
{
"ids": ids,
"embeddings": embeddings,
"metadatas": None,
"documents": documents,
},
10,
query_embeddings=[random_query_1, random_query_2],
)
|