一种基于滑动窗口扩展上下文的RAG（检索增强生成）优化实现方案实践

RAG（检索增强生成）是一种结合了检索（通常是知识库或数据库）和生成模型（大语言模型）的技术，目的是在生成文本的时候能够参考相关的外部知识。这样，即使生成模型在训练时没有看到某些信息，它也能在生成时通过检索到的知识来生成更加准确和丰富的回答，这篇文章实现一种基于动态上下文窗口的方案，能够处理大规模文档，保留重要的上下文信息，提升检索效率，同时保持灵活性和可配置性。

](https://link.juejin.cn/?target=https%3A%2F%2Fu.jd.com%2FV8pkqFY)

先看效果

演示效果

整体方案

整体方案在于文档预处理阶段实现满足上下文窗口的原始文本分块，文档检索阶段实现文本的三次检索，下面逐一进行说明。测试文章来自大语言模型的安全问题探究-提示词攻击

文档预处理过程

小文本块拆分

以 50 token大小（可根据自身文档之间的组织规律动态调整粒度）对文本做首次分割

ini复制代码# 小文档块大小
BASE_CHUNK_SIZE = 50
# 小块的重叠部分大小
CHUNK_OVERLAP = 0
def split_doc(
    doc: List[Document], chunk_size=BASE_CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP, chunk_idx_name: str
):
    data_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        # 使用了 tiktoken 来确保分割不会在一个 token 的中间发生
        length_function=tiktoken_len,
    )
    doc_split = data_splitter.split_documents(doc)
    chunk_idx = 0
    for d_split in doc_split:
        d_split.metadata[chunk_idx_name] = chunk_idx
        chunk_idx += 1
    return doc_split

下面示例显示了前 7 个分块信息，结果如下：

css

复制代码[Document(page_content='LLM 安全专题 提示词', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 0}),Document(page_content='是指在训练或与⼤型语⾔模型（Claude，ChatGPT等）进⾏交互时，提供给模', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 1}),Document(page_content='型的输⼊⽂本。通过给定特定的', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 2}),Document(page_content='提示词，可以引导模型⽣成特定主题或类型的⽂本。在⾃然语⾔处理（NLP）任务中，提', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 3}),Document(page_content='示词充当了问题或输⼊的⻆⾊，⽽模型的输出是对这个问题的回答或完成的任务。关于', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 4}),Document(page_content='怎样设计好的', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 5}),Document(page_content='Prompt，查看Prompt专题章节内容就可以了，我不在这⾥过多阐述，个⼈⽐较感兴趣针对', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 6}),...]

添加窗口

以步长为 3，窗口大小为 6，将上述步骤的小块匹配到不同的上下文窗口。

ini复制代码# 步长定义了窗口移动的速度，具体来说，它是上一个窗口中第一个块和下一个窗口中第一个块之间的距离
WINDOW_STEPS = 3
# 窗口大小直接影响到每个窗口中的上下文信息量，窗口大小= BASE_CHUNK_SIZE * WINDOW_SCALE
WINDOW_SCALE = 6
def add_window(
    doc: Document, window_steps=WINDOW_STEPS, window_size=WINDOW_SCALE, window_idx_name: str
):
    window_id = 0
    window_deque = deque()

    for idx, item in enumerate(doc):
        if idx % window_steps == 0 and idx != 0 and idx < len(doc) - window_size:
            window_id += 1
        window_deque.append(window_id)

        if len(window_deque) > window_size:
            for _ in range(window_steps):
                window_deque.popleft()

        window = set(window_deque)
        item.metadata[f"{window_idx_name}_lower_bound"] = min(window)
        item.metadata[f"{window_idx_name}_upper_bound"] = max(window)

下面示例显示了前 7 个增加窗口信息后的分块内容，结果如下

css

复制代码[Document(page_content='LLM 安全专题 提示词', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 0, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 0}),Document(page_content='是指在训练或与⼤型语⾔模型（Claude，ChatGPT等）进⾏交互时，提供给模', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 1, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 0}),Document(page_content='型的输⼊⽂本。通过给定特定的', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 2, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 0}),Document(page_content='提示词，可以引导模型⽣成特定主题或类型的⽂本。在⾃然语⾔处理（NLP）任务中，提', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 3, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 1}),Document(page_content='示词充当了问题或输⼊的⻆⾊，⽽模型的输出是对这个问题的回答或完成的任务。关于', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 4, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 1}),Document(page_content='怎样设计好的', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 5, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 1}),Document(page_content='Prompt，查看Prompt专题章节内容就可以了，我不在这⾥过多阐述，个⼈⽐较感兴趣针对', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 6, 'large_chunks_idx_lower_bound': 1, 'large_chunks_idx_upper_bound': 2}),Document(page_content='Prompt的攻击，随着⼤语⾔模型的⼴泛应⽤，安全必定是⼀个⾮常值', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'small_chunk_idx': 7, 'large_chunks_idx_lower_bound': 1, 'large_chunks_idx_upper_bound': 2}),...]

中等文本块

以小文本块 3 倍（可动态配置），即150 token大小对文本做二次分割，形成中等文本块

ini复制代码# 中等大小的文档块大小=BASE_CHUNK_SIZE * CHUNK_SCALE
CHUNK_SCALE = 3
def merge_metadata(dicts_list: dict):
    merged_dict = {}
    bounds_dict = {}
    keys_to_remove = set()

    for dic in dicts_list:
        for key, value in dic.items():
            if key in merged_dict:
                if value not in merged_dict[key]:
                    merged_dict[key].append(value)
            else:
                merged_dict[key] = [value]

    for key, values in merged_dict.items():
        if len(values) > 1 and all(isinstance(x, (int, float)) for x in values):
            bounds_dict[f"{key}_lower_bound"] = min(values)
            bounds_dict[f"{key}_upper_bound"] = max(values)
            keys_to_remove.add(key)

    merged_dict.update(bounds_dict)

    for key in keys_to_remove:
        del merged_dict[key]

    return {
        k: v[0] if isinstance(v, list) and len(v) == 1 else v
        for k, v in merged_dict.items()
    }

def merge_chunks(doc: Document, scale_factor=CHUNK_SCALE, chunk_idx_name: str):
    merged_doc = []
    page_content = ""
    metadata_list = []
    chunk_idx = 0

    for idx, item in enumerate(doc):
        page_content += item.page_content
        metadata_list.append(item.metadata)
        if (idx + 1) % scale_factor == 0 or idx == len(doc) - 1:
            metadata = merge_metadata(metadata_list)
            metadata[chunk_idx_name] = chunk_idx
            merged_doc.append(
                Document(
                    page_content=page_content,
                    metadata=metadata,
                )
            )
            chunk_idx += 1
            page_content = ""
            metadata_list = []

    return merged_doc

下面示例显示了前 3 个中等分块信息，结果如下：

css

复制代码[Document(page_content='LLM 安全专题 提示词是指在训练或与⼤型语⾔模型（Claude，ChatGPT等）进⾏交互时，提供给模型的输⼊⽂本。通过给定特定的', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 0, 'small_chunk_idx_lower_bound': 0, 'small_chunk_idx_upper_bound': 2, 'medium_chunk_idx': 0}),Document(page_content='提示词，可以引导模型⽣成特定主题或类型的⽂本。在⾃然语⾔处理（NLP）任务中，提示词充当了问题或输⼊的⻆⾊，⽽模型的输出是对这个问题的回答或完成的任务。关于怎样设计好的', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'large_chunks_idx_lower_bound': 0, 'large_chunks_idx_upper_bound': 1, 'small_chunk_idx_lower_bound': 3, 'small_chunk_idx_upper_bound': 5, 'medium_chunk_idx': 1}),Document(page_content='Prompt，查看Prompt专题章节内容就可以了，我不在这⾥过多阐述，个⼈⽐较感兴趣针对Prompt的攻击，随着⼤语⾔模型的⼴泛应⽤，安全必定是⼀个⾮常值得关注的领域。提示攻击', metadata={'source': './data/一文带你了解提示攻击.pdf', 'page': 0, 'large_chunks_idx_lower_bound': 1, 'large_chunks_idx_upper_bound': 2, 'small_chunk_idx_lower_bound': 6, 'small_chunk_idx_upper_bound': 8, 'medium_chunk_idx': 2}),...]

文档检索过程

检索器声明

首先声明一个检索器，用于检索文档，这里将 BM25 检索器和嵌入式检索器组合成一个集成检索器，用于检索和评估文档相似度。下面是一些需要相关知识：

BM25 是一种基于词袋模型的检索方法，它通过考虑单词在文档中的频率和在整个文档集合中的逆文档频率来计算文档之间的相似度
嵌入式检索器通常使用预训练的嵌入模型（本案例使用 OpenAI 的 text-embedding-ada-002 模型）将文档转换为密集向量，然后通过计算这些向量之间的相似度来评估文档之间的相似性
emb_filter: 用于在嵌入式检索过程中过滤结果。例如，可以根据某些标准排除不相关的文档
k: 这是一个整数，表示要返回的最匹配的前几个结果数量

weights: 包含两个权重值，分别用于 BM25 检索器和嵌入式检索器在集成检索中的权重。

ini复制代码def get_retriever(
    self,
    docs_chunks,
    emb_chunks,
    emb_filter=None,
    k=2,
    weights=(0.5, 0.5),
):
    bm25_retriever = BM25Retriever.from_documents(docs_chunks)
    bm25_retriever.k = k

    emb_retriever = emb_chunks.as_retriever(
        search_kwargs={
            "filter": emb_filter,
            "k": k,
            "search_type": "mmr",
        }
    )
    return MyEnsembleRetriever(
        retrievers={"bm25": bm25_retriever, "chroma": emb_retriever},
        weights=weights,
    )

检索相关文档

文档检索通过采用多阶段（三次）的方式进行

第一阶段：小分块检索 使用小文档块(docs_index_small)和小嵌入块(embedding_chunks_small)初始化一个检索器(first_retriever)，使用这个检索器检索与查询相关的文档，并将结果存储在 first 变量中，对检索到的文档 ID 进行清理和过滤，确保它们是相关的，并存储在 ids_clean 变量中。
第二阶段：移动窗口检索 针对每个唯一的源文档，使用小文档块检索与该源文档相关的所有文档块。使用包含这些文档块的新检索器(second_retriever)，再次进行检索，以进一步缩小相关文档的范围，将检索到的文档添加到 docs 列表中。

第三阶段：中等分块检索 使用过滤条件从中等文档块(docs_index_medium)检索相关文档，使用包含这些文档块的新检索器(third_retriever)进行检索。从检索到的文档中选择前 third_num_k 个文档，存储在 third 变量中，清理文档的元数据，删除不需要的内容，将最终检索到的文档按文件名分类，并存储在 qa_chunks 字典中。

ini复制代码def get_relevant_documents(
    self,
    query: str,
    num_query: int,
    *,
    run_manager: Optional[CallbackManagerForChainRun] = None,
) -> List[Document]:
    # 第一次检索，小分块信息
    first_retriever = self.get_retriever(
        docs_chunks=self.docs_index_small.documents,
        emb_chunks=self.embedding_chunks_small,
        emb_filter=None,
        k=self.first_retrieval_k,
        weights=self.retriever_weights,
    )
    first = first_retriever.get_relevant_documents(
        query, callbacks=run_manager.get_child()
    )
    ids_clean = self.get_relevant_doc_ids(first, query)
    qa_chunks = {}
    if ids_clean and isinstance(ids_clean, list):
        source_md5_dict = {}
        for ids_c in ids_clean:
            if ids_c < len(first):
                if ids_c not in source_md5_dict:
                    source_md5_dict[first[ids_c].metadata["source_md5"]] = [
                        first[ids_c]
                    ]
        if len(source_md5_dict) == 0:
            source_md5_dict[first[0].metadata["source_md5"]] = [first[0]]
        num_docs = len(source_md5_dict.keys())
        third_num_k = max(
            1,
            (
                int(
                    (
                        MAX_LLM_CONTEXT
                        / (BASE_CHUNK_SIZE * CHUNK_SCALE)
                    )
                    // (num_docs * num_query)
                )
            ),
        )

        for source_md5, docs in source_md5_dict.items():
            second_docs_chunks = self.docs_index_small.retrieve_metadata(
                {
                    "source_md5": (IndexerOperator.EQ, source_md5),
                }
            )
            # 第二次检索
            second_retriever = self.get_retriever(
                docs_chunks=second_docs_chunks,
                emb_chunks=self.embedding_chunks_small,
                emb_filter={"source_md5": source_md5},
                k=self.second_retrieval_k,
                weights=self.retriever_weights,
            )
            second = second_retriever.get_relevant_documents(
                query, callbacks=run_manager.get_child()
            )
            docs.extend(second)
            docindexer_filter, chroma_filter = self.get_filter(
                self.num_windows, source_md5, docs
            )
            third_docs_chunks = self.docs_index_medium.retrieve_metadata(
                docindexer_filter
            )
            # 第三次检索
            third_retriever = self.get_retriever(
                docs_chunks=third_docs_chunks,
                emb_chunks=self.embedding_chunks_medium,
                emb_filter=chroma_filter,
                k=third_num_k,
                weights=self.retriever_weights,
            )
            third_temp = third_retriever.get_relevant_documents(
                query, callbacks=run_manager.get_child()
            )
            third = third_temp[:third_num_k]
            for doc in third:
                mtdata = doc.metadata
                mtdata["page_content"] = None
            file_name = third[0].metadata["source"].split("/")[-1]
            if file_name not in qa_chunks:
                qa_chunks[file_name] = third
            else:
                qa_chunks[file_name].extend(third)
    return qa_chunks

整个过程是一个分层的检索过程，首先在小文档块中进行粗略检索，然后在特定的源文档中进行更精确的检索，最后在中等文档块中进行最终的检索。这种分层的方法有助于提高检索的效率和准确性，因为它允许系统在更小的文档集上进行更精确的检索，从而减少了在大文档集上进行复杂检索所需的计算量。