TensorFlow.jsのUniversal Sentence Encoderで、ブラウザ上でEmbeddingsを作る
ということをやってみた。
できたにはできた、が・・・、って感想だった。
モチベーション
前に、OpenAIのAPIを使って、テキストの特徴を表すEmbeddingsを作ってた。
で、1536次元のベクトルはなかなかの精度で良いものの・・・、
っていうところで、ブラウザで動くモデルがあれば、もっと手軽にできるのでは?!という。
ブラウザで動くモデルを探す
Embeddingsを作りたいのは、自然言語による検索に利用したいから。
となると、検索クエリもEmbeddingsにする必要があり、それをやれる場所となるとブラウザかCloudflare Workersになる、まずはサイズに縛られないブラウザで、という縛りで。
あれこれ探してみたものの、これだ!っていうのはそうそう見つからず。なんとか見つけて動かせたのが、Tensorflow.jsのUniversal Sentence Encoderを使った実装だったというわけ。
tfjs-models/universal-sentence-encoder at master · tensorflow/tfjs-models · GitHub
やってみた
GitHub Pagesにもアップロードしてあるので、誰でも試せるようになってます(が、通信料に注意)。
- Viteでビルドしていて
- TensorFlow.js関連は遅延ロードさせてある
- ガワだけで230KBもあるし、モデルは28MBとかある
- ロード後、テキストを入力するとEmbeddingsにその場で変換
- 2つ以上のEmbeddingsがあれば、それらのコサイン類似度を計算
って感じ。
2015年の古いMBPでも、1件あたり1秒くらいで512次元のベクトルにしてくれる。
結果に関して
元リポでも提示されてるような、英語に関する類似度はそれっぽい結果になる。
ふ〜む pic.twitter.com/g8KKjc0ny5
— りぃ / Yuji Sugiura (@leader22) May 31, 2023
けど、日本語にすると精度はもちろん良くないし、本来の目的である技術系のコンテキストも汲んでくれてるとは思えなかった。
でもまぁ元のモデルも学習データも英語のみが対象で、OpenAIのそれに比べるとモデルのサイズも小さいし、当然といえば当然か。
ちなみに、デフォルトで使われるモデルはコレだそうな。
https://tfhub.dev/tensorflow/tfjs-model/universal-sentence-encoder-lite/1/default/1
APIとしては他のモデルも使える風になってたけど、何をどう指定してもどこかでエラーになってうまくいかなかった。修行が足りぬ。