RAG 的后置分块：基于 Jina AI 的实现

朴素分块及其在 RAG 中的局限

在 RAG 流水线中，文档会在嵌入并存储到向量数据库之前被拆分为更小的块。每个块独立处理，并在查询时用于检索。然而，这种“朴素分块”往往会丢失重要的长距离上下文。

问题在于，传统分块在切分文档时并未考虑信息之间的关联。例如，在一篇关于巴黎的文档中，“the city（这座城市）”可能与“Paris（巴黎）”被分到不同的块。缺乏完整上下文时，检索模型可能难以将这些指代对应起来，导致结果不够准确。在长文档中这一问题更为突出，因为关键信息往往分散在多个部分。

后置分块：在文档切分中保留上下文

后置分块通过改变切分时机来解决问题。它不是先把文档拆分成块，而是先使用长上下文模型对整篇文档进行嵌入，之后再进行切分。

后置分块的关键优势包括：

• 保留语境： 通过先对整篇文档进行嵌入，后置分块确保每个块都能保留整体语境。文本中的跨段落指代与关联因此能在块的嵌入中保持完整。
• 更优检索： 后置分块得到的块嵌入更丰富、更准确，模型对文档理解更好，从而提升 RAG 系统的检索效果。
• 应对长文本： 非常适合传统模型因 token 限制而无法一次性处理的超长文档。

借助如 Jina 的 jinaai/jina-embeddings-v2-base-en 等长上下文模型（支持最多 8192 个 token），后置分块可以在切分前高效地为大段文本生成嵌入。

实现后置分块

以下是使用 Jina 的长上下文嵌入模型实现后置分块的分步指南。您可以在 此处免费获取 Jina 的 API 密钥，我们将使用如下输入文本作为示例：

input_text = """Berlin is the capital and largest city of Germany, both by area and by population.
Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.
The city is also one of the states of Germany, and is the third smallest state in the country in terms of area."""技术支持

步骤 1：获取分块与区间标注

首先，使用您的 Jina API 密钥和下方辅助函数将输入文本拆分成块。这些块会带有区间标注，便于稍后基于文档嵌入进行切分。Jina 的 API 会利用段落或句子等自然边界，确保分块合理且语义完整。

import json
import requests

def custom_tokenize_jina_api(input_text: str):
    url = '<https://segment.jina.ai/>'
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ENTER_YOUR_JINA_API_KEY'
    }
    data = {
        "content": input_text,
        "tokenizer": "o200k_base",
        "return_tokens": "true",
        "return_chunks": "true",
        "max_chunk_length": "1000"
    }
    # Make the API request
    response = requests.post(url, headers=headers, json=data)
    response_data = response.json()
    chunks = response_data.get("chunks", [])
    i = 1
    j = 1
    span_annotations = []
    for x in response_data['tokens']:
        if j == 1:
            j = len(x)
        else:
            j = len(x) + i
        span_annotations.append((i, j))
        i = j
    return chunks, span_annotations
chunks, span_annotations = custom_tokenize_jina_api(input_text)

print(chunks)
print(span_annotations)技术支持

['Berlin is the capital and largest city of Germany, both by area and by population.\\n\\n', "Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.\\n\\n", 'The city is also one of the states of Germany, and is the third smallest state in the country in terms of area.']
[(1, 17), (17, 44), (44, 69)]技术支持

步骤 2：对文本进行分词并生成 token 级文档嵌入

先使用兼容长上下文模型的分词器（如 Jina 的 embeddings-v2-base-en）将整篇文档切分为 token。接着，使用长上下文 Transformer 模型为每个 token 生成嵌入。这意味着文档中的每个词或 token 都会拥有捕捉其语义的专属嵌入。

from transformers import AutoModel
from transformers import AutoTokenizer

# load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
inputs = tokenizer(input_text, return_tensors='pt')
model_output = model(**inputs)
model_output[0].shape技术支持

torch.Size([1, 71, 768]) # 71 represents number of tokens in the entire document技术支持

步骤 3：执行后置分块

当您获得整篇文档的 token 嵌入后，即可进行后置分块。使用第一步得到的区间标注，将这些 token 嵌入切分为较小的块。随后，对每个块内的嵌入进行均值池化（mean pooling），得到每个块的单一嵌入。至此，我们获得了兼具整篇文档强语境信息的块嵌入。

def late_chunking(
    model_output: 'BatchEncoding', span_annotation: list, max_length=None
):
    token_embeddings = model_output[0]
    outputs = []
    for embeddings, annotations in zip(token_embeddings, span_annotation):
        if (
            max_length is not None
        ):  # remove annotations which go bejond the max-length of the model
            annotations = [
                (start, min(end, max_length - 1))
                for (start, end) in annotations
                if start < (max_length - 1)
            ]
        pooled_embeddings = [
            embeddings[start:end].sum(dim=0) / (end - start)
            for start, end in annotations
            if (end - start) >= 1
        ]
        pooled_embeddings = [
            embedding.detach().cpu().numpy() for embedding in pooled_embeddings
        ]
        outputs.append(pooled_embeddings)
    return outputs技术支持

embeddings = late_chunking(model_output, [span_annotations])[0]
len(embeddings)技术支持

3 # matches number of chunks in Step 1技术支持

步骤 4：后置分块 vs. 传统分块的结果对比

为理解后置分块的优势，我们将其与传统分块进行比较：

embeddings_traditional_chunking = model.encode(chunks)技术支持

import numpy as np

cos_sim = lambda x, y: np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))
q = "Berlin"
berlin_embedding = model.encode(q)

print(q)
print('\n')
for chunk, new_embedding, trad_embeddings in zip(chunks, embeddings, embeddings_traditional_chunking):
  print(chunk.strip())
  print(f'Late chunking:', cos_sim(berlin_embedding, new_embedding))
  print(f'Traditional chunking:', cos_sim(berlin_embedding, trad_embeddings))
  print("------------------------------------------------------------------")技术支持

Berlin
Berlin is the capital and largest city of Germany, both by area and by population.
Late chunking: 0.84954596
Traditional chunking: 0.84862185
------------------------------------------------------------------
Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.
Late chunking: 0.82489026
Traditional chunking: 0.70843375
------------------------------------------------------------------
The city is also one of the states of Germany, and is the third smallest state in the country in terms of area.
Late chunking: 0.84980094
Traditional chunking: 0.7534553
------------------------------------------------------------------技术支持

如您在第二、第三个分块中所见，与“Berlin（柏林）”相比，传统分块的相似度约为 70–75%。而保持整篇文档语境的后置分块则提升至 82–84%。这表明后置分块在保留上下文、生成更有意义的嵌入方面效果更佳，从而带来更准确的检索结果。

结语

后置分块对文档检索系统，尤其是 RAG 流水线，是一次重要改进。通过在整篇文档嵌入完成后再进行切分，后置分块让每个块都能保留完整语境，从而得到更准确、更有意义的嵌入。