Zenn
Closed6

QdrantのFastEmbedを試す

Qdrant
embedding
VectorDB
fastembed
kun432kun432

ChatGPTに聞いてみた。

FastEmbedは、テキストのための高速かつ軽量な埋め込み生成を実現するPythonライブラリです。以下がその概要です：

どんなもの？

  • FastEmbedは、NLP埋め込みの80%のユースケースを簡単に扱うために設計された、速度、効率、使いやすさを重視したPythonライブラリです。

そもそも本手法が必要となった既存の課題や理由は？

  • 従来の埋め込み生成では、PyTorchやTensorFlowのモデルを使うが、これらは使いやすさや計算速度の点で不便がありました。FastEmbedはこれらの問題を解決するために開発されました。

技術や手法のキモはどこ？

  • FastEmbedは、限られたモデルを使い、モデルの重みを量子化し、ONNX Runtimeとシームレスに統合しています。これにより、推論時間、リソースの使用量、パフォーマンス（リコール/精度）のバランスを取っています。

どうやって有効だと検証した？具体的にどの程度有効？

  • このライブラリは、PyTorch Transformersより50%速く、Sentence TransformersやOpenAI Ada-002よりも優れたパフォーマンスを提供し、量子化されたモデルとオリジナルのモデルのベクトル間のコサイン類似度が0.92に達しています。

具体的なメリットやユースケースは？

  • FastEmbedはQdrantのベクトルストア機能と統合されており、埋め込みの生成、保存、取得の透明なワークフローを提供します。これによりAPI設計が簡素化され、柔軟性が高まります。

本手法を評価するにあたり、課題や不足している情報は？

  • ライブラリの使いやすさやパフォーマンスに関するユーザーの評価や、長期にわたる性能維持の実証データが不足している可能性があります。また、異なるデータセットや環境での適用性についての情報も限られています。
kun432kun432

https://qdrant.github.io/fastembed/Getting Started/

Colaboratoryでやってみる。

FastEmbedのインストール。

!pip install fastembed --upgrade --quiet

from typing import List
import numpy as np
from fastembed.embedding import DefaultEmbedding

# Example list of documents
documents: List[str] = [
    "ハローワールド",
    "これはサンプルのドキュメントです。",
    "fastembed は Qdrant がサポート・保守を行っています。",
]

embedding_model = DefaultEmbedding()
embeddings: List[np.ndarray] = list(embedding_model.embed(documents))

for i in embeddings:
    print(i.shape)
    print(i[:5])

初回はモデルが自動的にダウンロードされる。そして、embeddingsが生成されている。

100%|██████████| 76.7M/76.7M [00:01<00:00, 71.1MiB/s]
(384,)
[-0.03255257 -0.03828212  0.04274148 -0.01026607  0.04006985]
(384,)
[-5.1185773e-03  2.1692269e-02  1.0441988e-02 -6.5478118e-05
 -1.3540394e-02]
(384,)
[-0.08520549  0.0472231  -0.0190189   0.05070178  0.02181093]

デフォルトだと"BAAI/bge-small-en-v1.5"が使用される。

embedding_model.model_name

BAAI/bge-small-en-v1.5

他のモデルを使う場合はFlagEmbeddingを使う。"intfloat/multilingual-e5-large"の場合。

from typing import List
import numpy as np
from fastembed.embedding import FlagEmbedding as Embedding

documents: List[str] = [
    "ハローワールド",
    "これはサンプルのドキュメントです。",
    "fastembed は Qdrant がサポート・保守を行っています。",
]

embedding_model = Embedding(model_name="intfloat/multilingual-e5-large", max_length=512)
embeddings: List[np.ndarray] = list(embedding_model.embed(documents))

for i in embeddings:
    print(i.shape)
    print(i[:5])

ダウンロードが行われる。ちょっとこのメッセージの意味がわからないのだけど、githubレポジトリのissue（こことかここ）を見ていると、どうも"fast-*"は量子化に加えて一定の高速化以上を確認できたモデル、というようなイメージに思える。つまり、multilingual-e5-largeは単に量子化しただけの可能性がある。

Was not able to download fast-multilingual-e5-large.tar.gz, trying intfloat-multilingual-e5-large.tar.gz
100%|██████████| 1.30G/1.30G [00:30<00:00, 43.0MiB/s]

ちなみに、mode初期化時にcache_dirを指定しておけばダウンロード先を指定できる。ただしディレクトリは作成してくれないので、予め用意しておく必要がある。

import os

cache_dir = "./models2"
os.makedirs(cache_dir, exist_ok=True)
embedding_model = Embedding(model_name="intfloat/multilingual-e5-large", max_length=512, cache_dir=cache_dir)

embeddingsが生成されている。

(1024,)
[ 0.03247941  0.03794141 -0.02999929 -0.0523814   0.01580315]
(1024,)
[ 0.02261261 -0.00554694 -0.03062819 -0.05118655  0.02342548]
(1024,)
[ 0.03745371  0.00901454  0.00649533 -0.05110265  0.02917754]

"query" や "passage" のプリフィクスにも対応している。

%%time

from typing import List
import numpy as np
from fastembed.embedding import FlagEmbedding as Embedding

documents: List[str] = [
    "query: 今日はいいお天気ですね",
    "passage: 今日はいいお天気ですね",
    "今日はいいお天気ですね",
]

embedding_model = Embedding(model_name="intfloat/multilingual-e5-large", max_length=512)
embeddings: List[np.ndarray] = list(embedding_model.embed(documents))

for i in embeddings:
    print(i.shape)
    print(i[:5])

当然ながら全部embeddingsが違うのがわかる。

(1024,)
[ 0.04578432 -0.02660514 -0.01958469 -0.04749198  0.04541172]
(1024,)
[ 0.04025473 -0.00037204 -0.01240254 -0.04469366  0.03982825]
(1024,)
[ 0.0341006   0.00427205 -0.01494944 -0.04360745  0.0257986 ]

確かにちょっと遅いかも。

CPU times: user 2.22 s, sys: 1.67 s, total: 3.9 s
Wall time: 3.69 s

一応比較してみる。コードは以下を参考にした。

https://qiita.com/Qiitaman/items/fa393d93ce8e61a857b1

import numpy as np

def cos_sim(v1, v2):
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))

print(cos_sim(embeddings[0],embeddings[1]))
print(cos_sim(embeddings[0],embeddings[2]))
print(cos_sim(embeddings[1],embeddings[2]))

"query: " とプレフィクスなしが一番近いという結果になったけど、あくまでも例なので。

0.92616284
0.9311075
0.9144577

Retrieval向けにquery_embedとpassage_embedというメソッドがある模様。

https://qdrant.github.io/fastembed/examples/Retrieval_with_FastEmbed/

まずドキュメントのembeddingsを生成。個々で作ったものを使ってみた。

https://zenn.dev/kun432/scraps/a3cc525a092f20

from typing import List
import numpy as np
from fastembed.embedding import FlagEmbedding as Embedding

# ドキュメントのリスト
documents: List[str] = [
    "イクイノックスはキタサンブラック産駒である。",
    "イクイノックスは2022年に初のGI制覇を果たした。",
    "イクイノックスは2023年に秋春グランプリ制覇を達成した。",
    "馬名イクイノックスの意味は「昼と夜の長さがほぼ等しくなる時」である。",
    "イクイノックスは2022年度のJRA賞年度代表馬である。",
    "イクイノックスは2022年度の最優秀3歳牡馬である。",
    "イクイノックスの主な勝ち鞍には2022年の天皇賞（秋）が含まれる。",
    "イクイノックスは2023年の天皇賞（秋）でも勝利を収めた。",
    "イクイノックスは2022年の有馬記念に勝利した。",
    "イクイノックスは2023年のドバイシーマクラシックに勝利した。",
    "イクイノックスは2023年の宝塚記念に勝利した。",
    "イクイノックスは2023年のジャパンカップに勝利した。"
]

# モデルはintfloat/multilingual-e5-largeを使う
embedding_model = Embedding(model_name="intfloat/multilingual-e5-large", max_length=512)

# ドキュメントはpassage_embedメソッドを使ってembeddingsを生成する
embeddings: List[np.ndarray] = list(
    embedding_model.passage_embed(documents)
)

print(embeddings[0].shape, len(embeddings))

クエリのembeddingsを生成して、ドキュメントを検索。query_embedとembedの両方で比較している。

query = "勝ったレース名を教えて"

# query_embedを使ってクエリのembeddingsを生成
query_embedding = list(embedding_model.query_embed(query))[0]

# 通常のembedを使ってクエリのembeddingsを生成
plain_query_embedding = list(embedding_model.embed(query))[0]


def print_top_k(query_embedding, embeddings, documents, k=10):
    # ドキュメントの最大を超えないようにする。
    k = min(k, len(documents))
    # numpyでクエリとドキュメントのコサイン類似度を計算
    scores = np.dot(embeddings, query_embedding)
    # スコアを降順でソート
    sorted_pairs = sorted(zip(scores, range(len(scores))), reverse=True)
    # 上位top-kを表示
    for i in range(k):
        score, doc_index = sorted_pairs[i]
        print(f"Rank {i+1}: {score:.10f} - {documents[doc_index]}")

# query_embedの場合の検索結果
print_top_k(query_embedding, embeddings, documents)

print()

# embedの場合の検索結果
print_top_k(plain_query_embedding, embeddings, documents)

結果

Rank 1: 0.8486520052 - イクイノックスは2022年の有馬記念に勝利した。
Rank 2: 0.8480774164 - イクイノックスは2023年に秋春グランプリ制覇を達成した。
Rank 3: 0.8429265618 - イクイノックスは2023年の宝塚記念に勝利した。
Rank 4: 0.8400874138 - イクイノックスは2022年に初のGI制覇を果たした。
Rank 5: 0.8390868902 - イクイノックスはキタサンブラック産駒である。
Rank 6: 0.8373890519 - イクイノックスの主な勝ち鞍には2022年の天皇賞（秋）が含まれる。
Rank 7: 0.8370204568 - イクイノックスは2023年のドバイシーマクラシックに勝利した。
Rank 8: 0.8369027972 - イクイノックスは2022年度の最優秀3歳牡馬である。
Rank 9: 0.8368648291 - イクイノックスは2023年の天皇賞（秋）でも勝利を収めた。
Rank 10: 0.8350150585 - 馬名イクイノックスの意味は「昼と夜の長さがほぼ等しくなる時」である。

Rank 1: 0.8323908448 - イクイノックスは2023年に秋春グランプリ制覇を達成した。
Rank 2: 0.8306561708 - イクイノックスは2022年の有馬記念に勝利した。
Rank 3: 0.8262840509 - イクイノックスはキタサンブラック産駒である。
Rank 4: 0.8260025382 - イクイノックスは2023年の宝塚記念に勝利した。
Rank 5: 0.8257783651 - イクイノックスは2022年に初のGI制覇を果たした。
Rank 6: 0.8225906491 - 馬名イクイノックスの意味は「昼と夜の長さがほぼ等しくなる時」である。
Rank 7: 0.8222481012 - イクイノックスは2022年度の最優秀3歳牡馬である。
Rank 8: 0.8193450570 - イクイノックスは2023年の天皇賞（秋）でも勝利を収めた。
Rank 9: 0.8177816868 - イクイノックスの主な勝ち鞍には2022年の天皇賞（秋）が含まれる。
Rank 10: 0.8168077469 - イクイノックスは2023年のドバイシーマクラシックに勝利した。

なんとなくquery_embedのほうが良いかなというところだけど、ここはクエリとドキュメントによって変わると思う。

kun432kun432

LlamaIndexでも使える

https://docs.llamaindex.ai/en/stable/examples/embeddings/fastembed.html

from llama_index.embeddings import FastEmbedEmbedding
import os

cache_dir = "./models"
os.makedirs(cache_dir, exist_ok=True)

embed_model = FastEmbedEmbedding(
    model_name="intfloat/multilingual-e5-large",
    doc_embed_type="passage",
    max_length=512,
    cache_dir=cache_dir
)

embeddings = embed_model.get_text_embedding("こんにちは。")
print(len(embeddings))
print(embeddings[:5])

embeddings = embed_model.get_query_embedding("こんにちは。")
print(len(embeddings))
print(embeddings[:5])

Was not able to download fast-multilingual-e5-large.tar.gz, trying intfloat-multilingual-e5-large.tar.gz
100%|██████████| 1.30G/1.30G [00:22<00:00, 57.0MiB/s]
1024
[0.043494462966918945, 0.0250167865306139, -0.010338474065065384, -0.04223291948437691, 0.04607989639043808]
1024
[0.035603806376457214, -0.02402173914015293, -0.008755884133279324, -0.04367868974804878, 0.03649580851197243]

モデル初期化時にdoc_embed_typeで"passage"を指定すると、get_text_embeddingメソッド内でFastEmbedのpassage_embedメソッドが使用される。何も指定しなければ、通常のembedメソッドになる。

クエリのembeddingsはget_query_embeddingを使えば常に"query"プレフィクスがつく様子なので、ローレベルでembeddingsを生成する場合は意識して使い分ける必要があると思うけど、ざっとコード見た限り、この辺の違いはよしなにやってくれる模様（どのembeddingsモジュールにもget_text_embeddingとget_query_embeddingは実装されているので、クエリエンジンやらインデックスとかでもその辺は見ているんだろうという勝手な推測）

kun432kun432

多少遅くなるかもしれないけども、量子化でリソースも抑えれそうなので良いなーと言う感はある。

LlamaIndexだとHuggingFaceEmbeddingを使ってもシンプルに書けるんだけど、VRAM5GBぐらいは使ってたような気がする。

from llama_index.embeddings import HuggingFaceEmbedding

embed_model = HuggingFaceEmbedding(
    model_name="./multilingual-e5-large",
    max_length=512,
)

HuggingFaceEmbeddingでCPUに限定させるならこういう感じ

embed_model = HuggingFaceEmbedding(
    model_name="./multilingual-e5-large",
    device="cpu",
    max_length=512,
)

我が家のサーバではCPU動作させてもクエリ程度なら全然遅いとは感じないけど、インデックス作成するときは多少の遅さを感じるし、CPUリソースも結構食ってた。

このスクラップは2024/01/25にクローズされました