Google Colab で BAAI/bge-m3 を試す

2024年4月1日 17:04

「Google Colab」で「BAAI/bge-m3」を試したので、まとめました。

1. BAAI/bge-m3

「BAAI/bge-m3」は、多機能、多言語、多粒度の特徴を持つ埋め込みモデルです。「JapaneseEmbeddingEval」でも高スコアをマークしています。

・多機能 : 埋め込みモデルの3つの一般的な検索機能 (dense retrieval, multi-vector retrieval, and sparse retrieval) を同時に実行可能
・多言語 : 100 以上の使用言語をサポート
・多粒度 : 短い文から最大8192トークンの長文まで、さまざまな粒度を入力可能

2. Colabでの実行

Colabでの実行手順は、次のとおりです。

(1) パッケージのインストール。

# パッケージのインストール
!pip install -U FlagEmbedding

(2) モデルの準備。
use_fp16=True を指定すると、パフォーマンスが若干低下しますが、計算が高速化されます。

from FlagEmbedding import BGEM3FlagModel

# モデルの準備
model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)

(3) 入力テキストのEmbeddingの生成。

# 入力テキストのEmbeddingの生成
input_texts = ["今日は雨振らなくてよかった"]
input_embeds = model.encode(input_texts)["dense_vecs"]
print(input_embeds)

[[-0.00291017  0.03561941 -0.04315067 ...  0.00856238 -0.0259888
  -0.03954376]]

model.encode()には次の引数を指定することもできます。

・batch_size : バッチサイズ (複数同時生成する際の速度とメモリのトレードオフ)
・max_length : 最大長 (長さが必要ない場合、小さい値を設定して高速化)

(4) Embeddingの次元数の確認。

# Embeddingの次元数の確認
print(len(input_embeds[0]))

(5) 対象テキストのEmbeddingの生成。
入力テキストと類似度を比較する対象テキストを準備します。

# 対象テキストからのEmbeddingの生成
target_texts = [
    "好きな食べ物は何ですか?",
    "どこにお住まいですか?",
    "朝の電車は混みますね",
    "今日は良いお天気ですね",
    "最近景気悪いですね"]
target_embeds = model.encode(target_texts)["dense_vecs"]

(6) 類似度の計算。
@は行列積演算子、.Tは転置操作です。

# 類似度の計算
similarity = input_embeds @ target_embeds.T
print(similarity)

(7) 類似度が高いテキストの確認。

# 類似度が高いテキストの確認
import numpy as np
print(target_texts[np.argmax(similarity)])

今日は良いお天気ですね

対象テキスト内で、「今日は雨振らなくてよかった」に最も関連するものが「今日は良いお天気ですね」であることがわかりました。