import pandas as pd
import tiktoken
import openai
import os
from openai.embeddings_utils import get_embedding, get_embeddings
openai.api_key = os.environ.get("OPENAI_API_KEY")
# embedding model parameters
embedding_model = "text-embedding-ada-002"
embedding_encoding = "cl100k_base" # this the encoding for text-embedding-ada-002
max_tokens = 8000 # the maximum for text-embedding-ada-002 is 8191
# import data/toutiao_cat_data.txt as a pandas dataframe
df = pd.read_csv('data/toutiao_cat_data.txt', sep='_!_', names=['id', 'code', 'category', 'title', 'keywords'])
df = df.fillna("")
df["combined"] = (
"标题: " + df.title.str.strip() + "; 关键字: " + df.keywords.str.strip()
)
print("Lines of text before filtering: ", len(df))
encoding = tiktoken.get_encoding(embedding_encoding)
# omit reviews that are too long to embed
df["n_tokens"] = df.combined.apply(lambda x: len(encoding.encode(x)))
df = df[df.n_tokens <= max_tokens]
print("Lines of text after filtering: ", len(df))
OpenAI 提供的接口限制了每条数据的长度。这里使用的 text-embedding-ada-002 的模型，支持的长度是每条记录 8191 个 Token。所以我们在实际发送请求前，需要计算一下每条记录有多少 Token，超过 8000 个的需要过滤掉。不过，在我们这个数据集里，只有新闻的标题，所以不会超过这个长度。但是你在使用其他数据集的时候，可能就需要过滤下数据，或者采用截断的方法，只用文本最后 8000 个 Token。我们在这里，调用了 Tiktoken 这个库，使用了 cl100k_base 这种编码方式，这种编码方式和 text-embedding-ada-002 模型是一致的。如果选错了编码方式，你计算出来的 Token 数量可能和 OpenAI 的不一样。
如果你直接一条条调用 OpenAI 的 API，很快就会遇到报错。这是因为 OpenAI 对 API 的调用进行了限速（Rate Limit）。如果你过于频繁地调用，就会遇到限速的报错。而如果你在报错之后继续持续调用，限速的时间还会被延长。那怎么解决这个问题呢？我习惯选用 backoff 这个 Python 库，在调用的时候如果遇到报错了，就等待一段时间，如果连续报错，就拉长等待时间。通过 backoff 改造的代码: （需要提前install backoff）
import backoff
@backoff.on_exception(backoff.expo, openai.error.RateLimitError)
def get_embedding_with_backoff(**kwargs):
return get_embedding(**kwargs)
# randomly sample 10k rows
df_10k = df.sample(10000, random_state=42)
df_10k["embedding"] = df_10k.combined.apply(lambda x : get_embedding_with_backoff(text=x, engine=embedding_model))
df_10k.to_csv("data/toutiao_cat_data_10k_with_embeddings.csv", index=False)
此文章为3月Day30学习笔记，内容来源于极客时间《05｜善用Embedding，我们来给文本分分类 (geekbang.org)》
