見出し画像

ChromaDBを試してみて気付きました…。

ベクトルデータベースChromaDBの実験をしてみました。

import chromadb

chroma_client = chromadb.Client()
collection = chroma_client.create_collection(name="my_collection")
collection.add(
    documents=["american shorthair", "bulldog"],
    metadatas=[{"source": "cat"}, {"source": "dog"}],
    ids=["id1", "id2"]
)
results = collection.query(
    query_texts=["cat"],
    n_results=2
)

print(results)

軽くソースコード解説すると
collection.addはCRUD操作のCreateにあたる。
documents=が登録したいデータの内容で
metadatasがそのメタデータ。
idsはたぶんPK
collection.queryはCRUD操作のReadにあたる。
query_textsで検索したい語句をいれてユークリッド距離で
DBの中から近いデータを探し出す。
(ユークリッド距離の計算にmetadatasは関係ないっぽいです。)
結果

{'ids': [['id2', 'id1']], 'distances': [[1.048373818397522, 1.4675920009613037]], 'metadatas': [[{'source': 'dog'}, {'source': 'cat'}]], 'embeddings': None, 'documents': [['bulldog', 'american shorthair']], 'uris': None, 'data': None}

id2はbulldogです。
id1はamerican shorthairつまり猫です。
クエリはcat、猫です。
id1の方がユークリッド距離が近いと思ったら、違いました。
これは単語の意味の類似性ではなく単語の綴りの類似性だけを
比較しているということだと思います。
試しに…

    query_texts=["american shorthair"],
    n_results=2
)

こうすると、

{'ids': [['id1', 'id2']], 'distances': [[0.0, 1.326899766921997]], 'metadatas': [[{'source': 'cat'}, {'source': 'dog'}]], 'embeddings': None, 'documents': [['american shorthair', 'bulldog']], 'uris': None, 'data': None}

こうなります。
0.0つまり完全一致。

ChromaDB公式によると、all-MiniLM-L6-v2
というembeddingモデルをデフォルトで使っているそうです。
しかしモデルの切り替えもできるっぽいです。

ChromaDBの公式Doc↑

↑ChromaDBで使えるモデル一覧

all-mpnet-base-v2が最高品質らしいので使ってみた

import chromadb
from chromadb.utils import embedding_functions

# 選択したモデルを使用して埋め込み関数を作成
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-mpnet-base-v2")

# Chromaクライアントを作成
chroma_client = chromadb.Client()

# コレクションの作成時に埋め込み関数を指定
collection = chroma_client.create_collection(name="my_collection", embedding_function=sentence_transformer_ef)

# ドキュメントの追加
collection.add(
    documents=["american shorthair", "bulldog"],
    metadatas=[{"source": "cat"}, {"source": "dog"}],
    ids=["id1", "id2"]
)

# クエリの実行
results = collection.query(
    query_texts=["cat"],
    n_results=2
)

print(results)

結果

{'ids': [['id2', 'id1']], 'distances': [[1.1602305173873901, 1.2963690757751465]], 'metadatas': [[{'source': 'dog'}, {'source': 'cat'}]], 'embeddings': None, 'documents': [['bulldog', 'american shorthair']], 'uris': None, 'data': None}

んん。
クエリcatなのにまだブルドッグが近いぞ。
データを変えてみよう。

# ドキュメントの追加
collection.add(
    documents=["bengal", "bulldog"],
    metadatas=[{"source": "cat"}, {"source": "dog"}],
    ids=["id1", "id2"]
)

ベンガル猫。
結果↓

{'ids': [['id1', 'id2']], 'distances': [[1.1358128786087036, 1.160230040550232]], 'metadatas': [[{'source': 'cat'}, {'source': 'dog'}]], 'embeddings': None, 'documents': [['bengal', 'bulldog']], 'uris': None, 'data': None}

やっとcatのクエリに対してbengalが近くなったけど
単純に文字数が短くなって近くなっているように思える。

色々調べていった結果、
僕がやりたいことはたぶんセマンティック検索で
やり方を間違えていたっぽいことに気づきました。
これは、ただ単語とデータのユークリッド距離を計算しているだけっぽいです。
次回はChromaDBでセマンティック検索を実施してみます。

この記事が気に入ったらサポートをしてみませんか?