跳到主要内容

嵌入的文本超过模型的最大上下文长度

nbviewer

OpenAI的嵌入模型无法嵌入超过最大长度的文本。最大长度因模型而异,以 tokens 为单位,而不是字符串长度。如果您对标记化不熟悉,请查看如何使用tiktoken计算标记数

本笔记本展示了如何处理长于模型最大上下文长度的文本。我们将演示如何使用text-embedding-3-small中的嵌入,但相同的思路也适用于其他模型和任务。要了解更多关于嵌入的信息,请查看OpenAI嵌入指南

1. 模型上下文长度

首先,我们选择模型并定义一个从API获取嵌入的函数。

from openai import OpenAI
import os
import openai
from tenacity import retry, wait_random_exponential, stop_after_attempt, retry_if_not_exception_type

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "<your OpenAI API key if not set as env var>"))

EMBEDDING_MODEL = 'text-embedding-3-small'
EMBEDDING_CTX_LENGTH = 8191
EMBEDDING_ENCODING = 'cl100k_base'

# 我们确保不对无效请求进行重试,因为这正是我们想要展示的情况。
@retry(wait=wait_random_exponential(min=1, max=20), stop=stop_after_attempt(6), retry=retry_if_not_exception_type(openai.BadRequestError))
def get_embedding(text_or_tokens, model=EMBEDDING_MODEL):
return client.embeddings.create(input=text_or_tokens, model=model).data[0].embedding

text-embedding-3-small 模型使用 cl100k_base 编码,具有 8191 个标记的上下文长度,我们可以看到超过这个限制会导致错误。

long_text = 'AGI ' * 5000
try:
get_embedding(long_text)
except openai.BadRequestError as e:
print(e)

Error code: 400 - {'error': {'message': "This model's maximum context length is 8192 tokens, however you requested 10001 tokens (10001 in your prompt; 0 for the completion). Please reduce your prompt; or completion length.", 'type': 'invalid_request_error', 'param': None, 'code': None}}

显然,我们希望避免这些错误,特别是在处理大量嵌入时。然而,我们仍然可能会遇到比最大上下文长度更长的文本。下面我们将描述并提供处理这些较长文本的主要方法:(1) 简单地将文本截断到允许的最大长度,以及(2) 将文本分块,并分别对每个块进行嵌入。

1. 截断输入文本

最简单的解决方案是将输入文本截断到允许的最大长度。由于上下文长度是以标记(tokens)来衡量的,所以我们必须在截断之前先对文本进行标记化。API接受文本或标记形式的输入,只要确保你使用了适当的编码,就无需将标记转换回字符串形式。下面是这样一个截断函数的示例。

import tiktoken

def truncate_text_tokens(text, encoding_name=EMBEDDING_ENCODING, max_tokens=EMBEDDING_CTX_LENGTH):
"""根据给定的编码,将字符串截断至 `max_tokens` 的长度。"""
encoding = tiktoken.get_encoding(encoding_name)
return encoding.encode(text)[:max_tokens]

我们之前的示例现在可以正常运行了。

truncated = truncate_text_tokens(long_text)
len(get_embedding(truncated))

1536

2. 将输入文本分块

尽管截断可以起到作用,但丢弃可能相关的文本是一个明显的缺点。另一种方法是将输入文本分成块,然后分别嵌入每个块。然后,我们可以单独使用块嵌入,或者以某种方式将它们组合在一起,比如平均值(按每个块的大小加权)。

我们将使用来自Python自己的食谱中的一个函数,该函数将一个序列分解成块。

from itertools import islice

def batched(iterable, n):
"""将数据分批处理成每批长度为 n 的元组。最后一批的长度可能较短。"""
# batched('ABCDEFG', 3) --> ABC DEF G
if n < 1:
raise ValueError('n must be at least one')
it = iter(iterable)
while (batch := tuple(islice(it, n))):
yield batch

现在我们定义一个函数,将一个字符串编码为标记,然后将其分成块。

def chunked_tokens(text, encoding_name, chunk_length):
encoding = tiktoken.get_encoding(encoding_name)
tokens = encoding.encode(text)
chunks_iterator = batched(tokens, chunk_length)
yield from chunks_iterator

最后,我们可以编写一个函数,安全地处理嵌入请求,即使输入文本超过最大上下文长度,也可以通过对输入标记进行分块处理,并分别嵌入每个分块。average 标志可以设置为 True,以返回分块嵌入的加权平均值,或设置为 False,以简单地返回未修改的分块嵌入列表。

import numpy as np


def len_safe_get_embedding(text, model=EMBEDDING_MODEL, max_tokens=EMBEDDING_CTX_LENGTH, encoding_name=EMBEDDING_ENCODING, average=True):
chunk_embeddings = []
chunk_lens = []
for chunk in chunked_tokens(text, encoding_name=encoding_name, chunk_length=max_tokens):
chunk_embeddings.append(get_embedding(chunk, model=model))
chunk_lens.append(len(chunk))

if average:
chunk_embeddings = np.average(chunk_embeddings, axis=0, weights=chunk_lens)
chunk_embeddings = chunk_embeddings / np.linalg.norm(chunk_embeddings) # 将长度归一化至1
chunk_embeddings = chunk_embeddings.tolist()
return chunk_embeddings

我们现在可以处理长输入文本。

average_embedding_vector = len_safe_get_embedding(long_text, average=True)
chunks_embedding_vectors = len_safe_get_embedding(long_text, average=False)

print(f"Setting average=True gives us a single {len(average_embedding_vector)}-dimensional embedding vector for our long text.")
print(f"Setting average=False gives us {len(chunks_embedding_vectors)} embedding vectors, one for each of the chunks.")


Setting average=True gives us a single 1536-dimensional embedding vector for our long text.
Setting average=False gives us 2 embedding vectors, one for each of the chunks.

在某些情况下,根据段落边界或句子边界来拆分块可能有助于保留文本的含义。