Atom - 树叶书屋

Atom-7B-chat

环境准备

在autodl平台中租一个3090等24G显存的显卡机器，如下图所示镜像选择PyTorch-->2.0.0-->3.8(ubuntu20.04)-->11.8

接下来打开刚刚租用服务器的JupyterLab，并且打开其中的终端开始环境配置、模型下载和运行demo。

pip换源和安装依赖包

# 升级pip
python -m pip install --upgrade pip
# 更换 pypi 源加速库的安装
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
#安装项目相关依赖包
pip install modelscope==1.9.5 transformers==4.35.2 gradio==4.4.1 SentencePiece==0.1.99 accelerate==0.24.1 bitsandbytes==0.41.2.post2

# 如遇 flash_attn 安装失败，请根据环境手动安装匹配的 wheel 包
# 例如：PyTorch 2.0.0 + Python 3.8 + CUDA 11.8（Ubuntu 20.04）
# 可直接使用下述预编译 wheel：
# 参考来源：https://github.com/Dao-AILab/flash-attention/releases
pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.6.3/flash_attn-2.6.3+cu118torch2.0cxx11abiFALSE-cp38-cp38-linux_x86_64.whl

# 其他环境请从 flash-attention Releases 页面选择与自身 CUDA / PyTorch / Python 版本匹配的 wheel
# 若无匹配 wheel，亦可参考官方文档源码编译安装

模型下载

使用 modelscope 中的snapshot_download函数下载模型，第一个参数为模型名称，参数cache_dir为模型的下载路径。

在 /root/autodl-tmp 路径下新建 download.py 文件并在其中输入以下内容，粘贴代码后记得保存文件，如下图所示。并运行 python /root/autodl-tmp/download.py执行下载，模型大小为 13 GB，下载模型大概需要 10~20 分钟

import torch
from modelscope import snapshot_download, AutoModel, AutoTokenizer
import os
model_dir = snapshot_download('FlagAlpha/Atom-7B-Chat', cache_dir='/root/autodl-tmp', revision='master')

代码准备

首先clone代码，打开autodl平台自带的学术镜像加速。学术镜像加速详细使用请看：https://www.autodl.com/docs/network_turbo/

source /etc/network_turbo

然后切换路径, clone代码.

cd /root/autodl-tmp
git clone https://github.com/FlagAlpha/Llama2-Chinese.git

切换commit版本，与教程commit版本保持一致，可以让大家更好的复现。

cd Llama2-Chinese
git checkout 0a2b588c5716f26f1e37affa308283354b3612be

最后取消镜像加速，因为该加速可能对正常网络造成一定影响，避免对后续下载其他模型造成困扰。

unset http_proxy && unset https_proxy

demo运行

进入代码目录,运行demo启动脚本，在--model_name_or_path 参数后填写下载的模型目录

cd /root/autodl-tmp/Llama2-Chinese/
python examples/chat_gradio.py --model_name_or_path /root/autodl-tmp/FlagAlpha/Atom-7B-Chat/

启动成功后终端显示如下：

设置代理访问

在Autodl容器实例页面找到自定义服务，下载对应的代理工具

启动代理工具，拷贝对应的ssh指令及密码，设置代理端口为7860，点击开始代理

代理成功后点击下方链接即可访问web-demo

Atom-7B-Chat 的 Lora 指令微调

概述

本节我们简要介绍如何基于 transformers、peft 等框架，对 Atom-7B-Chat 模型进行 Lora 微调。Lora 是一种高效微调方法，深入了解其原理可参见博客：知乎|深入浅出Lora。

本节所讲述的代码脚本在同级目录 02-Atom-7B-Chat Lora 下，可以通过运行目录下 train.sh 脚本来执行微调过程，但注意，本文代码未使用分布式框架，微调 Atom-7B 模型至少需要 32G 及以上的显存。

环境配置

在完成基本环境配置和本地模型部署的情况下，你还需要安装一些第三方库，可以使用以下命令：

pip install transformers==4.36.0.dev0
pip install peft==0.4.0.dev0
pip install datasets==2.10.1
pip install accelerate==0.20.3

在本节教程里，我们将微调数据集放置在根目录 /dataset，将基座模型参数放置在根目录 /model。

指令集构建

LLM 的微调一般指指令微调过程。所谓指令微调，是说我们使用的微调数据形如：

{
    "instrution":"回答以下用户问题，仅输出答案。",
    "input":"1+1等于几?",
    "output":"2"
}

其中，instruction 是用户指令，告知模型其需要完成的任务；input 是用户输入，是完成用户指令所必须的输入内容；output 是模型应该给出的输出。

即我们的核心训练目标是让模型具有理解并遵循用户指令的能力。因此，在指令集构建时，我们应针对我们的目标任务，针对性构建任务指令集。例如，在本节我们使用由笔者合作开源的 Chat-甄嬛项目作为示例，我们的目标是构建一个能够模拟甄嬛对话风格的个性化 LLM，因此我们构造的指令形如：

{
    "instruction": "请参考下面内容中的甄嬛的说话风格和语气，回答我的问题。甄嬛的说话风格需要是口语化的，回复内容不要超过30个字，尽可能字数简短一些。\n对话风格案例内容：\n```用户：这首歌虽未直写男女相悦，可字字写着两心相悦后女子的欢喜神态，而且'双双金鹧鹄'也是并蒂成双之意。\n用户：既然如此，安常在怎么就没唱出花好之情？难不成是看见皇上跟本宫在一起，心有不悦才唱不好的吗？\n甄嬛：回禀华妃娘娘，安常在早上受了风寒，嗓子有些不适。\n\n用户：你不是要看院子里的白梅吗，怎么那么快就回来了？\n甄嬛：雪景看久了反倒眼晕，四郎本是好意在园子里种植白梅，可是一下雪反倒与雪景融为一色，倒看不出来了。",
    "input":"你是谁？",
    "output":"家父是大理寺少卿甄远道。"

}

我们所构造的全部指令数据集在根目录下。

数据格式化

Lora 训练的数据是需要经过格式化、编码之后再输入给模型进行训练的，如果是熟悉 Pytorch 模型训练流程的同学会知道，我们一般需要将输入文本编码为 input_ids，将输出文本编码为 labels，编码之后的结果都是多维的向量。我们首先定义一个预处理函数，这个函数用于对每一个样本，编码其输入、输出文本并返回一个编码后的字典：

def preprocess(tokenizer, config, example, max_seq_length):
    '''
    args:
    tokenizer：分词器，导入的 Atom 模型分词器
    config：模型配置，导入的 Atom 模型配置
    example: 待处理的样本
    max_seq_length：文本的最大长度
    returns：字典，包括 inputs_id 和 seq_len
    '''
    # 将 instruction 和 input 按照 Atom SFT 时的格式拼接起来
    prompt = "<s>Human: " + example["instruction"] + "请回答用户问题: " + example["input"] + "\n" + "</s><s>Assistant:"
    target = example["output"]
    # 使用分词器进行编码，设置 truncation 为 True，避免出现过长的样本
    prompt_ids = tokenizer.encode(prompt, max_length=max_seq_length, truncation=True)
    target_ids = tokenizer.encode(
        target,
        max_length=max_seq_length,
        truncation=True,
        add_special_tokens=False)
    # 加入结束符 EOS
    input_ids = prompt_ids + target_ids + [config.eos_token_id]
    # 将 inputs_ids 和 seq_len 一起传回，后续会根据 seq_len 来切分 inputs 和 labels
    return {"input_ids": input_ids, "seq_len": len(prompt_ids)}

上述代码会对每一条样本进行格式化处理，我们再定义一个函数，这个函数基于上文函数，对源训练数据进行处理：

# 读取源训练数据并处理
def read_jsonl(path, max_seq_length, model_path, skip_overlength=False):
    '''
    args:
    path：训练数据路径
    max_seq_length：文本的最大长度
    model_path：模型路径，此处主要是为了加载分词器和配置
    returns：使用 yield 返回格式化的特征
    '''
    # 加载模型的分词器和配置参数
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_path, trust_remote_code=True)
    config = transformers.AutoConfig.from_pretrained(
        model_path, trust_remote_code=True, device_map='auto')
    # 读取源文件
    with open(path, "r") as f:
        # jsonl 数据需要先 readlines 读取成字符转，再使用 json 加载
        lst = [json.loads(line) for line in f.readlines()]
        print("加载jsonl数据集，数据总量为{}".format(len(lst)))
        # 依次处理每一个样本
        for example in tqdm(lst):
            # 调用上文的预处理函数
            feature = preprocess(tokenizer, config, example, max_seq_length)
            # 如果设置了跳过过长的样本
            if skip_overlength and len(feature["input_ids"]) > max_seq_length:
                continue
            # 截断过长的样本
            feature["input_ids"] = feature["input_ids"][:max_seq_length]
            # 通过 yield 返回迭代器
            yield feature

完成上述函数后，我们使用 datasets 库提供的 from_generator 函数来根据上述函数生成我们数据的 Dataset 对象：

# 通过 read_jsonl 函数返回的迭代器来生成 Dataset 对象，这个 Dataset 对象可以直接用在 transformers 框架中
dataset = datasets.Dataset.from_generator(
        lambda: read_jsonl(
            finetune_args.dataset_path, finetune_args.max_seq_length, finetune_args.model_path, finetune_args.skip_overlength
            )
    )

采样函数

为了对每一个 batch 的数据进行动态补齐，避免造成资源浪费，我们没有在生成函数中进行补齐操作，因此我们需要定义一个自定义采样函数，这个函数代替了 torch 中默认的采样函数功能，并自定义地实现了补齐、labels 遮蔽等操作，后续其会以 lambda 函数的方式传入 trainer：

# 自定义采样函数
def data_collator(features: list, tokenizer) -> dict:
    '''
    args:
    features: 一个批量的数据
    tokenizer：分词器
    returns：格式化的特征
    '''
    # 统计 batch 内所有数据的长度，将它们补齐
    len_ids = [len(feature["input_ids"]) for feature in features]
    # 补齐至最大长度
    longest = max(len_ids)
    # 分别存放 input_ids 和 labels
    input_ids = []
    labels_list = []
    # 有的模型没有定义 PAD，那么我们就用 UNK 来代替
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.unk_token_id
    # 从最长的文本开始处理，可以优化内存使用
    for ids_l, feature in sorted(zip(len_ids, features), key=lambda x: -x[0]):
        ids = feature["input_ids"]
        seq_len = feature["seq_len"]
        # labels 是将输入 PAD 之后保留输出的结果，用-100表示遮蔽，并且进行补齐，计算 loss 时会自动忽略 -100
        labels = (
            [-100] * (seq_len - 1) + ids[(seq_len - 1) :] + [-100] * (longest - ids_l)
        )
        ids = ids + [tokenizer.pad_token_id] * (longest - ids_l)
        _ids = torch.LongTensor(ids)
        labels_list.append(torch.LongTensor(labels))
        input_ids.append(_ids)
    # 在第0维进行拼接，也就是组成 batch_size*n*n 的矩阵
    input_ids = torch.stack(input_ids)
    labels = torch.stack(labels_list)
    return {
        "input_ids": input_ids,
        "labels": labels,
    }

自定义 Trainer

对于 Lora 微调，我们需要在基础 Trainer 的基础上继承一个自定义 Trainer，实现 Loss 计算（部分模型需要）和 Lora 参数的保存：

# 自定义 Trainer，继承自 transformers.trainer
class ModifiedTrainer(Trainer):
    # 重写损失计算函数，避免 LLaMA 类模型未定义 loss 的计算
    def compute_loss(self, model, inputs, return_outputs=False):
        # 7B
        return model(
            input_ids=inputs["input_ids"],
            labels=inputs["labels"],
        ).loss
    # 重写模型保存函数，从而保存模型的 Lora 参数
    def save_model(self, output_dir=None, _internal_call=False):
        from transformers.trainer import TRAINING_ARGS_NAME
        # 如果输出路径不存在，创建一个
        os.makedirs(output_dir, exist_ok=True)
        # 保存了模型训练的各种超参数
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        # 选出了所有梯度没有被冻结的参数，也就是所有参与更新的 Lora 参数
        saved_params = {
            k: v.to("cpu") for k, v in self.model.named_parameters() if v.requires_grad
        }
        # 保存所有 Lora 参数
        torch.save(saved_params, os.path.join(output_dir, "adapter_model.bin"))

参数解析

由于模型微调需要用到众多参数，最好以命令行方式传入，再通过 bash 脚本调用（例如 train.sh 脚本）。微调参数众多，transformers 提供了众多参数的解析，我们需要额外定义一个参数解析类，解析 transformers 没有提供的、用于 Lora 微调的参数解析：

# dataclass：Python 类修饰符，数据类，封装了__init__()、 __repr__()和__eq__()函数
@dataclass
class FinetuneArguments:
    # 微调参数
    # field：dataclass 函数，用于指定变量初始化
    # 训练集路径
    dataset_path: str = field(default="../../dataset/huanhuan.jsonl")
    # 基座模型参数路径
    model_path: str = field(default="../../dataset/model")
    # Lora 秩
    lora_rank: int = field(default=8)
    # 最大文本长度
    max_seq_length: int = field(default=256)
    # 是否跳过超长文本
    skip_overlength: bool = field(default=False)
    # 是否从断点继续训练
    continue_training: bool = field(default=False)
    # 断点路径，如果从断点继续训练需要传入
    checkpoint: str = field(default=None)

训练

完成上述定义和实现之后，我们可以正式开始我们的训练流程。首先我们需要解析脚本传入的训练参数，我们使用了 transformers 提供的 HfArgumentParser 函数，解析的参数包括 transformers 提供的 TrainingArguments 类（包括了一些常用训练参数）和我们自定义的 FinetuneArguments 类：

# 解析命令行参数
finetune_args, training_args = HfArgumentParser(
    (FinetuneArguments, TrainingArguments)
).parse_args_into_dataclasses()

接下来加载底座模型并进行一定的配置：

# 初始化底座模型 
tokenizer = AutoTokenizer.from_pretrained(finetune_args.model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    finetune_args.model_path, trust_remote_code=True, device_map="auto")
print("从{}加载模型成功".format(finetune_args.model_path))

# 启用梯度检查点，允许模型在前向计算时丢弃一些中间激活值，并在反向传播中重新计算，从而优化内存使用
model.gradient_checkpointing_enable()
# 确保输入向量能够计算梯度
model.enable_input_require_grads()
# 在训练过程中关闭缓存，提高计算效率，推理时应该开启
model.config.use_cache = (
    False  
)

然后设定 Lora 参数：

# 设定 peft 参数
# 手动确定 LoRA 层（注：理论上我们可以自动查找所有 Lora 层，但是在 LLaMA 类模型上出现 bug）
target_modules = ['W_pack', 'down_proj', 'o_proj', 'gate_proj', 'up_proj']
# 配置 Lora 参数
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM, # 任务为语言模型建模
    inference_mode=False, # 训练模式
    r=finetune_args.lora_rank, # Lora 秩
    lora_alpha=32, # Lora alaph，具体作用参见 Lora 原理
    lora_dropout=0.1,# Dropout 比例
    target_modules= target_modules # Lora 层
)

再基于 Lora 配置和底座模型，得到待训练的 Lora 模型（即冻结了非 Lora 层）。同时，需要判断是否是断点继续训练，如果是则要加载断点信息：

# 是否从断点继续训练
# 源点训练
if not finetune_args.continue_training:
    # 对基座模型进行 Lora 融合
    model = get_peft_model(model, peft_config)
    print("加载 LoRA 参数成功")
else:
    if finetune_args.check_point == None:
        print("断点训练需要给出 checkpoint 地址")
        raise ValueError("断点训练需要给出 checkpoint 地址")
    # 断点继续训练则直接加载断点的 Lora 参数
    model = PeftModel.from_pretrained(model, finetune_args.check_point, is_trainable=True)
    print("从{}加载断点成功".format(finetune_args.check_point))

然后基于上述定义加载数据集，在这一部分，我们使用了 try except 来捕捉异常：

# 加载数据集
try:
    # 调用上述定义函数生成迭代器
    dataset = datasets.Dataset.from_generator(
            lambda: read_jsonl(finetune_args.dataset_path, finetune_args.max_seq_length, finetune_args.model_path, finetune_args.skip_overlength)
        ) 
except Exception as e:
    print("从{}加载数据集失败".format(finetune_args.dataset_path))
    print("错误信息为：")
    print(e.__repr__())
    raise e   
print("从{}加载数据集成功".format(finetune_args.dataset_path))

最后，加载一个自定义的 trainer 并开始训练：

# 加载自定义 trainer
trainer = ModifiedTrainer(
    model=model, # 待训练模型
    train_dataset=dataset, # 数据集
    args=training_args, # 训练参数
    data_collator=lambda x : data_collator(x, tokenizer), # 自定义采样函数
)

print("成功加载 Trainer")
# 进行训练
trainer.train()
print("训练完成，训练结果保存在{}".format(training_args.output_dir))
# 保存模型
model.save_pretrained(training_args.output_dir)
print("模型参数保存在{}".format(training_args.output_dir))

通过上述代码，我们即可完成 Atom-7B-Chat 模型的 Lora 微调。我们将上述代码封装为 train.py 脚本，同时，提供一个启动训练的 bash 脚本：

python train.py \
    --dataset_path ../../dataset/huanhuan.jsonl \ # 数据集路径
    --model_path /root/autodl-tmp/data/model/Atom \ # 基座模型路径
    --lora_rank 8 \ # lora 秩
    --per_device_train_batch_size 16 \ # batch_size
    --gradient_accumulation_steps 1 \ # 梯度累积轮次
    --max_steps 120000 \ # 训练最大步数，训练 epoch 数 = max_steps / (num_whole_data / batch_size)
    --save_steps 40000 \ # 每训练多少步保存一次参数
    --save_total_limit 3 \ # 最多保存多少个参数
    --learning_rate 1e-4 \ # 学习率
    --fp16 \ # 使用 float16 的精度
    --remove_unused_columns false \ # 数据集处理时是否去除没有使用的特征
    --logging_steps 10 \ # 每训练多少步输出一次
    --output_dir ../../output # 输出路径

直接在目录下运行该脚本（bash train.sh）即可开始训练。

Atom-7B-Chat 接入langchain搭建知识库助手

环境准备

在autodl平台中租一个3090等24G显存的显卡机器，如下图所示镜像选择PyTorch-->2.0.0-->3.8(ubuntu20.04)-->11.8

接下来打开刚刚租用服务器的JupyterLab，并且打开其中的终端开始环境配置、模型下载和运行demo。

pip换源和安装依赖包

# 升级pip
python -m pip install --upgrade pip
# 更换 pypi 源加速库的安装
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

pip install modelscope==1.9.5
pip install "transformers>=4.32.0" accelerate tiktoken einops scipy transformers_stream_generator==0.0.4 peft deepspeed
pip install -U huggingface_hub
pip install -i https://test.pypi.org/simple/ bitsandbytes

模型下载

在已完成Atom-7B-chat部署的基础上，我们还需要还需要安装以下依赖包。请在终端复制粘贴以下命令，并按回车运行：

pip install langchain==0.0.292
pip install gradio==4.4.0
pip install chromadb==0.4.15
pip install sentence-transformers==2.2.2
pip install unstructured==0.10.30
pip install markdown==3.3.7

同时，我们还需要使用到开源词向量模型 Sentence Transformer

这里使用huggingface镜像下载到本地 /root/autodl-tmp/embedding_model，你也可以选择其它的方式下载

在 /root/autodl-tmp 路径下新建 download.py 文件并在其中输入以下内容，粘贴代码后记得保存文件，如下图所示。并运行 python /root/autodl-tmp/download.py执行下载。

import os
# 设置环境变量
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
# 下载模型
os.system('huggingface-cli download --resume-download sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 --local-dir /root/autodl-tmp/embedding_model')

使用 modelscope 中的snapshot_download函数下载模型，第一个参数为模型名称，参数cache_dir为模型的下载路径。

在 /root/autodl-tmp 路径下新建 download.py 文件并在其中输入以下内容，粘贴代码后记得保存文件，如下图所示。并运行 python /root/autodl-tmp/download.py执行下载，模型大小为 13 GB，下载模型大概需要 10~20 分钟

import torch
from modelscope import snapshot_download, AutoModel, AutoTokenizer
import os
model_dir = snapshot_download('FlagAlpha/Atom-7B-Chat', cache_dir='/root/autodl-tmp', revision='master')

知识库建设

我们选用以下两个开源仓库作为知识库来源

Atom-7B-Chat
Llama2-Chinese

首先我们需要将上述远程开源仓库 Clone 到本地，可以使用以下命令：

# 进入到数据库盘
cd /root/autodl-tmp
# 打开学术资源加速
source /etc/network_turbo
# clone 开源仓库
git clone https://github.com/FlagAlpha/Llama2-Chinese
# 关闭学术资源加速
unset http_proxy && unset https_proxy

接着，为语料处理方便，我们将选用上述仓库中所有的 markdown、txt 文件作为示例语料库。注意，也可以选用其中的代码文件加入到知识库中，但需要针对代码文件格式进行额外处理。

我们首先将上述仓库中所有满足条件的文件路径找出来，我们定义一个函数，该函数将递归指定文件夹路径，返回其中所有满足条件（即后缀名为 .md 或者 .txt 的文件）的文件路径：

import os 
def get_files(dir_path):
    # args：dir_path，目标文件夹路径
    file_list = []
    for filepath, dirnames, filenames in os.walk(dir_path):
        # os.walk 函数将递归遍历指定文件夹
        for filename in filenames:
            # 通过后缀名判断文件类型是否满足要求
            if filename.endswith(".md"):
                # 如果满足要求，将其绝对路径加入到结果列表
                file_list.append(os.path.join(filepath, filename))
            elif filename.endswith(".txt"):
                file_list.append(os.path.join(filepath, filename))
    return file_list

得到所有目标文件路径之后，我们可以使用 LangChain 提供的 FileLoader 对象来加载目标文件，得到由目标文件解析出的纯文本内容。由于不同类型的文件需要对应不同的 FileLoader，我们判断目标文件类型，并针对性调用对应类型的 FileLoader，同时，调用 FileLoader 对象的 load 方法来得到加载之后的纯文本对象：

from tqdm import tqdm
from langchain.document_loaders import UnstructuredFileLoader
from langchain.document_loaders import UnstructuredMarkdownLoader

def get_text(dir_path):
    # args：dir_path，目标文件夹路径
    # 首先调用上文定义的函数得到目标文件路径列表
    file_lst = get_files(dir_path)
    # docs 存放加载之后的纯文本对象
    docs = []
    # 遍历所有目标文件
    for one_file in tqdm(file_lst):
    file_type = one_file.split('.')[-1]
        if file_type == 'md':
            loader = UnstructuredMarkdownLoader(one_file)
        elif file_type == 'txt':
            loader = UnstructuredFileLoader(one_file)
        else:
            # 如果是不符合条件的文件，直接跳过
            continue
        docs.extend(loader.load())
    return docs

使用上文函数，我们得到的 docs 为一个纯文本对象对应的列表。得到该列表之后，我们就可以将它引入到 LangChain 框架中构建向量数据库。由纯文本对象构建向量数据库，我们需要先对文本进行分块，接着对文本块进行向量化。

LangChain 提供了多种文本分块工具，此处我们使用字符串递归分割器，并选择分块大小为 500，块重叠长度为 150：

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500, chunk_overlap=150)
split_docs = text_splitter.split_documents(docs)

接着我们选用开源词向量模型 Sentence Transformer 来进行文本向量化

LangChain 提供了直接引入 HuggingFace 开源社区中的模型进行向量化的接口：

from langchain.embeddings.huggingface import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(model_name="/root/autodl-tmp/embedding_model")

同时，我们选择 Chroma 作为向量数据库，基于上文分块后的文档以及加载的开源向量化模型，将语料加载到指定路径下的向量数据库：

from langchain.vectorstores import Chroma

# 定义持久化路径
persist_directory = 'data_base/vector_db/chroma'
# 加载数据库
vectordb = Chroma.from_documents(
    documents=split_docs,
    embedding=embeddings,
    persist_directory=persist_directory  # 允许我们将persist_directory目录保存到磁盘上
)
# 将加载的向量数据库持久化到磁盘上
vectordb.persist()

将上述代码整合在一起为知识库搭建的脚本：

# 首先导入所需第三方库
from langchain.document_loaders import UnstructuredFileLoader
from langchain.document_loaders import UnstructuredMarkdownLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from tqdm import tqdm
import os
import nltk
nltk.download('punkt')

# 获取文件路径函数
def get_files(dir_path):
    # args：dir_path，目标文件夹路径
    file_list = []
    for filepath, dirnames, filenames in os.walk(dir_path):
        # os.walk 函数将递归遍历指定文件夹
        for filename in filenames:
            # 通过后缀名判断文件类型是否满足要求
            if filename.endswith(".md"):
                # 如果满足要求，将其绝对路径加入到结果列表
                file_list.append(os.path.join(filepath, filename))
            elif filename.endswith(".txt"):
                file_list.append(os.path.join(filepath, filename))
    return file_list

# 加载文件函数
def get_text(dir_path):
    # args：dir_path，目标文件夹路径
    # 首先调用上文定义的函数得到目标文件路径列表
    file_lst = get_files(dir_path)
    # docs 存放加载之后的纯文本对象
    docs = []
    # 遍历所有目标文件
    for one_file in tqdm(file_lst):
        file_type = one_file.split('.')[-1]
        if file_type == 'md':
            loader = UnstructuredMarkdownLoader(one_file)
        elif file_type == 'txt':
            loader = UnstructuredFileLoader(one_file)
        else:
            # 如果是不符合条件的文件，直接跳过
            continue
        docs.extend(loader.load())
    return docs

# 目标文件夹
tar_dir = [
    "/root/autodl-tmp/Llama2-Chinese",
]

# 加载目标文件
docs = []
for dir_path in tar_dir:
    docs.extend(get_text(dir_path))

# 对文本进行分块
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500, chunk_overlap=150)
split_docs = text_splitter.split_documents(docs)

# 加载开源词向量模型
embeddings = HuggingFaceEmbeddings(model_name="/root/autodl-tmp/embedding_model")

# 构建向量数据库
# 定义持久化路径
persist_directory = 'data_base/vector_db/chroma'
# 加载数据库
vectordb = Chroma.from_documents(
    documents=split_docs,
    embedding=embeddings,
    persist_directory=persist_directory  # 允许我们将persist_directory目录保存到磁盘上
)
# 将加载的向量数据库持久化到磁盘上
vectordb.persist()

运行上述脚本，即可在本地构建已持久化的向量数据库，后续直接导入该数据库即可，无需重复构建。

Atom 接入LangChain

为便捷构建 LLM 应用，我们需要基于本地部署的 Atom，自定义一个 LLM 类，将 Atom 接入到 LangChain 框架中。完成自定义 LLM 类之后，可以以完全一致的方式调用 LangChain 的接口，而无需考虑底层模型调用的不一致。

基于本地部署的 Atom 自定义 LLM 类并不复杂，我们只需从 LangChain.llms.base.LLM 类继承一个子类，并重写构造函数与 _call 函数即可：

from langchain.llms.base import LLM
from typing import Any, List, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class Atom(LLM):
    # 基于本地 Atom 自定义 LLM 类
    tokenizer : AutoTokenizer = None
    model: AutoModelForCausalLM = None

    def __init__(self, model_path :str):
        # model_path: Atom 模型路径
        # 从本地初始化模型
        super().__init__()
        print("正在从本地加载模型...")
        model_dir = '/root/autodl-tmp/FlagAlpha/Atom-7B-Chat'
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True,torch_dtype=torch.float16,load_in_8bit=True).eval()
        print("完成本地模型的加载")

    def _call(self, prompt: str, stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any):
        input_ids = self.tokenizer([f'<s>Human: {prompt}\n</s><s>Assistant: '], return_tensors="pt", add_special_tokens=False).input_ids.to('cuda')
        generate_input = {
            "input_ids": input_ids,
            "max_new_tokens": 512,
            "do_sample": True,
            "top_k": 50,
            "top_p": 0.95,
            "temperature": 0.3,
            "repetition_penalty": 1.3,
            "eos_token_id": self.tokenizer.eos_token_id,
            "bos_token_id": self.tokenizer.bos_token_id,
            "pad_token_id": self.tokenizer.pad_token_id
        }
        generate_ids = self.model.generate(**generate_input)
        text = self.tokenizer.decode(generate_ids[0])
        return text

    @property
    def _llm_type(self) -> str:
        return "Atom"

在上述类定义中，我们分别重写了构造函数和 _call 函数：对于构造函数，我们在对象实例化的一开始加载本地部署的 Atom 模型，从而避免每一次调用都需要重新加载模型带来的时间过长；_call 函数是 LLM 类的核心函数，LangChain 会调用该函数来调用 LLM，在该函数中，我们调用已实例化模型的 generate 方法，从而实现对模型的调用并返回调用结果。

在整体项目中，我们将上述代码封装为 LLM.py，后续将直接从该文件中引入自定义的 LLM 类。

构建检索问答链

LangChain 通过提供检索问答链对象来实现对于 RAG 全流程的封装。即我们可以调用一个 LangChain 提供的 RetrievalQA 对象，通过初始化时填入已构建的数据库和自定义 LLM 作为参数，来简便地完成检索增强问答的全流程，LangChain 会自动完成基于用户提问进行检索、获取相关文档、拼接为合适的 Prompt 并交给 LLM 问答的全部流程。

首先我们需要将上文构建的向量数据库导入进来，我们可以直接通过 Chroma 以及上文定义的词向量模型来加载已构建的数据库：

from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
import os

# 定义 Embeddings
embeddings = HuggingFaceEmbeddings(model_name="/root/autodl-tmp/embedding_model")

# 向量数据库持久化路径
persist_directory = 'data_base/vector_db/chroma'

# 加载数据库
vectordb = Chroma(
    persist_directory=persist_directory, 
    embedding_function=embeddings
)

上述代码得到的 vectordb 对象即为我们已构建的向量数据库对象，该对象可以针对用户的 query 进行语义向量检索，得到与用户提问相关的知识片段。

接着，我们实例化一个基于 Atom 自定义的 LLM 对象：

from LLM import Atom
llm = Atom(model_path = "/root/autodl-tmp/FlagAlpha/Atom-7B-Chat")
llm.predict("你是谁")

构建检索问答链，还需要构建一个 Prompt Template，该 Template 其实基于一个带变量的字符串，在检索之后，LangChain 会将检索到的相关文档片段填入到 Template 的变量中，从而实现带知识的 Prompt 构建。我们可以基于 LangChain 的 Template 基类来实例化这样一个 Template 对象：

from langchain.prompts import PromptTemplate

# 我们所构造的 Prompt 模板
template = """使用以下上下文来回答最后的问题。如果你不知道答案，就说你不知道，不要试图编造答案。尽量使答案简明扼要。总是在回答的最后说“谢谢你的提问！”。
{context}
问题: {question}
有用的回答:"""

# 调用 LangChain 的方法来实例化一个 Template 对象，该对象包含了 context 和 question 两个变量，在实际调用时，这两个变量会被检索到的文档片段和用户提问填充
QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context","question"],template=template)

最后，可以调用 LangChain 提供的检索问答链构造函数，基于我们的自定义 LLM、Prompt Template 和向量知识库来构建一个基于 Atom 的检索问答链：

from langchain.chains import RetrievalQA

qa_chain = RetrievalQA.from_chain_type(llm,retriever=vectordb.as_retriever(),return_source_documents=True,chain_type_kwargs={"prompt":QA_CHAIN_PROMPT})

得到的 qa_chain 对象即可以实现我们的核心功能，即基于 Atom 模型的专业知识库助手。我们可以对比该检索问答链和纯 LLM 的问答效果：

question = "什么是Atom"
result = qa_chain({"query": question})
print("检索问答链回答 question 的结果：")
print(result["result"])

# 仅 LLM 回答效果
result_2 = llm(question)
print("大模型回答 question 的结果：")
print(result_2)

Alt text 可以看到，使用检索问答链生成的答案更接近知识库里的内容。

部署WebDemo

在完成上述核心功能后，我们可以基于 Gradio 框架将其部署到 Web 网页，从而搭建一个小型 Demo，便于测试与使用。

我们首先将上文的代码内容封装为一个返回构建的检索问答链对象的函数，并在启动 Gradio 的第一时间调用该函数得到检索问答链对象，后续直接使用该对象进行问答对话，从而避免重复加载模型：

# 导入必要的库
import gradio as gr
from langchain.vectorstores import Chroma
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
import os
from LLM import Atom
from langchain.prompts import PromptTemplate

def load_chain():
    # 加载问答链
    # 定义 Embeddings
    embeddings = HuggingFaceEmbeddings(model_name="/root/autodl-tmp/embedding_model")

    # 向量数据库持久化路径
    persist_directory = 'data_base/vector_db/chroma'

    # 加载数据库
    vectordb = Chroma(
        persist_directory=persist_directory,  # 允许我们将persist_directory目录保存到磁盘上
        embedding_function=embeddings
    )

    llm = Atom(model_path = "/root/autodl-tmp/FlagAlpha")

    template = """使用以下上下文来回答最后的问题。如果你不知道答案，就说你不知道，不要试图编造答
    案。尽量使答案简明扼要。总是在回答的最后说“谢谢你的提问！”。
    {context}
    问题: {question}
    有用的回答:"""

    QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context","question"],
                                    template=template)

    # 运行 chain
    from langchain.chains import RetrievalQA

    qa_chain = RetrievalQA.from_chain_type(llm,
                                        retriever=vectordb.as_retriever(),
                                        return_source_documents=True,
                                        chain_type_kwargs={"prompt":QA_CHAIN_PROMPT})
    
    return qa_chain

接着我们定义一个类，该类负责加载并存储检索问答链，并响应 Web 界面里调用检索问答链进行回答的动作：

class Model_center():
    """
    存储问答 Chain 的对象 
    """
    def __init__(self):
        self.chain = load_chain()

    def qa_chain_self_answer(self, question: str, chat_history: list = []):
        """
        调用不带历史记录的问答链进行回答
        """
        if question == None or len(question) < 1:
            return "", chat_history
        try:
            chat_history.append(
                (question, self.chain({"query": question})["result"]))
            return "", chat_history
        except Exception as e:
            return e, chat_history

    def clear_history(self):
        self.chain.clear_history()

然后我们只需按照 Gradio 的框架使用方法，实例化一个 Web 界面并将点击动作绑定到上述类的回答方法即可：

import gradio as gr
model_center = Model_center()

block = gr.Blocks()
with block as demo:
    with gr.Row(equal_height=True):   
        with gr.Column(scale=15):
            gr.Markdown("""<h1><center>Atom</center></h1>
                <center>Llama2-chinese(教程来自DataWhale Self-LM团队)</center>
                """)
        # gr.Image(value=LOGO_PATH, scale=1, min_width=10,show_label=False, show_download_button=False)

    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(height=450, show_copy_button=True)
            # 创建一个文本框组件，用于输入 prompt。
            msg = gr.Textbox(label="Prompt/问题")

            with gr.Row():
                # 创建提交按钮。
                db_wo_his_btn = gr.Button("Chat")
            with gr.Row():
                # 创建一个清除按钮，用于清除聊天机器人组件的内容。
                clear = gr.ClearButton(
                    components=[chatbot], value="Clear console")
                
        # 设置按钮的点击事件。当点击时，调用上面定义的 qa_chain_self_answer 函数，并传入用户的消息和聊天历史记录，然后更新文本框和聊天机器人组件。
        db_wo_his_btn.click(model_center.qa_chain_self_answer, inputs=[msg, chatbot], outputs=[msg, chatbot])
        
        # 点击后清空后端存储的聊天记录
        clear.click(model_center.clear_history)
    gr.Markdown("""提醒：<br>
    1. 初始化数据库时间可能较长，请耐心等待。
    2. 使用中如果出现异常，将会在文本输入框进行展示，请不要惊慌。 <br>
    """)
# threads to consume the request
gr.close_all()
# 启动新的 Gradio 应用，设置分享功能为 True，并使用环境变量 PORT1 指定服务器端口。
# demo.launch(share=True, server_port=int(os.environ['PORT1']))
# 直接启动
demo.launch()

通过将上述代码封装为 run_gradio.py 脚本，直接通过在终端运行命令 python run_gradio.py ，即可在本地启动知识库助手的 Web Demo，默认会在 7860 端口运行，使用类似于部署的方式将服务器端口映射到本地端口即可访问:

Atom-7B-chat 全量微调

修改代码

首先我们要准备训练模型的代码，这里我们使用的 modelscope 上的 Atom-7B-chat 模型，大家自行下载即可。

OK，模型下载完毕之后，我们就要准备代码文件。其实全量微调和 Lora 微调的代码基本一样，都采用了 Trainer 类来进行训练。只不过在全量微调的时候没有加载 LoraConfig，那我就直接给出代码，如果对代有什么问题，大家可以先自行探索Qwen lora的代码解释，有什么不懂的地方可以提Issue。

需要把代码中的模型地址修改一下，改成自己的模型地址。

from datasets import Dataset
import pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, HfArgumentParser, Trainer
import os
import torch
from dataclasses import dataclass, field
import deepspeed
deepspeed.ops.op_builder.CPUAdamBuilder().load()


@dataclass
class FinetuneArguments:
    # 微调参数
    # field：dataclass 函数，用于指定变量初始化
    model_path: str = field(default="../../model/FlagAlpha/Atom-7B-Chat/")

# 用于处理数据集的函数
def process_func(example):
    MAX_LENGTH = 128    # Llama分词器会将一个中文字切分为多个token，因此需要放开一些最大长度，保证数据的完整性
    input_ids, attention_mask, labels = [], [], []
    instruction = tokenizer("\n".join(["<s>Human:", "现在你要扮演皇帝身边的女人--甄嬛，请以甄嬛口吻回答用户问题：" + example["instruction"] + example["input"] + "</s>\n"]).strip(), add_special_tokens=False)  # add_special_tokens 不在开头加 special_tokens
    response = tokenizer("<s>Assistant:" + example["output"] + "<\s>\n", add_special_tokens=False)
    input_ids = instruction["input_ids"] + response["input_ids"]
    attention_mask = instruction["attention_mask"] + response["attention_mask"]
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"]
    if len(input_ids) > MAX_LENGTH:  # 做一个截断
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }


if "__main__" == __name__:
    # 解析参数
    # Parse 命令行参数
    finetune_args, training_args = HfArgumentParser(
        (FinetuneArguments, TrainingArguments)
    ).parse_args_into_dataclasses()

    # 处理数据集
    # 将JSON文件转换为CSV文件
    df = pd.read_json('./data/huanhuan.json')
    ds = Dataset.from_pandas(df)
    # 加载tokenizer
    tokenizer = AutoTokenizer.from_pretrained(finetune_args.model_path, use_fast=False, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    # 将数据集变化为token形式
    tokenized_id = ds.map(process_func, remove_columns=ds.column_names)

    # 创建模型并以半精度形式加载
    model = AutoModelForCausalLM.from_pretrained(finetune_args.model_path, trust_remote_code=True, torch_dtype=torch.half, device_map={"": int(os.environ.get("LOCAL_RANK") or 0)})
    
    # 使用trainer训练
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_id,
        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
        )
    trainer.train() # 开始训练
    response, history = model.chat(tokenizer, "你是谁", history=[], system="现在你要扮演皇帝身边的女人--甄嬛.")
    print(response)

DeepSpeed 环境配置

DeepSpeed 是微软开源的一个深度学习训练框架，可以用于分布式训练，同时还可以加速训练，减少显存占用。这里我们使用的是 DeepSpeed 的半精度训练，可以减少显存占用，加快训练速度。

首先我们需要安装 DeepSpeed，DeepSpeed 的安装很简单，但如果没有按照如下步骤安装，可能会出现一些问题。

首先创建一个崭新的，干净的conda环境，注意一定要使用当前目录下提供的environment.yml文件来创建环境，否则可能会出现一些问题。接着激活环境，安装deepspeed，使用DS_BUILD_OPS=1来安装deepspeed，这样会避免后续的很多报错。

conda env create -n deepspeed -f environment.yml --force
conda activate deepspeed 
DS_BUILD_OPS=1 pip install deepspeed

然后就是安装transformers等其他依赖，注意不需要再安装torch了，在创建环境的时候torch已经安装了。

pip install modelscope==1.9.5
pip install transformers==4.35.2
pip install streamlit==1.24.0
pip install sentencepiece==0.1.99
pip install accelerate==0.24.1
pip install datasets sentencepiece
pip install tiktoken
pip install transformers_stream_generator

注意：本环境是在aws服务器上安装并运行的，假如您在安装或者运行过程中遇到其他问题，欢迎提出issue，然后您解决之后，可以顺便提交PR，为项目添砖加瓦。

模型训练

首先创建deepspeed的config.json文件。我使用的是stage-2的配置。如果不懂也没关系，直接粘贴复制，创建为ds_config.json文件即可。

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "last_batch_iteration": -1,
            "total_num_steps": "auto",
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "contiguous_gradients": true
    },
    "activation_checkpointing": {
        "partition_activations": false,
        "cpu_checkpointing": false,
        "contiguous_memory_optimization": false,
        "number_checkpoints": null,
        "synchronize_checkpoint_boundary": false,
        "profile": false
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "min_lr": 5e-7,
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

然后我们来创建运行所需的bash脚本，创建一个train.sh文件，内容如下：

num_gpus=4

deepspeed --num_gpus $num_gpus train.py \
    --deepspeed ./ds_config.json \
    --output_dir="./output/Atom" \
    --per_device_train_batch_size=1 \
    --gradient_accumulation_steps=1 \
    --logging_steps=10 \
    --num_train_epochs=3 \
    --save_steps=100 \
    --learning_rate=1e-4 \
    --save_on_each_node=True \

接着在命令行输入：bash train.sh，开始训练。

注意：

因为本脚本使用了adam_cpu来加载优化器参数，所以全量微调所需的显存会比较小，但仍然需要使用至少4张24G显存的卡来训练。
如果第一步创建deepspeed环境时候，没有使用DS_BUILD_OPS=1，那么可能会出现一些问题，比如RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!，这个时候需要重新创建环境，然后再次运行。