【ベクトルデータベース】PythonからChromaDBを使用する際に抑えておくと良い内容まとめ

ChromaDBの概要

概要

ChromaDBはPythonやJavascriptなどから使うことのできるオープンソースのベクトルデータベースです。ChromaDBを用いることで単語や文書のベクトル化、作成したベクトルの保存、ベクトルの検索などが可能です。

インストール

ChromaDBはpip install chromadbを実行することでPyPIからインストールすることができます。

ChromaDBの使い方

サンプルコード

import chromadb
import pprint

chroma_client = chromadb.Client()

# switch `create_collection` to `get_or_create_collection` to avoid creating a new collection every time
collection = chroma_client.get_or_create_collection(name="my_collection")

# switch `add` to `upsert` to avoid adding the same documents every time
collection.upsert(
    documents=[
        "This is a document about sports",
        "This is a document about fruits"
    ],
    ids=["id1", "id2"]
)
results = collection.query(
    query_texts=["This is a query document about oranges"], # Chroma will embed this for you
    n_results=2 # how many results to return
)

pprint.pp(results)

・実行結果

{'ids': [['id2', 'id1']],
 'distances': [[0.7922683954238892, 1.5041061639785767]],
 'metadatas': [[None, None]],
 'embeddings': None,
 'documents': [['This is a document about fruits',
                'This is a document about sports']],
 'uris': None,
 'data': None,
 'included': ['metadatas', 'documents', 'distances']}

実行結果より、検索クエリのoranges(オレンジ)が、sports(スポーツ)よりもフルーツ(fruits)に関連していることが確認できます。

抑えておくと良い事項

類似度計算の手法の変更

client = chromadb.Client()
collection = client.create_collection(
        name="collection_name",
        metadata={"hnsw:space": "cosine"} # l2 is the default
    )

上記のように、client.create_collectionのmetadataでベクトルの類似度計算の手法を変えることが可能です。設定するパラメータと数式の対応は下記より確認できます。

ChromaDBのドキュメントより
https://docs.trychroma.com/guides

addの際に何を入力とするか

addやupsertを実行する際に、documentsを用いる場合は文書を入力します。この時、文書に対応するベクトルが生成され、それぞれDBに格納されます。一方で、embeddingsを用いる場合は文書の代わりにベクトルを入力にする必要があります。下記がそれぞれの例です。
・documentsを用いる場合

collection.add(
    documents=["lorem ipsum...", "doc2", "doc3", ...],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],
    ids=["id1", "id2", "id3", ...]
)

・embeddingsを用いる場合

collection.add(
    embeddings=[[1.1, 2.3, 3.2], [4.5, 6.9, 4.4], [1.1, 2.3, 3.2], ...],
    metadatas=[{"chapter": "3", "verse": "16"}, {"chapter": "3", "verse": "5"}, {"chapter": "29", "verse": "11"}, ...],
    ids=["id1", "id2", "id3", ...]
)

参考

・ChromaDB Document
https://docs.trychroma.com/getting-started
https://docs.trychroma.com/guides