OpenAI Embeddingの複数テキスト入力と埋め込み次元の削減

はじめに

Embeddingはテキストをベクトルで表すことができる技術で、検索やレコメンドなどで活用されます。今回はOpenAIのドキュメントをもとに備忘録としてまとめます。

OpenAIのEmbeddingの使い方

1. OpenAIのAPI keysからAPI keyを取得します
2. 環境変数にOPENAI_API_KEYを設定します
3. OpenAIライブラリをインストールします

   pip install openai
   

4. Pythonで実装します

   from openai import OpenAI
   client = OpenAI()
   
   response = client.embeddings.create(
     input="Your text string goes here",
     model="text-embedding-3-small"
   )
   
   print(response.data[0].embedding)
   

テキストの最大トークン数は8192という制限があるので注意が必要です。

複数テキストの入力

inputには複数のテキストを与えることができます。

response = client.embeddings.create(
    input=["input1", "input2", "input3"],
    model="text-embedding-3-small"
)

テキストの最大トークン数の制限はありますが、入力可能な最大件数はOpenAIのドキュメントには記載されていませんでした。そこで、試しに100文字の文字列を1000件入力してみましたが、4秒未満で実行が完了しました。恐らく内部でバッチ処理やメモリを考慮した切り分けがなされており、最大件数は気にする必要がないのではないかと勝手に解釈しています。（ただし、RPMやTPMといったリクエスト数やトークン数に関する制限は別途あるので注意が必要です。）

埋め込み次元の削減

埋め込み次元の削減により、費用削減や計算速度向上などの効果があります。例えば、検索システムでは埋め込み次元が大きいとベクトルストアの費用が高くなります。OpenAIのEmbeddingモデルの中ではtext-embedding-3-smallとtext-embedding-3-largeが次元削減に対応しており、引数にdimensionsを指定することができます。

response = client.embeddings.create(
    input="hello",
    model="text-embedding-3-large",
    dimensions=4
)
print(response) # [-0.9203252196311951, -0.2836061120033264, 0.14934176206588745, 0.22420111298561096]

これらのモデルの場合、dimensionsは1以上3072以下を指定することができます。ただし上例のように4まで小さくすると品質が悪化してしまうため、128くらいが良いのではないかと思います。（Matryoshka Representation Learningという論文では、ここまで次元を小さくしてもそれほど性能悪化しないことが報告されています。）
ちなみにtext-embedding-ada-002は埋め込み次元の削減に対応していないのでご注意ください。

ドキュメントには実装例が豊富にある

OpenAIのEmbeddingsのドキュメントには「埋め込みベースの検索を使用した質問応答」「2Dでのデータ視覚化」「クラスタリング」など用途ごとの実装例が豊富に紹介されています。

その他

ノルム1に正規化されている

OpenAIのEmbeddingモデルはすべて、ベクトルの長さであるノルムが1に正規化されています。そのため、cos類似度と内積が一致します。

カットオフは2021年9月まで

text-embedding-3-smallとtext-embedding-3-largeは2021年9月までの情報を学習しています。それ以降の情報は上手く表現できないことがあります。