見出し画像

ChromaDB セマンティック検索②OpenAIのEmbeddingモデル

前回の続き

catというクエリで
bengalよりもbulldogが近いと出てしまう問題があったので
モデルをOpenAIのEmbeddingモデルに変えてみた

import chromadb
import chromadb.utils.embedding_functions as embedding_functions

openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="APIKEY",
                model_name="text-embedding-3-small"
            )

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

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

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

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

print(results)

実行結果

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

おお!
良い感じ。
ちゃんとbengalの方がbulldogより近いです。

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

ベンガルをアメリカンショートヘアに変えてみる
実行結果

{'ids': [['id1', 'id2', 'id3']], 'distances': [[1.213362216949463, 1.2822411060333252, 1.4614238739013672]], 'metadatas': [[{'source': 'cat'}, {'source': 'dog'}, {'source': 'food'}]], 'embeddings': None, 'documents': [['american shorthair', 'bulldog', 'sausage']], 'uris': None, 'data': None}

おー!
ちゃんとできてる。
モデルを変えてみる。
次はada-002
ドキュメントとクエリは一緒
実行結果

{'ids': [['id1', 'id2', 'id3']], 'distances': [[0.36710622906684875, 0.4427233338356018, 0.45577099919319153]], 'metadatas': [[{'source': 'cat'}, {'source': 'dog'}, {'source': 'food'}]], 'embeddings': None, 'documents': [['american shorthair', 'bulldog', 'sausage']], 'uris': None, 'data': None}

ちゃんとできているけどさっきより数が小さい!
text-embedding-3-smallと計算方法が全然違う。

モデルを変えると全然違う値が返ってくる。
サービスに活かすか生かさないかの検討は
近いか遠いかが妥当であるかどうか。
別のドキュメントやクエリにすると結果は変わるかもしれない。

続く

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