nikkie-ftnextの日記

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

multilingual-e5のUsageと同じコードがsentence-transformersで数行で書けます！（ただしtokenizerの呼び出しの引数を変えるのが難しい）

NLP(自然言語処理)

はじめに

ジャングルパーティー！！
異次元フェス、行ってきました、nikkieです🙌

E5はtransformersから触ってもいいのですが、sentence-transformersからだと（特定のニーズを満たすコードは）数行で書けるということを見ていきます。

目次

multilingual-e5のUsage

文の埋め込み表現（embeddings）を得る手法の1つに、E5というモデルがあります。
サイズがいくつかありますが、largeはOpenAIのadaに匹敵するとか。
Hugging Face Hubで公開されています。

Usageのaverage_poolの実装について、先日読み解きました。

これと同じコードをsentence-transformersを使うと数行で書けることを示します。

sentence-transformersなら数行でできるよ

from sentence_transformers import SentenceTransformer

input_texts = [
    "query: how much protein should a female eat",
    "query: 南瓜的家",
]
model = SentenceTransformer("intfloat/multilingual-e5-small")
embeddings = model.encode(input_texts)

実はこちら、Hugging Face Hubにも記載があります。
Support for Sentence Transformersの箇所ですね

処理内容は同じですが、少しだけ異なる点があります

  • 返り値の型
    • Usageのコードはtorch.Tensorです
    • sentence-transformersを使うとNumPyのndarrayです
  • 勾配計算
    • Usageのコードは勾配計算が有効です
      • average_pool関数を@torch.inference_mode（や@torch.no_grad）でデコレートすると無効にできます1
    • sentence-transformersを使うと勾配計算が無効になります

同じembeddingsが得られていることの検証

Python 3.11.4

ライブラリのバージョン

certifi==2023.11.17
charset-normalizer==3.3.2
click==8.1.7
filelock==3.13.1
fsspec==2023.12.1
huggingface-hub==0.19.4
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.2
packaging==23.2
Pillow==10.1.0
PyYAML==6.0.1
regex==2023.10.3
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
sympy==1.12
threadpoolctl==3.2.0
tokenizers==0.15.0
torch==2.1.1
torchvision==0.16.1
tqdm==4.66.1
transformers==4.35.2
typing_extensions==4.9.0
urllib3==2.1.0

なぜ同じになるのか？

以前、SentenceTransformerクラスをインスタンス化して、encodeメソッドを呼ぶだけで済む仕組みについてソースコードを追いかけました。

SentenceTransformerはtorchのnn.Moduleを活用していて、3つのModuleからなります（順番にforwardを呼びます）。

  1. sentence-transformersのTransformerクラス
  2. Pooling
  3. Normalize

ここで、2はUsageのaverage_pool関数と同じです。
実装に違いがあるのですが、結果としては同じになります2。

3もUsageのF.normalizeと同じです。
同じ引数で呼んでいます3

1ですが、Transformerクラスがtokenizerも持ちます。
この呼び出しが同じです。

Usage

tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

sentence-transformers
https://github.com/UKPLab/sentence-transformers/blob/v2.2.2/sentence_transformers/models/Transformer.py#L113

self.tokenizer(*to_tokenize, padding=True, truncation='longest_first', return_tensors="pt", max_length=self.max_seq_length)

SentenceTransformerの実装を見ていくことで、Usageと同様の実装をしていることが確認できました！

注意点としては、tokenizerの呼び出しがハードコードになっているために、呼び出しの引数の変更が難しいです。
すなわち、Usageのコードと同じ引数で呼び出す分には数行のコードで簡単に使えますが、Usageのコードから引数を変える（例：padding=False）場合、変更が大変になっています4。
引数を変える場合、SentenceTransformerクラスを使う旨味が激減していて、私だったらUsageのコードを書くことになるかなと考えています（プルリクチャンスかも）

終わりに

E5のUsageのコードと全く同じコードに限れば、sentence-transformersで数行で書けることを見てきました。

最初E5をHugging Face Hubで見たときに、Usageとsentence-transformersのコードが本当に同じことをやっているかが分かりませんでした。
一つ一つの要素を見た今なら自信を持って言えます、同じです！

  1. https://pytorch.org/docs/stable/notes/autograd.html#grad-modes の表から、Data processingで使うのでinference modeを選びました↩
  2. 読み解き記事のP.S.に書いていますが、sentence-transformersはattention maskで内積をとる実装です↩
  3. https://github.com/UKPLab/sentence-transformers/blob/v2.2.2/sentence_transformers/models/Normalize.py#L14↩
  4. sentence-transformersはsimpleを謳っていますが、この点を考えると私の見地からはsimpleというよりもeasyであるように見受けられます（simpleというんだったら、tokenizer呼び出しの引数を変更する方法も提供されてほしいです）↩