これは、なにをしたくて書いたもの？
テキストの埋め込み（ベクトル化）の方法をいろいろ見ているのですが、SentenceTransformersというものを押さえておいた方が
よさそうに思ったので試してみることにしました。
SentenceTransformers
SentenceTransformersのWebサイトはこちら。
SentenceTransformers Documentation — Sentence-Transformers documentation
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
こうする必要があるのは、この後でコサイン類似度を計算するためのutil.cos_sim関数の引数がtorch.Tensorだからですね。
あとは、クエリーと用意していたテキストのそれぞれのベクトル化の結果から、コサイン類似度を算出してソートします。
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はテキスト埋め込みを使うところでは頻出しそうなので、覚えておきましょう。
