RAG入门 - Retriever（1）

💡 从上图能看出来，RAG 架构中有许多步骤是可以优化的，正确的优化会带来显著的效果提升。

在这篇文章中，我们会重点关注蓝色内容，来调整我们自己的 RAG 系统来获得最佳效果。

现在，让我们把手弄脏，直接跟着文章的思路来了解RAG的优化过程。

环境准备

在开始之前，我们需要装好如下依赖：

# 建议使用conda创建一个干净的虚拟环境
conda create --name huggingface python=3.10 -y
conda activate huggingface
pip install torch transformers accelerate bitsandbytes langchain sentence-transformers openpyxl pacmap datasets langchain-community ragatouille faiss

faiss安装可能会出现问题，解决方案:

$ pip install faiss-cpu
# or: 
$ pip install faiss-gpu-cu12 # CUDA 12.x, Python 3.8+
$ pip install faiss-gpu-cu11 # CUDA 11.x, Python 3.8+
$ pip install faiss-gpu # Python 3.6-3.10 (legacy, no longer available after version 1.7.3)

torch、cuda 等安装验证及版本查看：

import torch
# 获取 PyTorch 版本信息
print("PyTorch Version: {}".format(torch.version.__version__))  # 或直接使用 torch.__version__
print("PyTorch CUDA Version: {}".format(torch.version.cuda))  # 获取 CUDA 版本
print("PyTorch cuDNN Version: {}".format(torch.backends.cudnn.version()))  # 获取 cuDNN 版本

因为在文章代码中会自动下载huggingface上的模型和数据集，会默认存储在~/.cache/huggingface目录下。如果你担心系统盘不够存储这些数据，你也可以修改huggingface cache的默认根目录：

export HF_HOME="/{to_path}/huggingface"

# 为了不用每次都执行，你可以直接写入bash配置
echo 'export HF_HOME="/{to_path}/huggingface"' >> ~/.bashrc

另外，国内访问huggingface是受限的（墙），我们可以使用huggingface 国内镜像站运行python脚本:

HF_ENDPOINT=https://hf-mirror.com python advanced_rag.py

from tqdm.notebook import tqdm
import pandas as pd
from typing import Optional, List, Tuple
from datasets import Dataset
import matplotlib.pyplot as plt
pd.set_option("display.max_colwidth", None) # This will be helpful when visualizing retriever outputs 

知识库加载

import datasets
ds = datasets.load_dataset("m-ric/huggingface_doc", split="train")
from langchain.docstore.document import Document as LangchainDocument
RAW_KNOWLEDGE_BASE = [
        LangchainDocument(page_content=doc["text"], metadata={
   
   "source": doc["source"]}) for doc in tqdm(ds)
    ]

1. Retriever - embeddings 🗂️

retriever像一个内置的搜索引擎：接收用户的查询，返回知识库中的一些相关片段。这些片段会输入到Reader Model（如deepseek）中，来帮助它生成答案。因此，现在我们的目标就是，基于用户的问题，从我们的知识库中找到最相关的片段来回答这个问题。这是一个宽泛的目标，它引申出了一堆问题。比如我们应该检索多少个片段？这个关于片段数量的参数就被命名为 top_k 。再比如，每个片段应该有多长？这个片段长度的参数就被称为 chunk size 。这些问题没有唯一的适合所有情况的答案，但有一些相关知识我们可以了解下：

• 🔀 不同的片段可以有不同的chunk size 。

• 由于检索内容中总会有一些噪音，增加 top_k 的值会增加在检索到的片段中获得相关内容的机会。类似射箭🎯， 射出更多的箭会增加你击中目标的概率。

• 同时，检索到的文档的总长度不应太长：比如，对于目前大多数的模型，16k的token数量可能会让模型因为“Lost in the middle phenomemon”而淹没在信息中。所以，只给模型提供最相关的见解，而不是一大堆内容！

在这篇文章中，我们使用 Langchain 库，因为它提供了大量的向量数据库选项，并允许我们在处理过程中保持文档的元数据。

1.1 将文档拆分为chunks

在这一部分，我们将知识库中的文档拆分为更小的chunks，chat LLM 会基于这些chunks进行回答。我们的目标是得到一组语义相关的片段。因此，它们的大小需要适应具体的主题或者说是中心思想：太小的话会截断中心思想，太大可能就会稀释中心思想，被其他不相关内容干扰。

💡 现在有许多拆分文本内容的方案，比如：按词拆分、按句子边界拆分、递归拆分（以树状方式处理文档以保留结构信息）……要了解更多关于文本拆分的内容，可以参考这篇文档。

递归分块通过使用一组按重要性排序的分隔符，将文本逐步分解为更小的部分。如果第一次拆分没有给出正确大小的块，它就会在新的块上使用不同的分隔符来重复这个步骤。比如，我们可以使用这样的分隔符列表 ["\n\n", "\n", ".", ""] ：

这种方法很好的保留了文档的整体结构，但代价是块大小会有轻微的变化。

这里可以让你看到不同的拆分选项会如何影响你得到的块。

🔬 让我们先用一个任意大小的块来做一个实验，看看拆分具体是怎么工作的。我们直接使用 Langchain 的递归拆分类 RecursiveCharacterTextSplitter 。

from langchain.text_splitter import RecursiveCharacterTextSplitter

# We use a hierarchical list of separators specifically tailored for splitting Markdown documents
# This list is taken from LangChain's MarkdownTextSplitter class
MARKDOWN_SEPARATORS = [
    "\n#{1,6} ",
    "```\n",
    "\n\\*\\*\\*+\n",
    "\n---+\n",
    "\n___+\n",
    "\n\n",
    "\n",
    " ",
    "",
]

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,  # The maximum number of characters in a chunk: we selected this value arbitrarily
    chunk_overlap=100,  # The number of characters to overlap between chunks
    add_start_index=True,  # If `True`, includes chunk's start index in metadata
    strip_whitespace=True,  # If `True`, strips whitespace from the start and end of every document
    separators=MARKDOWN_SEPARATORS,
)

docs_processed = []
for doc in RAW_KNOWLEDGE_BASE:
    docs_processed += text_splitter.split_documents([doc])

• 其中，参数 chunk_size 控制单个块的长度：这个长度默认是按块中的字符数来计算的。

• 参数 chunk_overlap 是为了允许相邻的块之间有一些重叠，这能够减少一个主题可能在两个相邻块的分割中被切成两半的概率。我们把它设置为块大小的 1/10，当然你也可以自己尝试其他不同的值！

我们利用以下代码看看chunk的长度分布：

lengths = [len(doc.page_content) for doc in tqdm(docs_processed)]

# Plot the distribution of document lengths, counted as the number of chars
fig = pd.Series(lengths).hist()
plt.title("Distribution of document lengths in the knowledge base (in count of chars)")
plt.show()

💡如果你使用了远程设备不支持直接使用 plt.show() 可视化，可以换成用如下代码直接保存成图片。

plt.savefig("chunk_sizes_char.png", dpi=300, bbox_inches="tight")

可视化结果如下，我们可以看到最大的块的字符长度不会超过1000，这和我们预先设置的参数一致。

1.2 词嵌入

接下来，我们需要使用词嵌入模型来对分块进行向量化。在使用词嵌入模型时，我们需要知道模型能接受的最大序列长度max_seq_length（按照token数统计）。需要确保分块的token数低于这个值，因为超过max_seq_length的块在处理之前都会被截断，从而失去相关性。这里我们使用的嵌入模型是thenlper/gte-small， 下面代码先打印了该模型支持的最大长度，然后再对分块结果进行token数量的分布统计。

from sentence_transformers import SentenceTransformer

# To get the value of the max sequence_length, we will query the underlying `SentenceTransformer` object used in the RecursiveCharacterTextSplitter
print(f"Model's maximum sequence length: {
     
     SentenceTransformer('thenlper/gte-small').max_seq_length}")

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("thenlper/gte-small")
lengths = [len(tokenizer.encode(doc.page_content)) for doc in tqdm(docs_processed)]

# Plot the distribution of document lengths, counted as the number of tokens
fig = pd.Series(lengths).hist()
plt.title("Distribution of document lengths in the knowledge base (in count of tokens)")
plt.show()

上面代码会先输出：

Model's maximum sequence length: 512

表示thenlper/gte-small 支持的最大块长度是 512

可视化结果如下：

👀 可以看到，某些分块的token数量超过了 512 的限制，这样就会导致分块中的一部分内容会因截断而丢失！

• 既然是基于token数来统计，那我们就应该将 RecursiveCharacterTextSplitter 类更改为以token数量而不是字符数量来计算长度。

• 然后我们可以选择一个特定的块大小，这里我们选择一个低于 512 的阈值：

  • 较小的文档可以使分块更专注于特定的主题。

  • 但过小的块又会将完整的句子一分为二，从而再次失去意义，所以这也需要我们根据实际情况进行权衡。

from langchain.text_splitter import RecursiveCharacterTextSplitter
from transformers import AutoTokenizer

EMBEDDING_MODEL_NAME = "thenlper/gte-small"


def split_documents(
    chunk_size: int,
    knowledge_base: List[LangchainDocument],
    tokenizer_name: Optional[str] = EMBEDDING_MODEL_NAME,
) -> List[LangchainDocument]:
    """
    Split documents into chunks of maximum size `chunk_size` tokens and return a list of documents.
    """
    text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
        AutoTokenizer.from_pretrained(tokenizer_name),
        chunk_size=chunk_size,
        chunk_overlap=int(chunk_size / 10),
        add_start_index=True, # 是否在每个分块的 metadata 中添加该分块在原始文档中的起始字符索引
        strip_whitespace=True, # 是否去除每个分块开头和结尾的空白字符（如空格、换行等）
        separators=MARKDOWN_SEPARATORS,
    )

    docs_processed = []
    for doc in knowledge_base:
        docs_processed += text_splitter.split_documents([doc])

    # Remove duplicates
    unique_texts = {
   
   }
    docs_processed_unique = []
    for doc in docs_processed:
        if doc.page_content not in unique_texts:
            unique_texts[doc.page_content] = True
            docs_processed_unique.append(doc)

    return docs_processed_unique


docs_processed = split_documents(
    512,  # We choose a chunk size adapted to our model
    RAW_KNOWLEDGE_BASE,
    tokenizer_name=EMBEDDING_MODEL_NAME,
)

# Let's visualize the chunk sizes we would have in tokens from a common model
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained