【2024 Datawhale AI 夏令营】Task3:RAG实战

【2024 Datawhale AI 夏令营】Task3:RAG实战

Simuoss
2024-08-17 / 0 评论 / 14 阅读 / 正在检测是否收录...

RAG 技术简介

技术用途

RAG(检索增强生成)是为解决基础大模型在实际业务场景中所面临的一些局限性问题而提出的一种方法。主要解决的问题包括:

  1. 知识局限性:大模型的知识来源于训练数据,无法覆盖实时性或非公开的数据。
  2. 数据安全性:企业需要保护数据安全,而大模型训练往往需要接触大量数据,这可能带来数据泄露的风险。
  3. 大模型幻觉:大模型基于概率生成回答,有时会在没有知识的领域生成不准确的内容。

RAG 通过引入外部知识库来增强大模型的生成能力,通常分为以下三个步骤:

  • 索引:将文档分割为文本块,并构建向量索引。
  • 检索:根据问题检索与其相似的文本块。
  • 生成:将检索到的文本块作为背景信息,用于生成更加准确的回答。

当前,RAG 已成为解决知识盲区和生成准确答案的主流方法,在实际应用中非常有效。

构建流程

离线计算:

离线计算是 RAG 系统的基础,主要工作是对知识库中的文档进行解析、清理、切割、向量化,并将这些向量存储到数据库中,具体包括以下步骤:

  1. 文档解析与清洗:知识库包含多种文档格式(如 PDF、Word、PPT 等)。首先对文档进行解析,清理无关内容,切割成较短的文本块(Chunk),同时去除重复内容。这样可以提高系统的知识覆盖率和查询的准确性。
  2. 向量化(Vectorization):每个文本块会被转换为一个向量,这需要一个预训练的向量模型(Embedding Model)。该模型可以将具有相似语义的文本映射到语义空间中相近的位置,而语义不同的文本则会被映射到更远的地方。
  3. 存储:由于知识库的规模较大,向量化后的结果会存储在向量数据库中,如 Milvus、Faiss 等。向量数据库有助于高效管理这些向量,并在后续的检索过程中提高性能。

离线计算的核心在于提前做好大规模的计算和存储准备,确保在线计算时可以高效、准确地响应用户查询。


在线计算:

  1. 用户查询与检索:用户输入查询时,系统首先会将该查询向量化,并与数据库中的向量进行相似度计算,检索出一系列与查询最相关的文本块。这一过程的效率与数据库规模相关,因此可能会使用召回机制来缩小计算范围。
  2. 召回与精排:召回阶段快速获取大概率相关的文本块,通常采用简单的基于字符串的匹配算法(如 TF-IDF、BM25 等),然后在这些文本块中进一步进行精排,即基于向量相似度对文本块排序。精排确保相关性更高的文本块排在前列。
  3. 重排(Rerank):为提高检索精度,在精排之后还可以进行重排。重排模型通过更复杂的算法对精排结果重新排序,确保最相关的内容优先被选中。随着知识库规模的增加,重排能够有效提升整体检索效果。
  4. 生成回复:最终,将检索到的 k 个最相关文本块与用户的查询拼接成一个提示(prompt),输入到大模型中。大模型基于这些背景信息生成更加准确的回答,解决用户的问题。

通过离线与在线计算的配合,RAG 系统能够在保证数据安全的前提下,提供快速且精确的问答服务。

场景举例:

法律咨询服务系统

  • 离线计算:解析、切割和向量化法律法规、合同模板等法律文档,并将它们存入数据库。
  • 在线计算:用户输入具体问题(如“租赁合同中的解除条款是什么?”),系统检索相关法律条文和合同条款,将其作为背景信息,结合大模型生成具体的法律建议。

实战体验

注册社区

注册社区启动实例的部分可以参照前几天发布的[【2024 Datawhale AI 夏令营】Task1:借助魔塔(ModelScope)社区基于[源-2B]模型构建智能编程助手](https://blog.simuoss.cn/index.php/archives/15/)操作。

搭建Demo环境

进入实例,点击终端。
运行下面代码,下载文件,并将 Task 3:源大模型RAG实战中的内容拷贝到当前目录。

git lfs install
git clone https://www.modelscope.cn/datasets/Datawhale/AICamp_yuan_baseline.git
cp AICamp_yuan_baseline/Task\ 3:源大模型RAG实战/* .

双击打开 Task 3:源大模型RAG实战.ipynb,然后运行所有单元格。
通过下面的命令,我们可以看到 ModelScope 已经提供了所需的大部分依赖,如 torchtransformers 等。

pip list

但是为了进行模型微调以及 Demo 搭建,还需要在环境中安装 streamlit

pip install streamlit==1.24.0

安装成功后,我们的环境就准备好了。

模型下载

在 RAG 实战中,我们需要构建一个向量模型。

向量模型通常采用 BERT 架构,它是一个 Transformer Encoder。
输入向量模型前,首先会在文本的最前面额外加一个 [CLS] token,然后将该 token 最后一层的隐藏层向量作为文本的表示。

在本次学习中,我们选用基于 BERT 架构的向量模型 bge-small-zh-v1.5,它是一个 4 层的 BERT 模型,最大输入长度 512,输出的向量维度也为 512。

向量模型下载:

from modelscope import snapshot_download
model_dir = snapshot_download("AI-ModelScope/bge-small-zh-v1.5", cache_dir='.')

这里使用的是 modelscope 中的 snapshot_download 函数,第一个参数为模型名称 AI-ModelScope/bge-small-zh-v1.5,第二个参数 cache_dir 为模型保存路径,这里 . 表示当前路径。
模型大小约为 91.4M,由于是从魔搭直接进行下载,速度会非常快。
下载完成后,会在当前目录增加一个名为 AI-ModelScope 的文件夹,其中 bge-small-zh-v1___5 里面保存着我们下载好的向量模型。

另外,还需要下载源大模型 IEITYuan/Yuan2-2B-Mars-hf。下载方法和 Task 1:零基础玩转源大模型 一致。

from modelscope import snapshot_download
model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='.')

RAG实战

模型下载完成后,就可以开始 RAG 实战了。

索引

为了构造索引,这里我们封装了一个向量模型类 EmbeddingModel

# 定义向量模型类
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()

通过传入模型路径,新建一个 EmbeddingModel 对象 embed_model
初始化时自动加载向量模型的 tokenizer 和模型参数。

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

EmbeddingModel 类还有一个 get_embeddings() 函数,它可以获得输入文本的向量表示。
注意,这里为了充分发挥 GPU 矩阵计算的优势,输入和输出都是一个 List,即多条文本及其向量表示。

2.3.2 检索

为了实现向量检索,我们定义了一个向量库索引类 VectorStoreIndex

# 定义向量库索引类
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()

通过传入知识库文件路径,新建一个 VectorStoreIndex 对象 index
初始化时会自动读取知识库的内容,然后传入向量模型,获得向量表示。

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

VectorStoreIndex 类还有一个 get_similarity() 函数,用于计算两个向量之间的相似度,采用余弦相似度。

VectorStoreIndex 类的入口是查询函数 query()。传入用户的提问后,首先会送入向量模型获得其向量表示,然后与知识库中的所有向量计算相似度,最后将 k 个最相似的文档按顺序返回,k 默认为 1。

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

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

生成

为了实现基于 RAG 的生成,我们还需要定义一个大语言模型类 LLM

# 定义大语言模型类
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])

我们传入 Yuan2-2B-Mars 的模型路径,新建一个 LLM 对象 llm
初始化时自动加载源大模型的 tokenizer 和模型参数。

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

LLM 类的入口是生成函数 generate(),它有两个参数:

  • question: 用户提问,是一个 str
  • context: 检索到的上下文信息,是一个 List,默认是 [],代表没有使用 RAG

运行下面的代码,即可体验使用 RAG 技术之后 Yuan2-2B-Mars 模型的回答效果:

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

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

参考链接:
Datawhale官方教程

1

评论 (0)

取消