🧊

OpenAIのEmbeddingsで全文検索

https://platform.openai.com/docs/guides/embeddings

Embeddingsとは、テキストの類似性や関連性を、画一的なベクトルで表現することらしい。

それによってテキストの種類を分類したり、検索したり、いわゆるレコメンドに使えたりする。

今回は、このブログ記事をベクトル化して、それに対する全文検索を簡単にできる範囲でやってみた。

このドキュメントのユースケースでいうところの、Text search using embeddingsってやつ。

Step1. ブログ記事のエクスポート

はてなブログには記事のエクスポート機能があって、懐かしのMovable typeフォーマットで吐き出せるようになってる。

が、結局それは使わなかった。

  • 吐き出された記事データが、はてな記法やMarkdownではなく、変換後のHTMLだった
  • そもそもちゃんとしたものを作るよいうよりか試したかっただけ

という理由から、管理画面から愚直にコピペしてきて、`記事-{entryId}.md`みたいに保存した。

（そんなに詳しくはないけど、こういう機械学習用途に使う元データのフォーマットってたぶんすごい大事なはずで、HTMLのカッコとかただのノイズでしかないし、トークン数とかにも響いてきそうやなって思った。）

Step2. 記事本文のベクトル化

Embeddings用のOpenAI API使う。

https://platform.openai.com/docs/guides/embeddings/embeddings

コードで書くとこんな感じで、`data`がブログ記事の本文。

/** @type {(API_KEY: string, data: string) => Promise<number[] | Error>} */
export const fetchEmbeddingsVector = async (API_KEY, data) => {
  const res = await fetch("https://api.openai.com/v1/embeddings", {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${API_KEY}`,
    },
    body: JSON.stringify({
      input: data,
      model: "text-embedding-ada-002",
    }),
  });

  const json = await res.json();

  if (!res.ok) {
    console.error("ERR: OpenAI API request failed");
    console.error(res.status, json.error);
    return new Error(json.error.message);
  }

  const vector = json.data.at(0).embedding;
  if (vector.length === 0) {
    console.error("ERR: OpenAI API returned an empty vector");
    console.error(JSON.stringify(json, null, 2));
    return new Error("Empty vector");
  }

  return vector;
};

無料アカウントの場合のレート制限は、

  • 3 RPM
  • 150,000 TPM x 200(ada) = 30,000,000 TPM

という2軸だそうな。

https://platform.openai.com/docs/guides/rate-limits/what-are-the-rate-limits-for-our-api

愚直に並列実行すると即レートに引っかかるので、ちゃんとやるときはリトライとか20秒待ちとか考慮する。

こうして得られたベクトルも、`ベクトル-{entryId}.json`みたいに保存しておく。

このベクトル、なんと次元数が1536もあって、最初にレスポンス見たときはひゃ〜ってなった。

Step3. UIを作って質問を受け付ける

今回はとにかく試せればよかったので、CLIから`process.argv`で雑に取得した。

書くことなし。

Step4. 質問をベクトル化して、記事のベクトルと比較

先の記事本文をベクトル化したのと同様に、まずは質問文もベクトルに変換する。

こうすることで、ベクトル同士を比較することができるようになり、もっともそれらしいことが書いてある記事を見つけられると。

で、ベクトル同士を比較して、似たようなベクトルかどうか？を判定するには、`cosine_similarity`（コサイン類似度）なるものを計算するのが一般的らしい。

2つのベクトルの内積（＝向きと大きさを持つベクトル同士の掛け算）を、2つのベクトルの大きさ（＝L2ノルム）で割ることで計算される。
https://atmarkit.itmedia.co.jp/ait/articles/2112/08/news020.html

なんのこっちゃだがコードにするとこう。

/** @type {(x: number[], y: number[]) => number} */
const dot = (x, y) => {
  let sum = 0;
  for (let i = 0; i < x.length; i++) sum += x[i] * y[i];

  return sum;
};

/** @type {(arr: number[]) => number} */
const l2norm = (arr) => {
  let s = 1;
  let t = 0;
  for (let i = 0; i < arr.length; i++) {
    const val = arr[i];
    const abs = val < 0 ? -val : val;
    if (abs > 0) {
      if (abs > t) {
        const r = t / val;
        s = 1 + s * r * r;
        t = abs;
      } else {
        const r = val / t;
        s = s + r * r;
      }
    }
  }

  return t * Math.sqrt(s);
};

/** @type {(a: number[], b: number[]) => number} */
export const cosineSimilarity = (a, b) => dot(a, b) / (l2norm(a) * l2norm(b));

まあこれで準備は整ったので、

  • 質問文をEmbeddingsのAPIでベクトルにし
  • 保存しておいた記事のベクトルたちとコサイン類似度を比較していき
  • 一定以上の値かつ、大きい値が出た順に選ぶ
    • 手元でやってた感じだと、少なくとも`0.8`以上ないと信用できないなって感じ

という流れ。

適当なキーワードで検索してみると、確かにそれらしい記事がヒットする！

今回は記事数が大したことないので総当りで処理してみたけど、記事数が増えるとどう考えてもスケールしないので、それ用のソリューションが必要になる。

軽く調べてみると、ベクトルを型として扱える専用のデータベースもあるらしい。

知らないことばかりである。

というわけで

Embeddingsによる全文検索をやってみた。

とはいえここまでの内容であれば、下準備やDBなど仕組みを用意するコストを考えたときに、単なるキーワード検索でよくないか？ってなったのも正直なところで、そこはまあ一長一短で物は使いよう・・・という感じらしい。

全文検索といえばAlgoliaとかどうやってるのか調べてみたら、いろいろおもしろいことが書いてあった。

あとは、ChatGPTにもそういう特定のソースをつなげるプラグインがあったはずで、これでも同じようなことができるんかな？

https://github.com/openai/chatgpt-retrieval-plugin