🧊

TensorFlow.jsのUniversal Sentence Encoderで、ブラウザ上でEmbeddingsを作る

ということをやってみた。

できたにはできた、が・・・、って感想だった。

モチベーション

OpenAIのEmbeddingsで全文検索 - console.lealog();

前に、OpenAIのAPIを使って、テキストの特徴を表すEmbeddingsを作ってた。

で、1536次元のベクトルはなかなかの精度で良いものの・・・、

  • API経由なので大量作成が手間
  • それぞれのサイズも結構でかい（JSONにすると25KBくらい）
  • お金もかかる

っていうところで、ブラウザで動くモデルがあれば、もっと手軽にできるのでは？！という。

ブラウザで動くモデルを探す

Embeddingsを作りたいのは、自然言語による検索に利用したいから。
となると、検索クエリもEmbeddingsにする必要があり、それをやれる場所となるとブラウザかCloudflare Workersになる、まずはサイズに縛られないブラウザで、という縛りで。

あれこれ探してみたものの、これだ！っていうのはそうそう見つからず。なんとか見つけて動かせたのが、Tensorflow.jsのUniversal Sentence Encoderを使った実装だったというわけ。

tfjs-models/universal-sentence-encoder at master · tensorflow/tfjs-models · GitHub

やってみた

GitHub - leader22/tfjs-embeddings-by-use: Playground for embeddings by Tensorflow.js Universal Sentence Encoder

GitHub Pagesにもアップロードしてあるので、誰でも試せるようになってます（が、通信料に注意）。

  • Viteでビルドしていて
  • TensorFlow.js関連は遅延ロードさせてある
    • ガワだけで230KBもあるし、モデルは28MBとかある
  • ロード後、テキストを入力するとEmbeddingsにその場で変換
  • 2つ以上のEmbeddingsがあれば、それらのコサイン類似度を計算

って感じ。

2015年の古いMBPでも、1件あたり1秒くらいで512次元のベクトルにしてくれる。

結果に関して

元リポでも提示されてるような、英語に関する類似度はそれっぽい結果になる。


けど、日本語にすると精度はもちろん良くないし、本来の目的である技術系のコンテキストも汲んでくれてるとは思えなかった。

でもまぁ元のモデルも学習データも英語のみが対象で、OpenAIのそれに比べるとモデルのサイズも小さいし、当然といえば当然か。

ちなみに、デフォルトで使われるモデルはコレだそうな。

https://tfhub.dev/tensorflow/tfjs-model/universal-sentence-encoder-lite/1/default/1

APIとしては他のモデルも使える風になってたけど、何をどう指定してもどこかでエラーになってうまくいかなかった。修行が足りぬ。