从数据集获取嵌入

,
2022年3月10日
在 Github 中打开

此 notebook 示例展示了如何从大型数据集中获取嵌入。

1. 加载数据集

此示例中使用的数据集是来自 Amazon 的 fine-food reviews。该数据集总共包含截至 2012 年 10 月 Amazon 用户留下的 568,454 条食品评论。我们将使用该数据集的子集，包含 1,000 条最新的评论，以作演示说明之用。这些评论是英文的，并且倾向于正面或负面。每条评论都有 ProductId、UserId、Score、评论标题 (Summary) 和评论正文 (Text)。

我们将把评论摘要和评论文本合并成一个组合文本。模型将编码这个组合文本，并输出一个单向量嵌入。

要运行此 notebook，您需要安装：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  # 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)
时间 ProductId UserId 评分 摘要 文本 组合
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)