ChromaDB　セマンティック検索

前回、

セマンティック検索ができていないことにきづいて、
公式docがセマンティック検索ができるといっているモデルがあったので
試してみました

Pretrained Models — Sentence-Transformers documentation

import chromadb
from chromadb.utils import embedding_functions

# 選択したモデルを使用して埋め込み関数を作成
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="multi-qa-MiniLM-L6-cos-v1")

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

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

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

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

print(results)

multi-qa-MiniLM-L6-cos-v1というモデルを使ってみます。
ドキュメントを一つ増やしていてsausage、
クエリはhotdogです。
文字列だけみるとbulldogの方が近そうですが
意味的にはsausageのほうがhotdogに近いはず。
結果

{'ids': [['id3', 'id2', 'id1']], 'distances': [[1.1697618961334229, 1.3224105834960938, 1.5199581384658813]], 'metadatas': [[{'source': 'food'}, {'source': 'dog'}, {'source': 'cat'}]], 'embeddings': None, 'documents': [['sausage', 'bulldog', 'bengal']], 'uris': None, 'data': None}

おお。
ちゃんとsausageの方が距離が近くでています。
良い感じです
次はクエリをcatにしてみます。

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

結果

{'ids': [['id2', 'id1', 'id3']], 'distances': [[1.3627203702926636, 1.4854635000228882, 1.4922412633895874]], 'metadatas': [[{'source': 'dog'}, {'source': 'cat'}, {'source': 'food'}]], 'embeddings': None, 'documents': [['bulldog', 'bengal', 'sausage']], 'uris': None, 'data': None}

bulldogのほうがbengalより近くなってしまいました…。
続く。