DatawhaleAl夏令营第四期大模型应用开发TASK 3:源的RAG

在自然语言处理的领域中，源大模型 RAG（Retrieval-Augmented Generation）正逐渐展现出其强大的能力和潜力。RAG 作为一种创新的技术架构，将信息检索与文本生成相结合，为解决各种自然语言处理任务提供了全新的思路和方法。在这篇博客中，我们将深入探讨 RAG 的概念，剖析一个完整的 RAG 链路，并详细阐述离线计算和在线计算在其中的应用。

一、什么是 RAG

RAG 是一种将检索机制与语言生成模型相结合的架构。传统的语言生成模型通常基于模型内部的知识和参数来生成文本，而 RAG 则引入了外部的知识源，通过检索相关的信息来增强生成的准确性和合理性。

具体来说，RAG 在生成文本时，首先会根据输入的提示或问题，在预先构建的知识库或文档集中进行检索，获取与当前任务相关的信息片段。然后，将检索到的信息与模型内部的语言理解和生成能力相结合，生成更加准确、详细且与上下文相关的回答。

这种结合检索和生成的方式，使得 RAG 能够利用大规模的外部知识，弥补了模型内部知识的局限性，从而在处理复杂和需要特定领域知识的任务时表现出色。

二、完整的 RAG 链路

一个完整的 RAG 链路通常包括以下几个关键步骤：

1. 数据准备：

   • 收集和整理大量的文本数据，构建知识库或文档集。这些数据可以来自各种来源，如互联网、书籍、论文等。

   • 对数据进行预处理，包括清洗、分词、标记化等操作，以便后续的检索和模型处理。

2. 索引构建：

   • 使用合适的索引技术，如倒排索引、向量索引等，对准备好的数据进行索引构建。索引的目的是能够快速有效地根据输入的查询检索到相关的文档或片段。

3. 模型训练：

   • 选择合适的语言模型，如 Transformer 架构的预训练语言模型（如 GPT 系列）。

   • 在训练过程中，不仅要训练模型的语言生成能力，还要使其能够与检索到的信息进行融合和交互。

4. 检索模块：

   • 当接收到输入的提示或问题时，检索模块根据索引快速查找相关的文档或片段。

   • 检索的策略可以基于关键词匹配、语义相似度计算等方法，以确保获取到最相关的信息。

5. 融合与生成：

   • 将检索到的信息与模型的内部表示进行融合。

   • 利用融合后的信息，模型生成最终的回答或文本输出。

三、离线计算

在 RAG 中，离线计算主要涉及数据准备、索引构建和模型的预训练等阶段。

1. 数据预处理和索引构建：

   • 对大规模的文本数据进行清洗、分词等预处理操作，这通常需要大量的计算资源和时间。

   • 构建高效的索引结构，如倒排索引或向量索引，以便在后续的检索中能够快速响应。这些操作可以在离线阶段一次性完成，然后在在线服务中直接使用。

2. 模型预训练：

   • 选择大规模的语料库对语言模型进行无监督的预训练。这是一个计算密集型的过程，需要在强大的计算集群上进行，花费较长的时间来学习语言的通用模式和知识。

通过离线计算，可以提前完成这些耗时且资源密集的任务，为在线服务提供高效的基础支持。

四、在线计算

在线计算在 RAG 中主要包括实时的检索和文本生成阶段。

1. 实时检索：

   • 当用户输入问题或提示时，系统需要在毫秒级的时间内从预先构建的索引中检索到相关的信息。这要求检索算法具有高效性和实时性，能够快速响应并返回准确的结果。

1. 文本生成：

   • 利用检索到的信息和模型的内部状态，实时生成回答或文本输出。这需要模型能够快速处理输入，并在短时间内生成高质量的文本。

在线计算需要系统具备低延迟、高并发处理能力，以满足用户的实时需求。

五、RAG实战

1.导入所需的库

# 导入所需的库
from typing import List
import numpy as np

import torch
from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM

2.定义向量模型类

# 定义向量模型类
class EmbeddingModel:
    """
    class for EmbeddingModel
    """

    def __init__(self, path: str) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(path)

        self.model = AutoModel.from_pretrained(path).cuda()
        print(f'Loading EmbeddingModel from {path}.')

    def get_embeddings(self, texts: List) -> List[float]:
        """
        calculate embedding for text list
        """
        encoded_input = self.tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
        encoded_input = {k: v.cuda() for k, v in encoded_input.items()}
        with torch.no_grad():
            model_output = self.model(**encoded_input)
            sentence_embeddings = model_output[0][:, 0]
        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
        return sentence_embeddings.tolist()

3.模型导入

print("> Create embedding model...")
embed_model_path = './AI-ModelScope/bge-small-zh-v1___5'
embed_model = EmbeddingModel(embed_model_path)

4.定义向量库索引类

# 定义向量库索引类
class VectorStoreIndex:
    """
    class for VectorStoreIndex
    """

    def __init__(self, doecment_path: str, embed_model: EmbeddingModel) -> None:
        self.documents = []
        for line in open(doecment_path, 'r', encoding='utf-8'):
            line = line.strip()
            self.documents.append(line)

        self.embed_model = embed_model
        self.vectors = self.embed_model.get_embeddings(self.documents)

        print(f'Loading {len(self.documents)} documents for {doecment_path}.')

    def get_similarity(self, vector1: List[float], vector2: List[float]) -> float:
        """
        calculate cosine similarity between two vectors
        """
        dot_product = np.dot(vector1, vector2)
        magnitude = np.linalg.norm(vector1) * np.linalg.norm(vector2)
        if not magnitude:
            return 0
        return dot_product / magnitude

    def query(self, question: str, k: int = 1) -> List[str]:
        question_vector = self.embed_model.get_embeddings([question])[0]
        result = np.array([self.get_similarity(question_vector, vector) for vector in self.vectors])
        return np.array(self.documents)[result.argsort()[-k:][::-1]].tolist() 

5.创造索引

print("> Create index...")
doecment_path = './knowledge.txt'
index = VectorStoreIndex(doecment_path, embed_model)

6.查询问题

print("> Create index...")
doecment_path = './knowledge.txt'
index = VectorStoreIndex(doecment_path, embed_model)

7.问题内容查询转换生成

question = '介绍一下广州大学'
print('> Question:', question)

context = index.query(question)
print('> Context:', context)

8.大语言模型类的定义

# 定义大语言模型类
class LLM:
    """
    class for Yuan2.0 LLM
    """

    def __init__(self, model_path: str) -> None:
        print("Creat tokenizer...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, add_eos_token=False, add_bos_token=False, eos_token='<eod>')
        self.tokenizer.add_tokens(['<sep>', '<pad>', '<mask>', '<predict>', '<FIM_SUFFIX>', '<FIM_PREFIX>', '<FIM_MIDDLE>','<commit_before>','<commit_msg>','<commit_after>','<jupyter_start>','<jupyter_text>','<jupyter_code>','<jupyter_output>','<empty_output>'], special_tokens=True)

        print("Creat model...")
        self.model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()

        print(f'Loading Yuan2.0 model from {model_path}.')

    def generate(self, question: str, context: List):
        if context:
            prompt = f'背景：{context}\n问题：{question}\n请基于背景，回答问题。'
        else:
            prompt = question

        prompt += "<sep>"
        inputs = self.tokenizer(prompt, return_tensors="pt")["input_ids"].cuda()
        outputs = self.model.generate(inputs, do_sample=False, max_length=1024)
        output = self.tokenizer.decode(outputs[0])

        print(output.split("<sep>")[-1])

9.选择大模型

print("> Create Yuan2.0 LLM...")
model_path = './IEITYuan/Yuan2-2B-Mars-hf'
# model_path = './IEITYuan/Yuan2-2B-July-hf'
llm = LLM(model_path)

10.答案生成

print('> Without RAG:')
llm.generate(question, [])

print('> With RAG:')
llm.generate(question, context)

六、总结

源大模型 的RAG 蛮有意思的，这是一次有意思的RAG实战