🐕

openai apiのtext-embedding-ada-002でのRateLimitErrorの回避

2024/01/11に公開

text-embedding-ada-002で埋め込み表現を取得しようと、公式リファレンスのコードそのままに実行したところ、以下のエラーが出ました。
https://platform.openai.com/docs/guides/embeddings/use-cases

from openai import OpenAI
client = OpenAI()

def get_embedding(text, model="text-embedding-ada-002"):
   text = text.replace("\n", " ")
   return client.embeddings.create(input = [text], model=model).data[0].embedding

df['ada_embedding'] = df.embedding_text.apply(lambda x: get_embedding(x, model='text-embedding-ada-002'))
#RateLimitError: Error code: 429 - {'error': {'message': 'Rate limit reached for text-embedding-ada-002 in organization org-hL6roglYpeajaUkMcVDUvEtl on requests per min (RPM): Limit 3, Used 3, Requested 1. Please try again in 20s. Visit https://platform.openai.com/account/rate-limits to learn more. You can increase your rate limit by adding a payment method to your account at https://platform.openai.com/account/billing.', 'type': 'requests', 'param': None, 'code': 'rate_limit_exceeded'}}

原因はエラーメッセージの通り、レート制限。requests per min (RPM)が多すぎるのが問題。

解決策として、1つのリクエストで複数の行の埋め込み表現を取得することで、RPMを下げる。
まず、各行のトークン数を数える。

import tiktoken
embedding_encoding = "cl100k_base"  # this the encoding for text-embedding-ada-002
encoding = tiktoken.get_encoding(embedding_encoding)
df["num_tokens"]=df["embedding_text"].map(lambda x: len(encoding.encode(x)))

TMP上限ギリギリのトークン数でリクエストすることで、制限に引っ掛からないようにする。

import time
embedding_model="text-embedding-ada-002"
embedding_texts=[]
token_sum=0
max_input=2048 # inputのlistの最大長は2048 https://github.com/openai/openai-python/issues/519
max_tmp=150000*2/3
responses=[]
for i,row in df.iterrows():
    token_sum+=row["num_tokens"]
    if (token_sum>=max_tmp) or (len(embedding_texts)>=max_input):
        print(i,"/",df.shape[0],len(embedding_texts),token_sum-row["num_tokens"])
        response=client.embeddings.create(input = embedding_texts, model=embedding_model)
        responses+=response.data
        token_sum=row["num_tokens"]
        embedding_texts=[row["embedding_text"]]
        time.sleep(65)
    else:
        embedding_texts+=[row["embedding_text"]]
    
responses+=client.embeddings.create(input = embedding_texts, model=embedding_model).data
df["ada_embedding"]=[r.embedding for r in responses]

これでデータが多すぎる場合(200*150000トークン以上)以外はレート制限に引っ掛からないはず。
公式リファレンスに明示されていない色々な制約があるみたいなので少し大変。
https://github.com/openai/openai-python/issues/519
に細かい調査記録あり。
修正があればコメントお願いします。

Discussion