检索增强生成（RAG）已成为增强大型语言模型能力的一种强大范例。通过结合检索系统和生成模型的优势，RAG 系统可以产生更准确、更事实、更符合上下文的响应。当处理特定领域知识或需要最新信息时，这种方法尤其有价值。
在本帖中，您将探索如何使用 Hugging Face 库中的模型构建一个基本的 RAG 系统。您将构建从文档索引到检索和生成中的每个系统组件，并实现一个完整的端到端解决方案。具体来说，您将学习
- RAG 架构及其组件
- 如何构建文档索引和检索系统
- 如何实现基于 Transformer 的生成器
通过我的书籍《Hugging Face Transformers中的NLP》，快速启动您的项目。它提供了带有工作代码的自学教程。
让我们开始吧！
使用 Transformers 构建 RAG 系统
Tina Nord 拍摄的照片。保留部分权利。
概述
本帖分为五个部分
- 理解 RAG 架构
- 构建文档索引系统
- 实现检索系统
- 实现生成器
- 构建完整的 RAG 系统
理解 RAG 架构
RAG 系统包含两个主要组件
- 检索器：负责根据查询从知识库中查找相关文档或段落。
- 生成器：使用检索到的文档和原始查询来生成连贯且信息丰富的响应。
这些组件中的每一个都有许多细节。您需要 RAG 是因为生成器本身（即语言模型）无法生成准确且符合上下文的响应，这被称为幻觉。因此，您需要检索器来提供线索以帮助生成器。
这种方法结合了生成模型广泛的语言理解能力与从知识库访问特定信息的能力。这使得响应既流畅又事实准确。
让我们一步一步地实现 RAG 系统的每个组件。
构建文档索引系统
创建 RAG 系统的第一步是构建文档索引系统。该系统必须将文档编码为密集向量表示，并将其存储在数据库中。然后，我们可以根据上下文相似性检索文档。这意味着您需要能够通过向量相似性度量进行搜索，而不是精确匹配。这是一个关键点——并非所有数据库系统都可以用于构建文档索引系统。
当然，您可以收集文档，将它们编码为向量表示，并将它们保存在内存中。当请求检索时，您可以逐个计算相似性以找到最接近的匹配。但是，在循环中检查每个向量效率低下且不可扩展。FAISS 是一个针对此任务进行了优化的库。要安装 FAISS，您可以从源代码编译它，或使用 PyPI 上的预编译版本。
|
1 |
pip install faiss-cpu |
接下来，您将创建一个语言模型，将文档编码为密集向量表示，并将它们存储在 FAISS 索引中以进行高效检索。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import faiss import torch from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") def generate_embedding(docs, model, tokenizer): # Tokenize each text and convert to PyTorch tensors inputs = tokenizer(docs, padding=True, truncation=True, return_tensors="pt", max_length=512) with torch.no_grad(): outputs = model(**inputs) // 嵌入定义为所有 token 的平均池化 attention_mask = inputs["attention_mask"] embeddings = outputs.last_hidden_state expanded_mask = attention_mask.unsqueeze(-1).expand(embeddings.shape).float() sum_embeddings = torch.sum(embeddings * expanded_mask, axis=1) sum_mask = torch.clamp(expanded_mask.sum(axis=1), min=1e-9) mean_embeddings = sum_embeddings / sum_mask // 转换为 numpy 数组 return mean_embeddings.cpu().numpy() // 样本文档集合 documents = [ "Transformers are a type of deep learning model introduced in the paper 'Attention " "Is All You Need'.", "BERT (Bidirectional Encoder Representations from Transformers) is a " "transformer-based model designed to understand the context of a word based on " "its surroundings.", "GPT (Generative Pre-trained Transformer) is a transformer-based model designed for " "natural language generation tasks.", "T5 (Text-to-Text Transfer Transformer) treats every NLP problem as a text-to-text " "problem, where both the input and output are text strings.", "RoBERTa is an optimized version of BERT with improved training methodology and more " "training data.", "DistilBERT is a smaller, faster version of BERT that retains 97% of its language " "understanding capabilities.", "ALBERT reduces the parameters of BERT by sharing parameters across layers and using " "embedding factorization.", "XLNet is a generalized autoregressive pretraining method that overcomes the " "limitations of BERT by using permutation language modeling.", "ELECTRA uses a generator-discriminator architecture for more efficient pretraining.", "DeBERTa enhances BERT with disentangled attention and an enhanced mask decoder." ] // 为所有文档生成嵌入，然后创建 FAISS 索引以进行高效相似性搜索 document_embeddings = generate_embedding(documents, model, tokenizer) dimension = document_embeddings.shape[1] // 嵌入的维度 index = faiss.IndexFlatL2(dimension) // 使用 L2（欧几里得）距离 index.add(document_embeddings) // 将嵌入添加到索引 print(f"Created index with {index.ntotal} documents") |
这段代码的关键在于 generate_embedding() 函数。它接受一个文档列表，通过模型对它们进行编码，并使用来自每个文档的所有 token 嵌入的平均池化返回密集向量表示。文档不需要很长且完整。句子或段落是期望的，因为模型有上下文窗口限制。此外，您稍后将在另一个示例中看到，非常长的文档不适合 RAG。
您使用了一个预训练的 Sentence Transformer 模型 sentence-transformers/all-MiniLM-L6-v2，它专门用于生成句子嵌入。您不会将原始文档保留在 FAISS 索引中；您只保留嵌入向量。您预构建了这些向量之间的 L2 距离索引，以进行高效的相似性搜索。
您可以修改此代码以实现 RAG 系统的不同实现。例如，密集向量表示是通过平均池化获得的。但是，您可以使用第一个 token，因为分词器会在每个句子前面加上 [CLS] token，并且模型应该生成围绕这个特殊 token 的上下文嵌入。此外，这里使用 L2 距离是因为您声明 FAISS 索引时打算将其与 L2 度量一起使用。FAISS 中没有余弦相似性度量，但 L2 和余弦距离是相似的。请注意，对于归一化向量，
$$
\begin{align}
\Vert \mathbf{x} – \mathbf{y} \Vert_2^2
&= (\mathbf{x} – \mathbf{y})^\top (\mathbf{x} – \mathbf{y}) \\
&= \mathbf{x}^\top \mathbf{x} – 2 \mathbf{x}^\top \mathbf{y} + \mathbf{y}^\top \mathbf{y} \\
&= 2 – 2 \mathbf{x}^\top \mathbf{y} \\
&= 2 – 2 \cos \theta
\end{align}
$$
因此，当向量被归一化时（只要您记住当不相似性增加时，L2 从 0 跑到无穷大，而余弦距离从 +1 跑到 -1），L2 距离等同于余弦距离。如果您打算使用余弦距离，则应修改代码如下：
|
1 2 3 4 |
... document_embeddings = generate_embedding(documents, model, tokenizer) normalized = document_embeddings / np.linalg.norm(document_embeddings, axis=1, keepdims=True) index.add(normalized) |
基本上，您缩放了每个嵌入向量以使其成为单位长度。
实现检索系统
在索引了文档之后，让我们看看如何为给定的查询检索一些最相关的文档。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
... def retrieve_documents(query, index, documents, k=3): // 为查询生成嵌入 query_embedding = generate_embedding(query, model, tokenizer) // 1xD 矩阵 // 在索引中搜索相似文档 distances, indices = index.search(query_embedding, k) // 1xk 矩阵 // 返回检索到的文档及其距离 retrieved_docs = [(documents[idx], float(distances[0][i])) for i, idx in enumerate(indices[0])] return retrieved_docs # 示例查询 query = "What is BERT?" retrieved_docs = retrieve_documents(query, index, documents) // 打印检索到的文档 print(f"Query: {query}\n") for i, (doc, distance) in enumerate(retrieved_docs): print(f"Document {i+1} (Distance: {distance:.4f}):") print(doc) print() |
如果您运行此代码，您将看到以下输出：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
Query: What is BERT? Document 1 (Distance: 23.7060) BERT (Bidirectional Encoder Representations from Transformers) is a transformer-based model designed to understand the context of a word based on its surroundings. Document 2 (Distance: 28.0794) RoBERTa is an optimized version of BERT with improved training methodology and more training data. Document 3 (Distance: 29.5908) DistilBERT is a smaller, faster version of BERT that retains 97% of its language understanding capabilities. |
在 retrieve_documents() 函数中，您提供查询字符串、FAISS 索引和文档集合。然后，您像处理文档一样为查询生成嵌入。之后，您利用 FAISS 索引的 search() 方法来查找与查询嵌入最相似的 k 个文档。search() 方法返回两个数组：
distances：查询嵌入与索引嵌入之间的距离。由于这是您定义索引的方式，因此这些是 L2 距离。indices：索引嵌入的索引，这些索引与查询嵌入最相似，与距离数组匹配。
您可以使用这些数组从原始集合中检索最相似的文档。在此，您使用索引从列表中获取文档。之后，您打印检索到的文档及其在嵌入空间中与查询的距离，按相关性降序或距离升序排列。
请注意，文档的上下文向量应代表整个文档。因此，如果文档包含大量信息，查询与文档之间的距离可能会很大。理想情况下，您希望文档集中且简洁。如果您有很长的文本，您可能需要将其拆分成多个文档，以使 RAG 系统更准确。
这个检索系统构成了我们 RAG 架构的第一个组件。给定用户查询，它允许我们从知识库中查找相关信息。有许多其他方法可以实现相同的功能，但这突显了向量搜索的关键思想。
实现生成器
接下来，让我们实现 RAG 系统的生成器组件。
这是一个提示工程问题。在用户提供查询时，您首先从检索器中检索最相关的文档，并创建一个新的提示，其中包含用户的查询和检索到的文档作为上下文。然后，您使用预训练的语言模型根据新提示生成响应。
以下是实现方法：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
... from transformers import AutoModelForSeq2SeqLM gen_tokenizer = AutoTokenizer.from_pretrained("t5-small") gen_model = AutoModelForSeq2SeqLM.from_pretrained("t5-small") def generate_response(query, retrieved_docs, max_length=150): // 将查询和检索到的文档合并成一个提示 context = "\n".join(retrieved_docs) prompt = f"question: {query} context: {context}" // 生成响应 inputs = gen_tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True) with torch.no_grad(): outputs = gen_model.generate( inputs.input_ids, max_length=max_length, num_beams=4, early_stopping=True, no_repeat_ngram_size=2 ) response = gen_tokenizer.decode(outputs[0], skip_special_tokens=True) return response // 为示例查询生成响应 response = generate_response(query, [doc for doc, score in retrieved_docs]) print("Generated Response:") print(response) |
这是我们 RAG 系统的生成器组件。您实例化了预训练的 T5 模型（小版本，但您可以选择一个更大或不同的模型来运行）。这个模型是一个序列到序列模型，它从给定的序列生成新序列。如果您使用不同的模型，例如“因果 LM”模型，您可能需要更改提示以使其工作更有效。
在 generate_response() 函数中，您将查询和检索到的文档合并成一个单一的提示。然后，您使用 T5 模型生成响应。您可以调整生成参数以使其效果更好。在上面，为了简单起见，只使用了束搜索。模型输出然后被解码为文本字符串作为响应。由于您将多个文档合并到一个提示中，因此您需要注意提示不要超过模型的上下文窗口。
生成器利用检索到的文档中的信息来产生流畅且事实准确的响应。当您仅提出查询而不提供上下文时，模型的行为差异很大。
构建完整的 RAG 系统
这就是构建基本 RAG 系统所需的一切。让我们创建一个函数来封装检索和生成组件。
|
1 2 3 4 5 |
... def rag_pipeline(query, documents, retriever_k=3, max_length=150): retrieved_docs = retrieve_documents(query, index, documents, k=retriever_k) response = generate_response(query, retrieved_docs, max_length=max_length) return response, retrieved_docs |
然后，您可以循环使用 RAG 管道为一组查询生成响应。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
... // 示例查询 queries = [ "What is BERT?", "How does GPT work?", "What is the difference between BERT and GPT?", "What is a smaller version of BERT?" ] // 对每个查询运行 RAG 管道 for query in queries: response, retrieved_docs = rag_pipeline(query, documents) print(f"查询: {query}") print() print("检索到的文档:") for i, (doc, distance) in enumerate(retrieved_docs): print(f"文档 {i+1} (距离: {distance:.4f}):") print(doc) print() print("生成的响应:") print(response) print("-" * 20) |
您可以看到查询是如何在一个循环中一个接一个地回答的。然而，文档集是提前准备好的，并为所有查询重用。RAG系统通常是这样工作的。
以上所有内容的完整代码如下
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import faiss import torch from transformers import AutoTokenizer, AutoModel from transformers import AutoModelForSeq2SeqLM # 检索器中使用的模型 tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") # 生成器中使用的模型 gen_tokenizer = AutoTokenizer.from_pretrained("t5-small") gen_model = AutoModelForSeq2SeqLM.from_pretrained("t5-small") def generate_embedding(docs, model, tokenizer): # Tokenize each text and convert to PyTorch tensors inputs = tokenizer(docs, padding=True, truncation=True, return_tensors="pt", max_length=512) with torch.no_grad(): outputs = model(**inputs) // 嵌入定义为所有 token 的平均池化 attention_mask = inputs["attention_mask"] embeddings = outputs.last_hidden_state expanded_mask = attention_mask.unsqueeze(-1).expand(embeddings.shape).float() sum_embeddings = torch.sum(embeddings * expanded_mask, axis=1) sum_mask = torch.clamp(expanded_mask.sum(axis=1), min=1e-9) mean_embeddings = sum_embeddings / sum_mask // 转换为 numpy 数组 return mean_embeddings.cpu().numpy() def retrieve_documents(query, index, documents, k=3): // 为查询生成嵌入 query_embedding = generate_embedding(query, model, tokenizer) // 1xD 矩阵 // 在索引中搜索相似文档 distances, indices = index.search(query_embedding, k) // 1xk 矩阵 // 返回检索到的文档及其距离 retrieved_docs = [(documents[idx], float(distances[0][i])) for i, idx in enumerate(indices[0])] return retrieved_docs def generate_response(query, retrieved_docs, max_length=150): // 将查询和检索到的文档合并成一个提示 if retrieved_docs: context = "\n".join(retrieved_docs) prompt = f"question: {query} context: {context}" else: prompt = f"question: {query}" // 生成响应 inputs = gen_tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True) with torch.no_grad(): outputs = gen_model.generate( inputs.input_ids, max_length=max_length, num_beams=4, early_stopping=True, no_repeat_ngram_size=2 ) response = gen_tokenizer.decode(outputs[0], skip_special_tokens=True) return response def rag_pipeline(query, documents, retriever_k=3, max_length=150): retrieved_docs = retrieve_documents(query, index, documents, k=retriever_k) docs = [doc for doc, distance in retrieved_docs] response = generate_response(query, docs, max_length=max_length) return response, retrieved_docs // 样本文档集合 documents = [ "Transformers are a type of deep learning model introduced in the paper 'Attention " "Is All You Need'.", "BERT (Bidirectional Encoder Representations from Transformers) is a " "transformer-based model designed to understand the context of a word based on " "its surroundings.", "GPT (Generative Pre-trained Transformer) is a transformer-based model designed for " "natural language generation tasks.", "T5 (Text-to-Text Transfer Transformer) treats every NLP problem as a text-to-text " "problem, where both the input and output are text strings.", "RoBERTa is an optimized version of BERT with improved training methodology and more " "training data.", "DistilBERT is a smaller, faster version of BERT that retains 97% of its language " "understanding capabilities.", "ALBERT reduces the parameters of BERT by sharing parameters across layers and using " "embedding factorization.", "XLNet is a generalized autoregressive pretraining method that overcomes the " "limitations of BERT by using permutation language modeling.", "ELECTRA uses a generator-discriminator architecture for more efficient pretraining.", "DeBERTa enhances BERT with disentangled attention and an enhanced mask decoder." ] // 为所有文档生成嵌入，然后创建 FAISS 索引以进行高效相似性搜索 document_embeddings = generate_embedding(documents, model, tokenizer) dimension = document_embeddings.shape[1] // 嵌入的维度 index = faiss.IndexFlatL2(dimension) // 使用 L2（欧几里得）距离 index.add(document_embeddings) // 将嵌入添加到索引 print(f"Created index with {index.ntotal} documents") // 示例查询 queries = [ "What is BERT?", "How does GPT work?", "What is the difference between BERT and GPT?", "What is a smaller version of BERT?" ] // 对每个查询运行 RAG 管道 for query in queries: response, retrieved_docs = rag_pipeline(query, documents) print(f"查询: {query}") print() print("检索到的文档:") for i, (doc, distance) in enumerate(retrieved_docs): print(f"文档 {i+1} (距离: {distance:.4f}):") print(doc) print() print("生成的响应:") print(response) print("-" * 20) |
此代码是独立的。所有文档和查询都在代码中定义。这是一个起点，您可以扩展它以实现新功能，例如将索引的文档保存在一个文件中，以后无需每次都重新索引即可加载。
进一步阅读
以下是一些您可能觉得有用的进一步阅读资料：
总结
本文探讨了使用 Hugging Face 库中的 Transformer 模型构建检索增强生成（RAG）系统。我们实现了从文档索引到检索和生成的每个系统组件，并将它们组合成一个完整的端到端解决方案。
RAG 系统代表了一种增强语言模型能力的方法，方法是将其与外部知识相结合。RAG 系统通过检索相关信息并将其纳入生成过程，可以生成更准确、事实性更强、更具上下文相关性的响应。
想在您的NLP项目中使用强大的语言模型吗？
在您自己的机器上运行最先进的模型
...只需几行Python代码在我的新电子书中探索如何实现
使用 Hugging Face Transformers 进行自然语言处理
它涵盖了在以下任务上的实践示例和实际用例：文本分类、摘要、翻译、问答等等...
最终将高级NLP带入
您自己的项目
没有理论。只有实用的工作代码。
暂无评论。