Chroma 검색 — .query와 .get 활용하기

Chroma 검색 — .query와 .get 활용하기

Chroma의 검색은 크게 두 가지예요. .query 는 밀집 임베딩 기반 최근접 유사도 검색이고, .get 은 유사도 순위 없이 레코드를 그대로 가져옵니다. 상황에 따라 골라 쓰면 됩니다.

출처: https://docs.trychroma.com/docs/querying-collections/query-and-get

Query로 유사도 검색

.query에 질의 텍스트를 넘기면 컬렉션의 임베딩 함수로 벡터를 만든 뒤 유사도 검색을 수행해요.

collection.query(
    query_texts=["thus spake zarathustra", "the oracle speaks"]
)

임베딩 함수 없이 만들어진 컬렉션이라면 query_embeddings를 직접 넘겨야 합니다. 이때 차원이 컬렉션의 임베딩과 일치해야 해요.

collection.query(
    query_embeddings=[[11.1, 12.1, 13.1], [1.1, 2.3, 3.2]],
    n_results=100
)

n_results로 반환 개수를 조절하고, ids로 특정 ID만 검색 대상에 넣을 수 있어요. 메타데이터 필터링은 where, 문서 내용 필터링은 where_document로 합니다.

collection.query(
    query_embeddings=[[11.1, 12.1, 13.1], [1.1, 2.3, 3.2]],
    n_results=100,
    where={"page": 10}, # query records with metadata field 'page' equal to 10
    where_document={"$contains": "search string"} # query records with the search string in the records' document
)

Get으로 레코드 조회

유사도 순위 없이 ID나 필터로 그냥 가져오려면 .get을 씁니다. 페이징도 지원해요.

result = collection.get(include=["documents", "metadatas"])
for id, document, metadata in zip(result["ids"], result["documents"], result["metadatas"]):
    print(id, document, metadata)

반환되는 데이터 고르기

include 옵션으로 무엇을 돌려받을지 정할 수 있어요. 문서·메타데이터·임베딩 중 필요한 것만 선택합니다.

collection.query(
    query_texts=["my query"],
    include=["documents", "metadatas", "embeddings"],
)

.query는 배치 API라 결과가 질의 입력마다 묶여서 옵니다. 보통 각 질의의 묶음을 돌고, 그 안에서 결과를 다시 돌면서 처리하는 패턴을 씁니다.

result = collection.query(query_texts=["first query", "second query"])
for ids, documents, metadatas in zip(result["ids"], result["documents"], result["metadatas"]):
    for id, document, metadata in zip(ids, documents, metadatas):
        print(id, document, metadata)

더 알아보기