Spaces:
Sleeping
Sleeping
File size: 1,171 Bytes
9b45841 | 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 | import chromadb
import uuid
client = chromadb.PersistentClient(
path="./chroma_db"
)
collection = client.get_or_create_collection(
name="documents"
)
def add_chunks(chunks, embeddings, filename):
ids = [str(uuid.uuid4()) for _ in chunks]
metadatas = [
{
"source": filename,
"chunk_id": i
}
for i in range(len(chunks))
]
collection.add(
ids=ids,
documents=chunks,
embeddings=embeddings.tolist(),
metadatas=metadatas
)
def search(query_embedding, top_k=5):
results = collection.query(
query_embeddings=[query_embedding.tolist()],
n_results=top_k,
include=["documents", "metadatas", "distances"]
)
return {
"documents": results["documents"][0],
"metadatas": results["metadatas"][0],
"distances": results["distances"][0]
}
def reset_collection():
global collection
try:
client.delete_collection("documents")
except:
pass
collection = client.get_or_create_collection(
name="documents"
) |