ChromaDB セマンティック検索③

multilingual-e5-largeモデルを使ってみる

import chromadb
from transformers import AutoTokenizer, AutoModel

model = AutoModel.from_pretrained('intfloat/multilingual-e5-large')
tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-large')

class EmbeddingFunction:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def __call__(self, input):
        embeddings = []
        for text in input:
            inputs = self.tokenizer(text, return_tensors="pt")
            embedding = self.model(**inputs).last_hidden_state[:, 0, :].detach().numpy().squeeze()
            embeddings.append(embedding.tolist())  # ndarrayをリストに変換
        return embeddings

# 埋め込み関数のインスタンスを作成
embedding_function = EmbeddingFunction(model, tokenizer)

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

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

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

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

print(results)

実行結果

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

catに対してsausageが近いという実行結果になってしまった。
且つ、数字がとても大きい。
今までと違う。。。
やり方は合ってるのかな。。。