从零开始解析RAG（四）：索引优化——多表示、分层与嵌入

大家好，这是我们 RAG 系列文档的第四篇，聚焦于除分块外，构建索引阶段的其他优化技巧。

Multi-Representation Indexing 多表示索引

基本 RAG 流程中，嵌入阶段是用同一个嵌入模型将 query 和文档都编码到同一个语义空间。这里的重点在于使用了同一模型去对 query 和 document 去做嵌入。

记得我们之前在 HyDE 中提过，query 和文档是两种不同的文本对象，那么能否从文档生成与 query 更为相近的文本对象，将 query 和生成的文本嵌入到同一语义空间进行检索呢？这就是接下来要介绍的多表示索引。

最初的思想来自于一篇叫做 Proposition Indexing 的文章，它的核心思想是将原始文档与最终用于检索的文档单元进行解耦。Proposition Indexing 提出，针对文档产生命题 proposition。之后，使用命题进行嵌入，用于后续检索。

还记得 从零开始解析 RAG（一）中的 HyDE 吗？Proposition Indexing 与 HyDE 方法像是一个维度上两个方向的延伸，HyDE 提出原始 query 和知识库中的 document 是两种不同的文本对象，因此使用 query 生成假设文档，使用假设文档进行检索，其核心思想可以简化为 query 向文档的靠近。 Proposition Indexing 则是针对文档产生命题（也可以看做就是对文档进行摘要，让其变得精炼），最终使用命题进行检索。可以看做是 document 向 query 的靠近。

原理解析

在 Proposition Indexing 的基础上提出了多表示索引 Multi-Representation Indexing，其本质是为同一文档对象建立多种索引表示形式，通过摘要进行检索，將原始文档放入 LLM 的上下文中。其基本流程如下所示：

1. 从文档中产生命题：接收文档，在实际上提炼它，从文档中创建命题（在实现中，通常是产生文档摘要）。
2. 对文档摘要进行嵌入：摘要可能更适用于检索，摘要可能包括文档的许多关键字，或者文档的大意。因此，将嵌入摘要进行检索时，可能会找到更适合的文档。
3. 使用原始文档作为上下文：这里需要注意的是，虽然使用摘要进行检索，但是最终要使用原始的文档作为LLM的上下文。 这类方法适用于长下文的LLM中。

代码实现

首先，从外部加载两个文档，一篇关于 agent, 另一篇关于人类数据质量的。

from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
docs = loader.load()

loader = WebBaseLoader("https://lilianweng.github.io/posts/2024-02-05-human-data-quality/")
docs.extend(loader.load())

然后，对完整的文档进行摘要，这一步通过 prompt LLM 完成。

import uuid

from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

chain = (
    {"doc": lambda x: x.page_content}
    | ChatPromptTemplate.from_template("Summarize the following document:\n\n{doc}")
    | ChatOpenAI(model="gpt-3.5-turbo",max_retries=0)
    | StrOutputParser()
)

summaries = chain.batch(docs, {"max_concurrency": 5})

接下来，建立两重索引

• 使用向量数据库 Chroma 存储摘要的嵌入
• InMemoryByteStore 用于存储原始文档的
• 为两层索引建立链接，这里使用多向量检索器 MultiVectorRetriver ，doc_id 摘要与原始文档联系起来。

这里的关键是理解 MultiVectorRetriever 的工作原理：它通过摘要（子块）进行向量检索，再通过 doc_id 找到对应的原始文档（父文档）。

from langchain.storage import InMemoryByteStore
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.retrievers.multi_vector import MultiVectorRetriever

# The vectorstore to use to index the child chunks
vectorstore = Chroma(collection_name="summaries",
                     embedding_function=OpenAIEmbeddings())

# The storage layer for the parent documents
store = InMemoryByteStore()
id_key = "doc_id"

# The retriever
# 将向量数据库和原始文档数据的存储使用 doc_id 进行链接
retriever = MultiVectorRetriever(
    vectorstore=vectorstore,
    byte_store=store,
    id_key=id_key,
)
# 为每个文档生成唯一的ID
doc_ids = [str(uuid.uuid4()) for _ in docs]

# Docs linked to summaries
# 生成的摘要文档与对应的doc_id 关联
summary_docs = [
    Document(page_content=s, metadata={id_key: doc_ids[i]})
    for i, s in enumerate(summaries)
]

# Add
# 将 summary_docs 添加到vectorstore 中，用于相似度搜索
retriever.vectorstore.add_documents(summary_docs)
# 将原始文档存储docstore 中
retriever.docstore.mset(list(zip(doc_ids, docs)))

接下来，使用摘要进行语义搜索

query = "Memory in agents"
sub_docs = vectorstore.similarity_search(query,k=1)
sub_docs[0] 
# Document(page_content='The document discusses the concept of building autonomous agents powered by Large Language Models (LLMs) as their core controllers. It covers components such as planning, memory, and tool use, along with case studies and proof-of-concept examples like AutoGPT and GPT-Engineer. Challenges like finite context length, planning difficulties, and reliability of natural language interfaces are also highlighted. The document provides references to related research papers and offers a comprehensive overview of LLM-powered autonomous agents.', metadata={'doc_id': 'cf31524b-fe6a-4b28-a980-f5687c9460ea'})
# 这里检索出来的文档是与angent 相关，通过查看sub_doc[0] 的内容可知，用于检索的是摘要

最后，使用 get_relevant_documents 方法检索相关文档。 查看一下检索到文档的内容，是原始文档，而不是摘要。

retrieved_docs = retriever.get_relevant_documents(query,n_results=1)
retrieved_docs[0].page_content[0:500]
# 使用检索器的get_relevant_documents方法检索到的是完整的文档
# "\n\n\n\n\n\nLLM Powered Autonomous Agents | Lil'Log\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\nLil'Log\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\nPosts\n\n\n\n\nArchive\n\n\n\n\nSearch\n\n\n\n\nTags\n\n\n\n\nFAQ\n\n\n\n\nemojisearch.app\n\n\n\n\n\n\n\n\n\n      LLM Powered Autonomous Agents\n    \nDate: June 23, 2023  |  Estimated Reading Time: 31 min  |  Author: Lilian Weng\n\n\n \n\n\nTable of Contents\n\n\n\nAgent System Overview\n\nComponent One: Planning\n\nTask Decomposition\n\nSelf-Reflection\n\n\nComponent Two: Memory\n\nTypes of Memory\n\nMaximum Inner Product Search (MIPS)\n\n"

更多内容可参考：

• arxiv.org/abs/2312.06…
• blog.langchain.dev/semi-struct…
• python.langchain.com/docs/module…

Raptor

原理解析

将 RAG 中 query 分为两类：

• 需要几个文档或者几个文档 chunk 辅助才能回答，称之为低级问题。
• 有一些问题需要大范围整合不同的文档才能回答，称之为高级问题。

这就为检索过程带来了问题，在基本的 RAG 流程中，检索阶段通常使用 KNN 检索，也就是说对于每个问题，从文档中检索 K 个文档块用于回答。对于需要整合多个文档才能回答的高级问题，它需要的文档数可能超过检索过程中设置的参数 K。举个例子，假设有一个问题需要在 5 个文档中整合信息才能回答，但是 K=3，遇到这种情况怎么办？

Raptor 的出现就是为了解决这种情况，其核心思想是针对文档摘要创建分层索引。具体来说

1. 将原始文档称为叶子，对这些叶子节点进行嵌入并聚类。这样依赖，相似的文档会放在一个聚类中，这个聚类中，可能包含不同的 document，不同的文本 chunk。
2. 对每个聚类进行总结，生成聚类摘要。每个聚类中可能包含来自同一个文档不同的 chunks, 甚至是来自不同文档的 chunks。这个聚类的过程本质上是在捕获相似的信息，可以在聚类的摘要中进行查阅。
3. 在实践中，通过不断地重复聚类，摘要这两步，可以创建更多层次的索引，形成更加抽象的树的结构。 对于高级问题，它们在语义搜索上与聚类摘要应该更相似。对于低级问题，直接检索低级的文档块就可以。Raptor 提出的分层索引，可以更好地覆盖不同类型的问题。

原始论文

这里补充一些原始论文的细节。 RAPTOR 的英文全称是 Recursive Abstractive Processing for Tree-Organized Retrieval，将通过递归抽象化构建一个树结构的索引，用于解决检索增强语言模型在表示和利用大规模话语结构方面的不足。具体来说，

1. ​文本分段和嵌入：首先，将检索语料库分割成短的连续文本片段，每个片段长度为 100 个词。如果一个句子超过 100 个词的限制，则将整个句子移动到下一个片段，以保持文本的连贯性。然后，使用 SBERT（Sentence-BERT）对这些文本片段进行嵌入。

2. ​聚类算法：接下来，使用高斯混合模型（GMM）对这些文本片段进行聚类。GMM 假设数据点是由多个高斯分布生成的，公式如下：

其中，$x$ 表示文本向量，$k$ 表示高斯分布的索引，$\mu_k$ 和 $\boldsymbol{\Sigma}_k$ 分别表示均值和协方差矩阵。为了应对高维向量的挑战，使用了均匀流形近似和投影（UMAP）技术。

3. ​文本摘要：聚类后，使用 GPT-3.5-turbo 对每个聚类的文本进行摘要。这一步骤将大量检索到的信息压缩成可管理的摘要。

4. ​树结构构建：重复嵌入、聚类和摘要的过程，直到进一步聚类变得不可行，从而构建出一个多层次的树结构。树的节点根据语义相似性分组，而不仅仅是文本顺序。

5. 查询策略：引入了两种查询机制：树遍历和折叠树。

   • 树遍历方法逐层遍历树，选择最相关的节点；
   • 折叠树方法将所有节点折叠成单层，并根据余弦相似度选择最相关的节点。

代码实现

其代码实现在这里，由于作者的代码实现有详细的注释，在这里不过多赘述。 更多内容可参考

• www.youtube.com/watch?v=jbG…
• github.com/langchain-a…

ColBERT

原理解析

我们已经讨论了很多关于嵌入对于语义相似性和检索非常重要，在 RAG 的基本流程中，通常是将文档进行嵌入，转换为向量，这个向量需要表示该文档的所有语义。因此，引发一个文档，将文档表示单个向量是不是压缩了太多的信息呢？可不可想一种方法不要压缩太多？

ColBERT 就是这样这一种方法。不是直接使用文档进行嵌入，而是针对 token 进行嵌入，也就是说对于原始的 query 和 document 不直接生成嵌入，而是将其先进行分词之后在进行嵌入，然后再进行检索。 具体来说，

1. 独立编码：首先，ColBERT 使用 BERT 独立地对查询和文档进行编码，生成两组上下文化嵌入。查询嵌入记为 Eq​，文档嵌入记为 Ed​。
2. ​延迟交互：然后，ColBERT 通过最大相似度（MaxSim）操作来评估查询和文档之间的相关性。具体来说，对于每个查询嵌入 v∈Eq​，找到其在文档嵌入 Ed​ 中的最大余弦相似度，并将这些相似度相加。公式如下： $S_{q,d}:=\sum_{i\in[|E_q|]}\max_{j\in[|E_d|]}E_{q_i}\cdot E_{d_j}^T$ 其中，$E_{q_i}$ ​​ 和 $E_{d_j​​}$分别是查询嵌入和文档嵌入的第 i 个和第 j 个元素。
3. ​离线索引：为了进一步提高效率，ColBERT 在离线阶段预先计算并存储文档嵌入。使用批量处理和 GPU 加速来提高索引速度。
4. ​端到端检索：最后，ColBERT 利用 faiss 库进行大规模向量相似性搜索，直接从大型文档集合中检索前 k 个结果。具体步骤包括在离线索引阶段构建倒排索引，在在线检索阶段使用近似最近邻搜索和精细化重排序。

这个方法的特点：

• 使用 token 进行 embedding，代替原有的文档进行嵌入
• 使用不同的算法计算文档相似度
• 缺点是延迟太高了

代码实现

ragatouille 有现成的模型实现。 首先，安装库。

! pip install -U ragatouille

载入模型

from ragatouille import RAGPretrainedModel
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

加载一个网页作为要构建索引的文档，这里通过维基百科官方 API 获取“宫崎骏”的英文百科页面，并解析返回主要内容。

import requests

def get_wikipedia_page(title: str):
    """
    Retrieve the full text content of a Wikipedia page.

    :param title: str - Title of the Wikipedia page.
    :return: str - Full text content of the page as raw string.
    """
    # Wikipedia API endpoint
    URL = "https://en.wikipedia.org/w/api.php"

    # Parameters for the API request
    params = {
        "action": "query",
        "format": "json",
        "titles": title,
        "prop": "extracts",
        "explaintext": True,
    }

    # Custom User-Agent header to comply with Wikipedia's best practices
    headers = {"User-Agent": "RAGatouille_tutorial/0.0.1 (ben@clavie.eu)"}

    response = requests.get(URL, params=params, headers=headers)
    data = response.json()

    # Extracting page content
    page = next(iter(data["query"]["pages"].values()))
    return page["extract"] if "extract" in page else None

full_document = get_wikipedia_page("Hayao_Miyazaki")
full_document
"""
'Hayao Miyazaki (宮崎 駿 or 宮﨑 駿, Miyazaki Hayao, [mijaꜜzaki hajao]; born January 5, 1941) is a Japanese animator, filmmaker, and manga artist. ....
"""

构建索引

RAG.index(
    collection=[full_document],
    index_name="Miyazaki-123",
    max_document_length=180,
    split_documents=True,
)

创建检索器，并检索相关文档

retriever = RAG.as_langchain_retriever(k=3)
retriever.invoke("What animation studio did Miyazaki found?")

更多内容请查看：

• arxiv.org/abs/2004.12…
• hackernoon.com/how-colbert…
• python.langchain.com/docs/integr…
• til.simonwillison.net/llms/colber….

总结

本文探讨索引阶段的优化技巧，主要聚焦于 3 种方法：对同一文档生成多种表示的形式的多表示索引，构建分层索引的 Raptor 以及基于 token 嵌入的ColBERT。下一篇文章将探讨检索阶段的优化技巧，敬请期待。