赞
踩
本文首发于博客LLM 应用开发实践
RAG(检索增强生成)是一种结合了检索(通常是知识库或数据库)和生成模型(大语言模型)的技术,目的是在生成文本的时候能够参考相关的外部知识。这样,即使生成模型在训练时没有看到某些信息,它也能在生成时通过检索到的知识来生成更加准确和丰富的回答,这篇文章实现一种基于动态上下文窗口的方案,能够处理大规模文档,保留重要的上下文信息,提升检索效率,同时保持灵活性和可配置性。
整体方案在于文档预处理阶段实现满足上下文窗口的原始文本分块,文档检索阶段实现文本的三次检索,下面逐一进行说明。测试文章来自 大语言模型的安全问题探究-提示词攻击
以 50 token
大小(可根据自身文档之间的组织规律动态调整粒度)对文本做首次分割
# 小文档块大小 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 个分块信息,结果如下:
[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,将上述步骤的小块匹配到不同的上下文窗口。
# 步长定义了窗口移动的速度,具体来说,它是上一个窗口中第一个块和下一个窗口中第一个块之间的距离 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 个增加窗口信息后的分块内容,结果如下
[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
大小对文本做二次分割,形成中等文本块
# 中等大小的文档块大小=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 个中等分块信息,结果如下:
[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 检索器和嵌入式检索器组合成一个集成检索器,用于检索和评估文档相似度。下面是一些需要相关知识:
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, )
文档检索通过采用多阶段(三次)的方式进行
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
整个过程是一个分层的检索过程,首先在小文档块中进行粗略检索,然后在特定的源文档中进行更精确的检索,最后在中等文档块中进行最终的检索。这种分层的方法有助于提高检索的效率和准确性,因为它允许系统在更小的文档集上进行更精确的检索,从而减少了在大文档集上进行复杂检索所需的计算量。
由于知识库通常包含大量的文档,直接在这么大的文档集合上进行检索是非常耗时的。通过将文档分割成更小的块(chunk_small
),小块更容易被索引和检索。
通过为小块添加窗口信息(add_window
),可以确保在检索时不会丢失重要的上下文信息。这是因为有些信息可能分布在多个小块中,单独检索一个小块可能会遗漏这些信息,窗口机制确保在检索时考虑到足够的上下文,从而生成更准确的回答。
通过将相邻的小块合并成中等大小的块(chunk_medium
),在保留细粒度特性的同时增加了更大范围的上下文信息。这有助于提升检索的效率和准确性,因为中等大小的块既不会像大块那样导致检索效率下降,也不会像小块那样缺乏足够的上下文信息。
由于整个流程是模块化的,可以根据具体应用的需求灵活地配置每个步骤的参数(如块的大小、窗口的大小和步长等),以达到最佳的性能和效果平衡。
由于有了不同大小和包含窗口信息的文档块,可以根据查询的需要选择最适合的块来进行检索。对于需要广泛上下文的查询,可以使用包含更多上下文信息的中等大小或大块来进行检索,对于需要快速响应的查询,可以使用小块来提高检索速度。
整体方案在保证生成质量的同时,实现高效处理,在实际应用中效果明显。
因为检索阶段经过三次检索,加之使用本地矢量数据库使用 Chroma,整体的响应速率还需进一步改善。
回复 RAG获取完整源码。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。