CLOVER🍀

That was when it all began.

SentenceTransformersとintfloat/multilingual-e5でテキスト埋め込みを試してみる

これは、なにをしたくて書いたもの?

テキストの埋め込み(ベクトル化)の方法をいろいろ見ているのですが、SentenceTransformersというものを押さえておいた方が
よさそうに思ったので試してみることにしました。

SentenceTransformers

SentenceTransformersのWebサイトはこちら。

SentenceTransformers Documentation — Sentence-Transformers documentation

GitHubリポジトリーはこちら。

GitHub - UKPLab/sentence-transformers: Multilingual Sentence & Image Embeddings with BERT

紹介文を見てみます。

SentenceTransformersはセンテンス、テキスト、画像埋め込みのPythonフレームワークです。

SentenceTransformers is a Python framework for state-of-the-art sentence, text and image embeddings. The initial work is described in our paper Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks.

BERT-Networksの修正版である、Sentence-BERTというものらしいです。

[1908.10084] Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks

SentenceTransformersを使って埋め込みを計算してコサイン類似度で比較したりすることで、テキストの類似度、セマンティック類似度検索、
パラグラフマイニングなどで便利だそうです。

You can use this framework to compute sentence / text embeddings for more than 100 languages. These embeddings can then be compared e.g. with cosine-similarity to find sentences with a similar meaning. This can be useful for semantic textual similar, semantic search, or paraphrase mining.

実装としては、PyTorchとTransformersの上に構築されています。

今回はこちらを使って、テキストの埋め込みを計算してテキストの類似度を使用した簡単な検索をしてみたいと思います。

Computing Sentence Embeddings — Sentence-Transformers documentation

Semantic Textual Similarity — Sentence-Transformers documentation

他にも、セマンティック検索(こちらは今回は使わないことにします)

Semantic Search — Sentence-Transformers documentation

リランキング

Retrieve & Re-Rank — Sentence-Transformers documentation

クラスタリングなどの使い方のドキュメントが並んでいます。

Clustering — Sentence-Transformers documentation

使いたい用途に応じて、こちらを眺めていくのがいいんでしょうね。

モデルは、事前トレーニング済みのものがあります。多くはHugging Face Hubから取得することになります。

Pretrained Models — Sentence-Transformers documentation

バランスがよさそうなのは、all-MiniLM-L6-v2みたいですね。

The all-* models where trained on all available training data (more than 1 billion training pairs) and are designed as general purpose models. The all-mpnet-base-v2 model provides the best quality, while all-MiniLM-L6-v2 is 5 times faster and still offers good quality.

テキストにはできれば日本語を使いたいので、今回はMTEBを調べた時に気になったintfloat/multilingual-e5のものを使ってみたいと思います。

intfloat/multilingual-e5-large · Hugging Face

intfloat/multilingual-e5-base · Hugging Face

intfloat/multilingual-e5-small · Hugging Face

実際に使うのは、こちらにします。

intfloat/multilingual-e5-base · Hugging Face

環境

今回の環境は、こちら。

$ python3 --version
Python 3.10.12


$ pip3 --version
pip 22.0.2 from /usr/lib/python3/dist-packages/pip (python 3.10)

SentenceTransformersを使ってテキスト埋め込みを行ってみる

まずはSentenceTransformersを使ってテキスト埋め込みを行ってみます。

SentenceTransformersのインストール。

$ pip3 install sentence-transformers

依存関係。

$ pip3 list
Package                  Version
------------------------ ----------
certifi                  2023.11.17
charset-normalizer       3.3.2
click                    8.1.7
filelock                 3.13.1
fsspec                   2023.12.2
huggingface-hub          0.20.2
idna                     3.6
Jinja2                   3.1.2
joblib                   1.3.2
MarkupSafe               2.1.3
mpmath                   1.3.0
networkx                 3.2.1
nltk                     3.8.1
numpy                    1.26.3
nvidia-cublas-cu12       12.1.3.1
nvidia-cuda-cupti-cu12   12.1.105
nvidia-cuda-nvrtc-cu12   12.1.105
nvidia-cuda-runtime-cu12 12.1.105
nvidia-cudnn-cu12        8.9.2.26
nvidia-cufft-cu12        11.0.2.54
nvidia-curand-cu12       10.3.2.106
nvidia-cusolver-cu12     11.4.5.107
nvidia-cusparse-cu12     12.1.0.106
nvidia-nccl-cu12         2.18.1
nvidia-nvjitlink-cu12    12.3.101
nvidia-nvtx-cu12         12.1.105
packaging                23.2
pillow                   10.2.0
pip                      22.0.2
PyYAML                   6.0.1
regex                    2023.12.25
requests                 2.31.0
safetensors              0.4.1
scikit-learn             1.3.2
scipy                    1.11.4
sentence-transformers    2.2.2
sentencepiece            0.1.99
setuptools               59.6.0
sympy                    1.12
threadpoolctl            3.2.0
tokenizers               0.15.0
torch                    2.1.2
torchvision              0.16.2
tqdm                     4.66.1
transformers             4.36.2
triton                   2.1.0
typing_extensions        4.9.0
urllib3                  2.1.0

テキスト埋め込みのドキュメントを元に、ソースコードを書いてみます。

compute_embeddings.py

import time
from sentence_transformers import SentenceTransformer

start_time = time.perf_counter()

model = SentenceTransformer("intfloat/multilingual-e5-base")

sentences = [
    "passage: Hello World.",
    "passage: こんにちは、世界。"
]

embeddings = model.encode(sentences)

for sentence, embedding in zip(sentences, embeddings):
    print(f"Sentence: {sentence}")
    print(f"Embedding: {embedding}")
    print(f"Dimention: {len(embedding)}")
    print("")

elapsed_time = time.perf_counter() - start_time
print(f"elapsed time = {elapsed_time:.3f} sec")

Computing Sentence Embeddings — Sentence-Transformers documentation

まず、SentenceTransformerのインスタンスを作成。コンストラクターには、使用するモデルを指定します。
このモデルは、実行時にHugging Face Hubからダウンロードされます。

model = SentenceTransformer("intfloat/multilingual-e5-base")

今回はintfloat/multilingual-e5-baseを使います。

intfloat/multilingual-e5-base · Hugging Face

テキスト埋め込みはSentenceTransformer#encodeを呼び出して行うのですが、以下のようにリストを渡して一気にベクトル化
することもできるようです。

sentences = [
    "passage: Hello World.",
    "passage: こんにちは、世界。"
]

embeddings = model.encode(sentences)

passage:というのは、multilingual-e5を使う時に必要な接頭辞です。もうひとつquery:というものもあります。

詳しくはFAQへ。

intfloat/multilingual-e5-base · Hugging Face

SentenceTransformerのAPIドキュメントはこちらです。

SentenceTransformer — Sentence-Transformers documentation

encodeメソッドの型宣言は以下のようになっています。

encode(sentences: Union[str, List[str]], batch_size: int = 32, show_progress_bar: Optional[bool] = None, output_value: str = 'sentence_embedding', convert_to_numpy: bool = True, convert_to_tensor: bool = False, device: Optional[str] = None, normalize_embeddings: bool = False)→ Union[List[torch.Tensor], numpy.ndarray, torch.Tensor]

全体的に、内容自体はドキュメントの通りのプログラムなのですが、最後にベクトルに変換後の次元数と処理時間も出力するように
しています。

for sentence, embedding in zip(sentences, embeddings):
    print(f"Sentence: {sentence}")
    print(f"Embedding: {embedding}")
    print(f"Dimention: {len(embedding)}")
    print("")

実行結果。

$ python3 compute_embeddings.py
Sentence: passage: Hello World.
Embedding: [ 1.79633386e-02  2.24528369e-02 -8.03033356e-03  1.34270154e-02
  7.36045977e-03 -1.24695031e-02 -8.00018199e-03 -2.61394307e-02
  3.89445797e-02  3.66625227e-02 -9.62788798e-03 -2.38732807e-02

  〜省略〜

  1.16878161e-02  1.80561915e-02 -2.47767903e-02 -1.29564367e-02
  3.56676653e-02  3.35939322e-03  1.21231386e-02 -6.40960131e-03
  4.65008318e-02 -5.66708297e-02 -4.44268845e-02  2.82738358e-02]
Dimention: 768

Sentence: passage: こんにちは、世界。
Embedding: [ 1.83719657e-02  1.77575648e-02 -1.70415696e-02  2.91866586e-02
  1.19520919e-02  1.15471681e-04 -3.95672303e-03 -2.57793497e-02
  2.95180883e-02  4.05238494e-02 -5.44774905e-03 -4.13883068e-02
  1.82050198e-01  4.32411432e-02 -6.16173185e-02 -1.98859628e-02

  〜省略〜

  1.78205948e-02  1.36664063e-02 -1.75660066e-02 -1.06884157e-02
  4.39340137e-02  9.70378146e-03  3.32553945e-02 -9.03843204e-04
  2.76353564e-02 -4.93616387e-02 -3.61337587e-02  3.97946090e-02]
Dimention: 768

elapsed time = 1.791 sec

とりあえず、動かすことはできました、と。

次元数は768ですね。

ドキュメントを見てみると、次元削除の記述に埋め込みの次元はデフォルトでは768(baseモデル)または1024(largeモデル)と
ありました。

By default, the pretrained models output embeddings with size 768 (base-models) or with size 1024 (large-models).

Model Distillation / Dimensionality Reduction

テキストの類似度を使用した簡単な検索をしてみる

次に、テキストの類似度を使用して簡単な検索をしてみます。

こちらですね。

Semantic Textual Similarity — Sentence-Transformers documentation

やりたいこと自体はセマンティック検索な気もしますが、まずは基本的な動きを見たいということで。

Semantic Search — Sentence-Transformers documentation

作成したソースコードはこちら。

text_similarity_search.py

import sys
import time
from sentence_transformers import SentenceTransformer, util

start_time = time.perf_counter()

model = SentenceTransformer("intfloat/multilingual-e5-base")

documents = [
    "passage: 特急に乗っています。",
    "passage: 今から実家へ帰ります。",
    "passage: 九州へ行きます。",
    "passage: リンゴを食べます。",
    "passage: 釣りに行ってきます。",
    "passage: 魚を食べます。",
    "passage: 肉を食べます。",
    "passage: みかんが欲しいです。",
    "passage: セーターを着ます。",
    "passage: コートを着ます。"
]

documents_with_embedding = [{
    "document": document, "embedding": model.encode(document, convert_to_tensor=True)
} for document in documents]

query = f"query: {sys.argv[1]}"
query_embedding = model.encode(query, convert_to_tensor=True)

documents_with_similarity = [{
    "document": d["document"],
    "embedding": d["embedding"],
    "similarity": util.cos_sim(query_embedding, d["embedding"]).item()  # コサイン類似度をとった後でtorch.TensorからPythonの型へ変換
} for d in documents_with_embedding]

sorted_documents = sorted(documents_with_similarity, key=lambda d: d["similarity"], reverse=True)

print("ranking:")
for document in sorted_documents:
    print(f"  document: {document['document']}")
    print(f"  similarity: {document['similarity']:.3f}")

print()

elapsed_time = time.perf_counter() - start_time
print(f"elapsed time = {elapsed_time:.3f} sec")

適当なテキストに対して

documents = [
    "passage: 特急に乗っています。",
    "passage: 今から実家へ帰ります。",
    "passage: 九州へ行きます。",
    "passage: リンゴを食べます。",
    "passage: 釣りに行ってきます。",
    "passage: 魚を食べます。",
    "passage: 肉を食べます。",
    "passage: みかんが欲しいです。",
    "passage: セーターを着ます。",
    "passage: コートを着ます。"
]

テキスト埋め込みを行います。

documents_with_embedding = [{
    "document": document, "embedding": model.encode(document, convert_to_tensor=True)
} for document in documents]

今回はリストではなく、ひとつずつにしました。

検索クエリーはコマンドライン引数として受け取り、テキスト埋め込みを行います。

query = f"query: {sys.argv[1]}"
query_embedding = model.encode(query, convert_to_tensor=True)

この時のmultilingual-e5の接頭辞はquery:で、コマンドライン引数に追加するようにしています。

SentenceTransformer#encodeの引数でconvert_to_tensor=Trueというものを指定していますが、これは戻り値をtorch.Tensor
(呼び出し方によってはList[torch.Tensor])にする設定です。デフォルトではnumpy.ndarrayで返ってきます。

convert_to_tensor – If true, you get one large tensor as return. Overwrites any setting from convert_to_numpy

SentenceTransformer#encode

こうする必要があるのは、この後でコサイン類似度を計算するためのutil.cos_sim関数の引数がtorch.Tensorだからですね。

util.cos_sim

あとは、クエリーと用意していたテキストのそれぞれのベクトル化の結果から、コサイン類似度を算出してソートします。

documents_with_similarity = [{
    "document": d["document"],
    "embedding": d["embedding"],
    "similarity": util.cos_sim(query_embedding, d["embedding"]).item()  # コサイン類似度をとった後でtorch.TensorからPythonの型へ変換
} for d in documents_with_embedding]

sorted_documents = sorted(documents_with_similarity, key=lambda d: d["similarity"], reverse=True)

最後に結果表示。

print("ranking:")
for document in sorted_documents:
    print(f"  document: {document['document']}")
    print(f"  similarity: {document['similarity']:.3f}")

なお、この時torch.Tensorだとちょっと表示に違和感があるので、`torch.Tensor#itemでPythonの組み込み型に戻しておきました。

    "similarity": util.cos_sim(query_embedding, d["embedding"]).item()  # コサイン類似度をとった後でtorch.TensorからPythonの型へ変換

いくつか試してみましょう。

$ python3 text_similarity_search.py 帰省する
ranking:
  document: passage: 今から実家へ帰ります。
  similarity: 0.849
  document: passage: 九州へ行きます。
  similarity: 0.829
  document: passage: 釣りに行ってきます。
  similarity: 0.815
  document: passage: セーターを着ます。
  similarity: 0.807
  document: passage: 特急に乗っています。
  similarity: 0.807
  document: passage: みかんが欲しいです。
  similarity: 0.806
  document: passage: 肉を食べます。
  similarity: 0.801
  document: passage: コートを着ます。
  similarity: 0.798
  document: passage: 魚を食べます。
  similarity: 0.795
  document: passage: リンゴを食べます。
  similarity: 0.783

elapsed time = 2.248 sec


$ python3 text_similarity_search.py 食事
ranking:
  document: passage: 肉を食べます。
  similarity: 0.868
  document: passage: 魚を食べます。
  similarity: 0.856
  document: passage: リンゴを食べます。
  similarity: 0.833
  document: passage: みかんが欲しいです。
  similarity: 0.820
  document: passage: 釣りに行ってきます。
  similarity: 0.814
  document: passage: 九州へ行きます。
  similarity: 0.814
  document: passage: 今から実家へ帰ります。
  similarity: 0.809
  document: passage: 特急に乗っています。
  similarity: 0.803
  document: passage: セーターを着ます。
  similarity: 0.802
  document: passage: コートを着ます。
  similarity: 0.793

elapsed time = 2.130 sec


$ python3 text_similarity_search.py 寒い
ranking:
  document: passage: セーターを着ます。
  similarity: 0.815
  document: passage: コートを着ます。
  similarity: 0.814
  document: passage: みかんが欲しいです。
  similarity: 0.814
  document: passage: 特急に乗っています。
  similarity: 0.813
  document: passage: 九州へ行きます。
  similarity: 0.810
  document: passage: 今から実家へ帰ります。
  similarity: 0.807
  document: passage: 釣りに行ってきます。
  similarity: 0.807
  document: passage: 肉を食べます。
  similarity: 0.806
  document: passage: 魚を食べます。
  similarity: 0.795
  document: passage: リンゴを食べます。
  similarity: 0.782

elapsed time = 2.075 sec

それっぽい感じになっているような気がしますね。よさそうです。

おわりに

SentenceTransformersを使って、テキスト埋め込みを試してみました。
モデルはintfloat/multilingual-e5のものを使うことで、日本語のテキストに対しても確認しやすい結果になるなと。

SentenceTransformersはテキスト埋め込みを使うところでは頻出しそうなので、覚えておきましょう。