本笔记本展示了如何从大型数据集中获取嵌入向量的示例。
本示例使用的数据集是来自亚马逊的精细食品评论。该数据集包含截至2012年10月亚马逊用户留下的总计568,454条食品评论。为便于演示，我们将使用该数据集的一个子集，包含最近的1,000条评论。这些评论均为英文，且倾向于表达正面或负面评价。每条评论包含产品ID(ProductId)、用户ID(UserId)、评分(Score)、评论标题(Summary)和评论正文(Text)。
我们将把评论摘要和评论文本合并为一个统一的文本。模型会对这个合并后的文本进行编码，并输出一个单一的向量嵌入。
要运行此笔记本，您需要安装以下依赖项：pandas、openai、transformers、plotly、matplotlib、scikit-learn、torch（transformer依赖项）、torchvision以及scipy。
import pandas as pd
import tiktoken
from utils.embeddings_utils import get_embeddingembedding_model = "text-embedding-3-small"
embedding_encoding = "cl100k_base"
max_tokens = 8000 # the maximum for text-embedding-3-small is 8191# load & inspect dataset
input_datapath = "data/fine_food_reviews_1k.csv" # to save space, we provide a pre-filtered dataset
df = pd.read_csv(input_datapath, index_col=0)
df = df[["Time", "ProductId", "UserId", "Score", "Summary", "Text"]]
df = df.dropna()
df["combined"] = (
"Title: " + df.Summary.str.strip() + "; Content: " + df.Text.str.strip()
)
df.head(2)| 时间 | 产品ID | 用户ID | 评分 | 摘要 | 文本内容 | 合并字段 | |
|---|---|---|---|---|---|---|---|
| 0 | 1351123200 | B003XPF9BO | A3R7JR3FMEBXQB | 5 | 该从哪里开始...又该在哪里结束...关于一棵树... | 想留一些带给芝加哥的家人... | 标题：该从哪里开始...又该在哪里结束...关于一棵树... |
| 1 | 1351123200 | B003JK537S | A3JBPC3WFUT5ZP | 1 | 到货时已破碎 | 非常不满意。当我打开箱子时，大部分... | 标题：到货时已破碎；内容：非常不满意... |
# subsample to 1k most recent reviews and remove samples that are too long
top_n = 1000
df = df.sort_values("Time").tail(top_n * 2) # first cut to first 2k entries, assuming less than half will be filtered out
df.drop("Time", axis=1, inplace=True)
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].tail(top_n)
len(df)1000
# Ensure you have your API key set in your environment per the README: https://github.com/openai/openai-python#usage
# This may take a few minutes
df["embedding"] = df.combined.apply(lambda x: get_embedding(x, model=embedding_model))
df.to_csv("data/fine_food_reviews_with_embeddings_1k.csv")a = get_embedding("hi", model=embedding_model)