通过上下文检索优化RAG的语境理解

无论你的模型（大型语言模型LLM）有多先进，如果上下文信息块没有提供正确的信息，模型将无法生成准确的答案。在本教程中，我们将探索一种称为上下文检索的技术，以提高你的RAG系统中上下文信息块的质量。

为了更好地理解，让我们从一个简单的例子开始。设想你手头有一份包含多个信息段的文档，你希望根据其中的某个信息段来提出问题。接下来，我们来看一个具体的信息段示例：

For more information, please refer to 

[the documentation of `vllm`](https://docs.vllm.ai/en/stable/).

 

Now, you can have fun with Qwen2.5 models.

这是一个能很好体现从其他上下文中获益的信息块示例。单独来看，这个信息块的信息含量相对有限。接下来，我们来看看增加了上下文信息后的信息块：

带上下文的示例数据块：

For more information, please refer to 

[the documentation of `vllm`](https://docs.vllm.ai/en/stable/).

 

Now, you can have fun with Qwen2.5 models.

The chunk is situated at the end of the document, following the section on 

deploying Qwen2.5 models with vLLM, and serves as a concluding remark 

encouraging users to explore the capabilities of Qwen2.5 models.

你可以想象，当模型接收到这个块时，它对上下文有了更好的理解，并且可以提供更准确的答案。让我们构建管道来创建这些块。

什么是上下文检索？

上下文检索（由 Anthropic 引入）解决了传统检索增强生成 （RAG） 系统中的一个常见问题：单个文本块通常缺乏足够的上下文来准确检索和理解。

上下文检索通过在嵌入或索引之前添加特定的解释性上下文来增强每个块。这保留了块与其更广泛的文档之间的关系，从而显著提高了系统检索和使用最相关信息的能力。

根据 Anthropic 的实验：

• 上下文嵌入将前 20 个块检索失败率降低了 35%。
• 将上下文嵌入与上下文 BM25 相结合，进一步降低了49%的失败率。

这些改进凸显了上下文检索的潜力，可以提高 AI 驱动的问答系统的性能，使其更加准确和上下文感知。

我们将构建什么

我们将使用两个示例文档来演示上下文检索如何改进问答系统。我们的系统将执行以下操作：

1. 将文档拆分成更小的信息块。
2. 向每个信息块添加上下文信息，将其嵌入，并将它们存储在数据库中。
3. 执行相似性搜索以找到最相关的上下文。
4. 使用大型语言模型（LLM）根据检索到的上下文生成用户问题的答案。

设置环境

首先，让我们安装必要的库：

pip install -Uqqq pip --progress-bar off

pip install -qqq fastembed==0.3.6 --progress-bar off

pip install -qqq sqlite-vec==0.1.2 --progress-bar off

pip install -qqq groq==0.11.0 --progress-bar off

pip install -qqq langchain-text-splitters==0.3.0 --progress-bar off

现在，让我们导入所需的模块：

import sqlite3

from textwrap import dedent

from typing import List

 

import sqlite_vec

from fastembed import TextEmbedding

from google.colab import userdata

from groq import Groq

from groq.types.chat import ChatCompletionMessage

from langchain_text_splitters import RecursiveCharacterTextSplitter

from sqlite_vec import serialize_float32

from tqdm import tqdm

语言模型设置

我们将通过 Groq API 使用 Llama 3.1。首先，让我们设置客户端：

client = Groq(api_key=userdata.get("GROQ_API_KEY"))

MODEL = "llama-3.1-70b-versatile"

TEMPERATURE = 0

接下来，我们将创建一个辅助函数来与模型交互。此函数将接受提示和可选的消息历史记录：

def call_model(prompt: str, messages=[]) -> ChatCompletionMessage:

    messages.append({

        "role": "user",

        "content": prompt,

    })

    response = client.chat.completions.create(

        model=MODEL,

        messages=messages,

        temperature=TEMPERATURE,

    )

    return response.choices[0].message.content

此函数向模型发送提示并返回模型的响应。您还可以传递消息历史记录以维护对话的上下文。

数据库设置

我们将使用带有sqlite-vec扩展的SQLite来存储我们的文档及其嵌入。以下是设置数据库的方法：

db = sqlite3.connect("readmes.sqlite3")

db.enable_load_extension(True)

sqlite_vec.load(db)

db.enable_load_extension(False)

连接到数据库后，让我们创建必要的表：

db.execute("""

CREATE TABLE documents(

    id INTEGER PRIMARY KEY AUTOINCREMENT,

    text TEXT

);

""")

 

db.execute("""

CREATE TABLE chunks(

    id INTEGER PRIMARY KEY AUTOINCREMENT,

    document_id INTEGER,

    text TEXT,

    FOREIGN KEY(document_id) REFERENCES documents(id)

);

""")

 

db.execute(f"""

CREATE VIRTUAL TABLE chunk_embeddings USING vec0(

  id INTEGER PRIMARY KEY,

  embedding FLOAT[{document_embeddings[0].shape[0]}]

);

""")

以下是表格的分类：

• documents：存储每个文档的全文。
• chunks：存储从文档中拆分的较小文本块。
• chunk_embeddings：存储每个块的嵌入，以便进行相似性搜索。

这种数据库设置允许我们有效地存储、检索和嵌入块，从而便于以后执行相似性搜索。

创建数据块

为了将文档分解为可管理的块以便更好地进行上下文检索，我们将按照以下步骤操作：

1. 将文档文本拆分为较小的块。
2. 向每个数据块添加上下文信息。
3. 嵌入每个块并将其与文本一起存储在数据库中。

我们将使用的文档是Qwen 2.5模型和LangGraph项目的README文件。

首先，让我们将文档保存在数据库中：

documents = [qwen_doc, langgraph_doc]

 

with db:

    for doc in documents:

        db.execute("INSERT INTO documents(text) VALUES(?)", [doc])

为了将文档拆分成更小的信息块，我们将使用LangChain中的RecursiveCharacterTextSplitter3工具：

text_splitter = RecursiveCharacterTextSplitter(chunk_size=2048, chunk_overlap=128)

我们现在可以创建块并将它们存储在数据库中：

with db:

    document_rows = db.execute("SELECT id, text FROM documents").fetchall()

    for row in document_rows:

        doc_id, doc_text = row

        chunks = text_splitter.split_text(doc_text)

        contextual_chunks = create_contextual_chunks(chunks, doc_text)

        save_chunks(contextual_chunks)

为了给每个数据块提供额外的上下文，我们将使用以下提示生成简短的摘要：

CONTEXTUAL_EMBEDDING_PROMPT = """

Here is the chunk we want to situate within the whole document:

<chunk>

{chunk}

</chunk>

 

Here is the content of the whole document:

<document>

{document}

</document>

 

Please provide a short, succinct context to situate this chunk within the overall document to improve search retrieval. Respond only with the context.

"""

以下是该函数的工作原理：

def create_contextual_chunks(chunks: List[str], document: str) -> List[str]:

    contextual_chunks = []

    for chunk in chunks:

        prompt = CONTEXTUAL_EMBEDDING_PROMPT.format(chunk=chunk, document=document)

        chunk_context = call_model(prompt)

        contextual_chunks.append(f"{chunk}\n{chunk_context}")

    return contextual_chunks

此函数会将每个信息块连同整个文档一起发送到模型，模型会生成一个简短的上下文，以提高搜索检索的准确性。然后，将这个上下文前置到信息块的前面。

我们将使用fastembed4库来为文档的信息块创建嵌入表示：

embedding_model = TextEmbedding()

最后，让我们将块及其嵌入保存在数据库中：

def save_chunks(chunks: List[str]):

    chunk_embeddings = list(embedding_model.embed(chunks))

    for chunk, embedding in zip(chunks, chunk_embeddings):

        result = db.execute(

            "INSERT INTO chunks(document_id, text) VALUES(?, ?)", [doc_id, chunk]

        )

        chunk_id = result.lastrowid

        db.execute(

            "INSERT INTO chunk_embeddings(id, embedding) VALUES (?, ?)",

            [chunk_id, serialize_float32(embedding)],

        )

此函数将每个信息块及其嵌入表示保存到数据库中的chunks表和chunk_embeddings表中。serialize_float32函数用于将嵌入表示存储为一种可以稍后高效检索的格式。

检索上下文

一旦块及其嵌入存储在数据库中，我们就可以检索给定查询的最相关上下文。下面是实现这一点的函数：

def retrieve_context(query: str, k: int = 3, embedding_model: TextEmbedding = embedding_model) -> str:

    query_embedding = list(embedding_model.embed([query]))[0]

    results = db.execute(

        """

    SELECT

        chunk_embeddings.id,

        distance,

        text

    FROM chunk_embeddings

    LEFT JOIN chunks ON chunks.id = chunk_embeddings.id

    WHERE embedding MATCH ? AND k = ?

    ORDER BY distance

        """,

        [serialize_float32(query_embedding), k],

    ).fetchall()

    return "\n-----\n".join([item[2] for item in results])

1. 查询嵌入：该函数首先使用嵌入模型将输入查询转换为嵌入表示。
2. 数据库查询：然后，它通过以下方式检索与查询嵌入表示最相似的前k个信息块及其嵌入表示：
   • 计算查询嵌入表示与存储的信息块嵌入表示之间的余弦相似度（这由sqlite-vec扩展处理）。
   • 按相似度距离对结果进行排序（距离越低表示匹配越紧密）。
3. 返回结果：将检索到的信息块连接成一个单独的字符串，并用“\n—–\n”分隔以提高清晰度。

生成答案

为了生成答案，我们将系统提示符与检索到的上下文相结合。这可确保模型提供准确且与上下文相关的响应。

系统提示为模型应如何响应设定基调和期望：

SYSTEM_PROMPT = """

You're an expert AI/ML engineer with a background in software development.

You're answering questions about technical topics and projects.

If you don't know the answer, simply state that you don't know. 

Keep your answers brief and to the point. Be kind and respectful.

 

Use the provided context for your answers. The most relevant information is 

at the top. Each piece of information is separated by ---.

"""

以下是将所有内容联系在一起的函数：

def ask_question(query: str) -> str:

    messages = [

        {

            "role": "system",

            "content": SYSTEM_PROMPT,

        },

    ]

    context = retrieve_context(query)

    prompt = dedent(

        f"""

Use the following information:

 

```

{context}

```

 

to answer the question:

{query}

        """

    )

    return call_model(prompt, messages), context

1. 设置系统提示：SYSTEM_PROMPT指导模型如何回答问题——鼓励简洁、礼貌且考虑上下文的回答。如果模型不知道答案，它会按照指示承认这一点。
2. 检索相关上下文：retrieve_context(query)函数从数据库中为给定查询检索最相关的上下文信息块。
3. 创建最终提示：将检索到的上下文插入到提示中，然后指示模型使用该信息来回答用户的问题。
4. 调用模型：call_model(prompt, messages)函数将提示发送到大型语言模型（LLM）并生成答案。
5. 返回响应：该函数返回模型生成的答案以及检索到的上下文（供审查时可选）。

要回答问题，您可以像这样调用函数：

answer, context = ask_question("How does Contextual Retrieval improve RAG performance?")

print("Answer:", answer)

print("Context used:", context)

这既提供了答案，也提供了模型用来生成回答的上下文。

使用 RAG

现在我们可以用一些问题来测试我们的系统。让我们先问一个关于Qwen模型的简单问题：

query = "How many parameters does Qwen have?"

response, context = ask_question(query)

print(response)

输出：

Qwen2.5 models are available in various sizes, with the number of parameters 

ranging from 0.5B to 72B. The specific model mentioned in the text has 32.5B 

parameters, with 31.0B non-embedding parameters.

非常好，看起来模型是基于检索到的上下文提供了准确的信息。让我们尝试一些技术性更强的内容：

query = "How should one deploy Qwen model on a private server?"

response, context = ask_question(query)

print(response)

输出：

To deploy Qwen2.5 on a private server, you can use vLLM, a fast and easy-to-use 

framework for LLM inference and serving. First, install `vllm>=0.4.0` using 

pip. Then, run the following command to build up a vLLM service:

 

```bash

python -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-7B-Instruct

```

 

Alternatively, with `vllm>=0.5.3`, you can use:

 

```bash

vllm serve Qwen/Qwen2.5-7B-Instruct

```

 

This will start a service that you can interact with using the OpenAI API.

这是对文档部署部分的一个很好的总结。让我们再尝试一个问题：

query = "I have a RTX 4090 (24GB). Which version of the model can I run with good inference speed?"

response, context = ask_question(query)

print(response)

输出：

Based on the provided information, the model sizes available for Qwen2.5 are 

0.5B, 1.5B, 3B, 7B, 14B, 32B, and 72B. 

 

Considering your RTX 4090 has 24GB of memory, you can likely run the 7B or 14B 

models with good inference speed. However, the 14B model might be pushing the 

limits of your GPU's memory, so the 7B model would be a safer choice.

 

Keep in mind that the actual performance will also depend on other factors such 

as your system's CPU, RAM, and the specific use case.

此信息在文档中找不到，但该模型根据检索到的上下文及其推理能力提供了很好的答案。

结论

您已经构建了一个 RAG 系统，该系统使用：

• 上下文分块：将文档分解为有意义的信息块，从而提高检索准确性。
• 高效的相似性搜索：使用向量嵌入来查找最相关的信息。
• 语言模型集成：利用强大的模型根据检索到的上下文生成自然语言响应。

在继续完善此系统的过程中，您可以考虑通过以下方式进行增强：

• 缓存：为了更快的响应时间和更高的性能（如果使用提示缓存则可以实现).
• 多文档支持： 扩展以处理更多不同类型的文档。
• 用户友好的界面：使非技术用户能够访问系统。

告诉我您打算用这个系统来构建什么！

原文链接：https://www.mlexpert.io/blog/rag-contextual-retrieval