2023年1月18日

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

在Github中打开以Markdown格式查看

OpenAI的嵌入模型无法处理超过最大长度的文本。不同模型的最大长度各不相同，并以标记而非字符串长度来衡量。如果您不熟悉标记化处理，请查看How to count tokens with tiktoken。

本笔记本展示了如何处理超过模型最大上下文长度的文本。我们将使用text-embedding-3-small生成的嵌入向量进行演示，但同样的思路也适用于其他模型和任务。想了解更多关于嵌入向量的信息，请查看OpenAI Embeddings Guide。

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'

# let's make sure to not retry on an invalid request, because that is what we want to demonstrate
@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 模型的上下文长度为8191个token（使用cl100k_base编码），我们可以看到超过这个限制会导致错误。

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. 截断输入文本

最简单的解决方案是将输入文本截断至允许的最大长度。由于上下文长度是以token为单位计算的，我们必须在截断前先对文本进行token化处理。API既接受文本形式也接受token形式的输入，因此只要确保使用正确的编码方式，就无需将token转换回字符串形式。以下是一个此类截断函数的示例。

import tiktoken

def truncate_text_tokens(text, encoding_name=EMBEDDING_ENCODING, max_tokens=EMBEDDING_CTX_LENGTH):
    """Truncate a string to have `max_tokens` according to the given encoding."""
    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):
    """Batch data into tuples of length n. The last batch may be shorter."""
    # 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

现在我们定义一个函数，将字符串编码为令牌(token)然后将其分割成块。

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)  # normalizes length to 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.

在某些情况下，按段落边界或句子边界分割文本块有助于保留原文的语义。