从数据集中获取嵌入向量

本笔记本提供了如何从大型数据集中获取嵌入向量的示例。

1. 加载数据集

本示例中使用的数据集是来自亚马逊的精细食品评论。该数据集包含截至2012年10月用户在亚马逊上留下的共计 568,454 条食品评论。我们将使用该数据集的一个子集，包括最近的 1,000 条评论，用于说明目的。这些评论是用英语撰写的，往往是积极的或消极的。每条评论都有一个ProductId、UserId、Score、评论标题（Summary）和评论内容（Text）。

我们将把评论标题和评论内容合并为一个组合文本。模型将对这个组合文本进行编码，并输出一个单一的向量嵌入。

要运行这个笔记本，您需要安装：pandas、openai、transformers、plotly、matplotlib、scikit-learn、torch（transformer dep）、torchvision 和 scipy。

import pandas as pd
import tiktoken

from utils.embeddings_utils import get_embedding

embedding_model = "text-embedding-3-small"
embedding_encoding = "cl100k_base"
max_tokens = 8000  # text-embedding-3-small 的最大长度为 8191。

# 加载并检查数据集
input_datapath = "data/fine_food_reviews_1k.csv"  # 为了节省空间，我们提供了一个预先过滤的数据集。
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)

	Time	ProductId	UserId	Score	Summary	Text	combined
0	1351123200	B003XPF9BO	A3R7JR3FMEBXQB	5	where does one start...and stop... with a tre...	Wanted to save some to bring to my Chicago fam...	Title: where does one start...and stop... wit...
1	1351123200	B003JK537S	A3JBPC3WFUT5ZP	1	Arrived in pieces	Not pleased at all. When I opened the box, mos...	Title: Arrived in pieces; Content: Not pleased...

# 对最新评论进行子采样，选取1000条，并剔除过长样本。
top_n = 1000
df = df.sort_values("Time").tail(top_n * 2)  # 首先截取前2000条记录，假设其中不超过一半会被过滤掉。
df.drop("Time", axis=1, inplace=True)

encoding = tiktoken.get_encoding(embedding_encoding)

# 省略过长无法嵌入的评论
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

2. 获取嵌入并保存以便将来重用

# 请确保按照 README 文件中的说明，在你的环境中设置了 API 密钥：https://github.com/openai/openai-python#usage

# 这可能需要几分钟时间。
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)