nikkie-ftnextの日記

イベントレポートや読書メモを発信

ライブラリsentence-transformersのサンプルコードを動かし、英語や日本語テキストからembeddingsやその類似度を計算する

NLP(自然言語処理)

はじめに

アヤさん、たんじょーび、おめでとう！！ nikkieです。
みんなアイうた見ていて嬉しい限り♪

sentence-transformersというPythonのライブラリがあります。
こいつでembeddings（テキストの埋め込み表現）が計算できるらしく、気になったので触ってみました。
※レベル感としては使い出しレベル、やってみた系です。

目次

動作環境

  • macOS 12.6.6
    • CPU環境です
  • Python 3.10.9
  • sentence-transformers 2.2.2

pip install sentence-transformersで入ったライブラリのうち主なもののバージョンはこちら

  • torch 2.0.1
  • transformers 4.30.2
  • tokenizers 0.13.3
  • huggingface-hub 0.16.4
  • numpy 1.25.0
  • scikit-learn 1.3.0

ドキュメントの例でembeddingsを計算（英語テキスト）

ドキュメントトップのUsageの例を写経していきます。
これは英語テキストのembeddingsを計算します。
https://www.sbert.net/#usage

  • 指定しているモデルは all-MiniLM-L6-v2
    • テキストを384次元のembeddingsに変換
    • モデルサイズが非常に小さくダウンロードは一瞬でした
      • 小さいので性能は高くないのかもしれません
      • スクリプトの動作確認用途にぱっと使えそうなのが好感触
  • 英語テキスト3文
    • 1文目と2文目は近そうですよね（sentenceという共通語）
    • 3文目は1文目・2文目と全然関係なさそうです

人が読んで感じる文同士の意味的な近さがembeddingsに表れているか、コサイン類似度を求めてみます。

>>> from sklearn.metrics.pairwise import cosine_similarity
>>> cosine_similarity(embeddings[0].reshape(1, -1), embeddings[1].reshape(1, -1))
array([[0.53807926]], dtype=float32)
>>> cosine_similarity(embeddings[0].reshape(1, -1), embeddings[2].reshape(1, -1))
array([[0.11805625]], dtype=float32)

2つのベクトルのコサイン類似度1は-1から1の値です。
1に近いほど2つのベクトルの向きが揃っている（＝近い）ということですから、算出結果から

  • 1文目と2文目は近い
  • 1文目と3文目は遠い

と、人が読んだ時の感覚と一致すると言えそうな結果でした

日本語テキストからembeddingsを計算

英語について動かせたので、次は日本語の例を探します。
見つかったのがnpakaさん2のnote。

  • 多言語を扱えるモデル stsb-xlm-r-multilingual
    • 768次元のembeddings
    • モデルサイズは大きく、1GBありました（ダウンロードにちょっとかかります。ダウンロード後に保存されるので一度だけです）
  • sentense-transformersの中にコサイン類似度を計算する関数がある
    • sentence_transformers.util.cos_sim
      • npakaさんが使ったpytorch_cos_simは中でcos_simを呼ぶだけでした3
    • これを使うためにencodeメソッドにconvert_to_tensor=Trueを指定！
  • cos_simの出力する類似度（scores）のsizeは[1, 5]
    • 768次元の1つのベクトル（[1, 768]）と5つのベクトル（[5, 768]）の類似度なのでこのsizeですね
    • 類似度最大の文を知りたいのでargmaxでインデックスを得ています
      • 2つのベクトルと5つのベクトルになったら2つargmaxが返るようにdim=1と指定しています4

文: 今日は雨降らなくてよかった
類似文: 今日は良いお天気ですね
類似度: tensor([[ 0.1204, -0.0705,  0.1849,  0.5848,  0.5198]])
文: ハンバーガーは好きですか？
類似文: 好きな食べ物は何ですか？
類似度: tensor([[ 0.5839,  0.0653, -0.0964,  0.1122,  0.0838]])

類似文は近いものが返せていますね！
一度モデルをダウンロードした状態でスクリプトをCPUで動かすと、1回の類似度計算に5-6秒かかりました。

終わりに

sentence-transformersライブラリを使って、テキストのembeddingsや類似度を計算しました。
英語や日本語用のモデルをダウンロードしてきて、ローカル環境で（＝Web APIは呼び出さずに）embeddingsを計算できます。
Hugging Faceのtransformersライブラリのように、非常に簡単に動かせますね。

OpenAIのEmbeddings APIでもテキストからembeddingsを計算できますが、従量課金です。
embeddingsを保存するvector store関係の素振りのためにEmbeddings APIを毎回呼ぶというのは少し抵抗がありました（塵も積もれば山となりそう）。
そこで、費用をかけずにembeddingsを算出する方法（ローカル環境だけで完結する方法）を知りたく、sentence-transformersを触ってみました。
素振りにかかる費用はなるべく抑えて、精度を求める段階になったらぶっぱ（金の弾丸）というのが私の好むアプローチです。

  1. https://scikit-learn.org/stable/modules/metrics.html#cosine-similarity↩
  2. LLM関係のnoteの数がすごいですし、書籍も出たばかりで、信頼できる情報源という認識です。 ↩
  3. エイリアスにすぎないようです。ref: https://github.com/UKPLab/sentence-transformers/blob/v2.2.2/sentence_transformers/util.py#L23-L28↩
  4. torch.Tensor.argmaxのドキュメントから案内されるtorch.argmaxのドキュメントより↩