写点什么

三种高级 RAG 检索方法帮企业告别冗长文档!

作者:神州数码
  • 2024-07-03
    北京
  • 本文字数:5007 字

    阅读完需:约 16 分钟

前言


RAG(Retrieval-Augmented Generation)检索增强生成,是现如今基于企业私域知识的问答应用所使用的主流技术之一。相较于重新训练基于私域知识的大模型来说,RAG 没有额外的预训练成本,且回答效果与之相当。


但在实际应用场景中,RAG 所面临最大的问题是 LLM 的上下文长度限制。企业私域知识文本的数量十分庞大,不可能将其全部放在模型的 prompt 中,即使现在各类模型已经将上下文 token 从年初的 2k、4k 扩充到了 128k、192k,但是这可能也就是一份合同、一份标书的长度。因此,如何减少传递给模型的内容数量且同时提高内容质量,是提升基于 RAG 的 AI 应用回答准确度的一个重要方法。


本篇文章将基于 LangChain 实现三种高级检索方法,句子窗口检索和自动合并检索旨在改善 RAG 流程的召回过程中存在的信息残缺的问题,而多路召回检索则保证了在多个文档中检索召回的准确性。


先验知识


● RAG 简要流程


图片


加载文档——切分划片——嵌入为向量表示——存入数据库



向量化问题——向量召回文档——合并放入 Prompt——LLM 生成答案


句子窗口检索


一、概念


在文档进行切片工作后,文档被分为若干个 Langchain 自定义的 Document 对象,该对象有两个属性,一是 page_content 即该切片的文本内容,二是 meta_data 即有关该切片的一些信息和可自定义封装进入的信息。


句子窗口检索方法,将每个切片的相邻切片的内容封装在切片的 meta_data 中。在检索和召回过程时,根据命中文档的 meta_data 可获得此段落的上下文信息,并将其封装进入命中文档的 page_content 中。组合完成的文档列表即可作为 prompt 交付给大模型生成。


在实际问答任务中,我们建议使用切片器将文档切分为较短的分片,或使用依据标点符号进行切分的切片器。保证整片文档拥有较细的颗粒度。同时在封装和召回阶段,适当扩大窗口大小,保证召回段落的完整性。


二、BERT


(1)元数据封装

def metadata_format(self, ordered_text, **kwargs):        count = kwargs.get("split_count", 1)        for i, document in enumerate(ordered_text):            if i > 0:                document.metadata['previous_page'] = ordered_text[i-count].page_content            else:                document.metadata['previous_page'] = ''
if i < len(ordered_text) - 1: document.metadata['next_page'] = ordered_text[i+count].page_content else: document.metadata['next_page'] = '' return ordered_text
复制代码


(2)数据重构

def search_and_format(self, databases, query, **kwargs):        top_documents = []        for db in databases:            top_documents.append(db.similarity_search_with_score(query))        docs = []        for doc, _ in top_documents:            doc.page_content = doc.metadata.get("previous_page") + doc.page_content + doc.metadata.get("next_page")            docs.append(doc)        return docs
复制代码


(3)调用示例伪代码

# load document......
# split......# use smartvision sdk to formatsentence_window_retrival = SentenceWindow()formatted_documents = sentence_window_retrival.metadata_format(documents, split_count=2)
# embedding......
# load in local vector db......
# use smartvision sdk to do search and multiple recalldatabases = [db]query = "烟草专卖品的运输"top_documents = sentence_window_retrival.search_and_format(databases, query)print(top_documents)
复制代码



自动合并检索


一、概念


自动合并检索方法,实现方法源自 Llamaindex 所封装的自动合并检索,但 RAG 全流程需要制定一套准确的规范,因此在用户文档完成读取和切片工作后,所得到的 Langchain 格式的 Document 对象需转化为 Llamaindex 定义的 Document 对象,便可通过 Llamaindex 的自定义算法自动划分整个切片列表的子节点和父节点,最后鉴于规范再重新转化为 Langchain 格式的 Document 对象,并将父节点信息、深度信息等封装进每个节点。


在检索阶段,召回最相关的若干个节点,遍历这些节点和附加信息,如若超过 K 个节点同时属于同一个节点(这里的 K 为用户自定义阈值,通常为一个节点所有子节点的半数)则执行合并该父节点下属所有子节点,即返回整个父节点内容。这使我们能够将可能不同的较小上下文合并到一个可能有助于综合的更大上下文中。


二、代码实现和调用


(1)元数据封装


def auto_merge_format(documents, **kwargs):    if documents is None:        raise ValueError('documents is required')    formatted_documents = []    doc_text = "\n\n".join([d.page_content for d in documents])    docs = [Document(text=doc_text)]    node_parser = HierarchicalNodeParser.from_defaults(chunk_sizes=kwargs.get("pc_chunk_size", [2048, 512, 128]),chunk_overlap=kwargs.get("pc_chunk_overlap", 10))    nodes = node_parser.get_nodes_from_documents(docs)    leaf_nodes = get_leaf_nodes(nodes)    root_nodes = get_root_nodes(nodes)    middle_nodes = get_middle_node(nodes, leaf_nodes, root_nodes)    root_context_dict = {}    for root_node in nodes:        root_context_dict[root_node.node_id] = root_node.get_content()
for node in nodes: if node.parent_node: node_id = node.node_id root_node_id = node.parent_node.node_id root_node_content = root_context_dict.get(node.parent_node.node_id) root_node_child_count = 0 for parent_node in root_nodes + middle_nodes: if parent_node.node_id == node.parent_node.node_id: root_node_child_count = len(parent_node.child_nodes) break depth = 2 if node in middle_nodes else 3 child_count = len(node.child_nodes) if node.child_nodes is not None else 0 document = langchain.schema.Document(page_content=node.get_content(),metadata={"node_id": node_id, "root_node_id": root_node_id, "root_node_content": root_node_content, "root_node_child_count": root_node_child_count, "depth": depth, "child_count": child_count}) formatted_documents.append(document) return formatted_documents
复制代码

(2)数据重构

def search_and_format(self, databases, query, **kwargs):        top_documents = []        for db in databases:            top_document = db.similarity_search_with_score(query)            top_documents.append(top_document)        leaf_nodes = [doc for doc, _ in top_documents]        return do_merge(leaf_nodes, **kwargs)

def group_nodes_by_depth(nodes, depth): return [node for node in nodes if node.metadata.get("depth") == depth]
def process_group(nodes, threshold): grouped_by_root_id = {} for node in nodes: root_id = node.metadata.get("root_node_id") grouped_by_root_id.setdefault(root_id, []).append(node)
merge_context = [] for group in grouped_by_root_id.values(): node_count = len(group) child_count = group[0].metadata.get("root_node_child_count") if node_count / child_count >= threshold: merge_context.append(langchain.schema.Document( page_content=group[0].metadata.get("root_node_content") )) else: for document in group: merge_context.append(document) return merge_context
def do_merge(nodes, **kwargs) -> List[langchain.schema.Document]: threshold = kwargs.get("threshold", 0.5) leaf_nodes = group_nodes_by_depth(nodes, 3) middle_nodes = group_nodes_by_depth(nodes, 2) leaf_merge_context = process_group(leaf_nodes, threshold) middle_merge_context = process_group(middle_nodes, threshold) merge_content = leaf_merge_context + middle_merge_context return merge_content
def get_middle_node(nodes, leaf_nodes, root_nodes): middle_node = [] for node in nodes: if node not in leaf_nodes and node not in root_nodes: middle_node.append(node) return middle_node
复制代码

(3)调用示例伪代码

# load document......
#split......
# use smartvision sdk to formatauto_merge_retrival = AutoMergeRetrieval()formatted_documents = auto_merge_retrival.metadata_format(documents, pc_chunk_size=[1024, 128, 32], pc_chunk_overlap=4)#embedding......
#load in local vector db......
# use smartvision sdk to do search and multiple recalltop_documents = auto_merge_retrival.search_and_format(databases, query, threshold=0.5)print(top_documents)
复制代码



多路召回检索


一、概念


多路召回检索方法,在元数据封装环节并未做任何操作,而在检索阶段他允许用户上传多个数据集或不同类型的向量数据库作为检索对象,以适应用户私域知识库文档类型不同,文档数量庞大的问题。从多个数据源检索得到文档列表,而后通过 rerank 模型对文档与问题的相关性进行评分,筛选出大于一定分值的文档,组合成为 prompt。


由此可见,多路召回检索在数据源广而杂的情况下,富有更好的效果。此外,rerank 模型虽能进行再次的重排以提高准确性,但是在牺牲速度和效率的前提下进行的,因此需充分考虑这个问题。


二、代码实现


(1)元数据封装


def metadata_format(self, ordered_text, **kwargs):        """            默认rag,不做任何处理            """        return ordered_text
复制代码


(2)数据重构

def search_and_format(self, databases, query, **kwargs):        top_documents = []        result_data = []        for db in databases:            top_document = db.similarity_search_with_score(query)            top_documents.append(top_document)        pairs = [[query, item.page_content] for item in top_documents]        with torch.no_grad():            rerank_tokenizer = AutoTokenizer.from_pretrained(RERANK_FILE_PATH)            inputs = rerank_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)            rerank_model = AutoModelForSequenceClassification.from_pretrained(RERANK_FILE_PATH)            scores = rerank_model(**inputs, return_dict=True).logits.view(-1, ).float()            for i, score in enumerate(scores):                data = {                    "text": top_documents[i].page_content,                    "score": float(score)                }                result_data.append(data)        return result_data
复制代码


结语


本文提供的三种高级 RAG 检索方法,但仅改善了流程中检索召回环节的信息残缺问题,实质上 RAG 全流程均存在各种优化方法,但最有效的方法仍是改进或提供新的召回方式。


总结以上三种方法,均需要重点注意切片器的选用并控制切片大小,过大导致上下文长度过长,且有研究表明过长的 prompt 易使大模型忽略的中间部分的信息。过短则导致关键信息残缺,无法为大模型提供有效的上下文。因此开发者需根据文档类型和结构,谨慎选择并适当调节优化。


神州数码集团的神州问学平台不仅提供了本文所述的三种高级检索方法的 SDK,而且我们的开发团队正不断探索和研发新的、更高效的检索技术。我们致力于满足客户对于多样化私域知识库结构的需求,以实现更精准、更全面的搜索体验。同时,我们也欢迎您体验平台并提供宝贵意见。

用户头像

神州数码

关注

还未添加个人签名 2024-05-27 加入

还未添加个人简介

评论

发布
暂无评论
三种高级RAG检索方法帮企业告别冗长文档!_神州数码_InfoQ写作社区