ssbuild大佬的chatglm_finetuning项目---data_utils.py代码解读

# @Time    : 2023/1/22 16:22
# @Author  : tk
# @FileName: data_utils.py


import copy
import json
import os
import random
import typing
from enum import Enum

import numpy as np
import torch
from deep_training.data_helper import DataHelper, ModelArguments, TrainingArguments, DataArguments
from deep_training.nlp.models.chatglm import ChatGLMConfig
from deep_training.nlp.models.lora import LoraArguments
from deep_training.utils.func import is_chinese_char
from fastdatasets.record import load_dataset as Loader, RECORD, WriterObject, gfile
from tqdm import tqdm
from transformers import HfArgumentParser

from data_processer import DataStrategy, TokenTruncation, TokenSingleSliding, TokenDoubleSliding
from models import ChatGLMTokenizer

train_info_args = {
    'devices': 1,
    'data_backend': 'record',
    'model_type': 'chatglm',
    # 预训练模型路径 , 从0训练，则置空
    'model_name_or_path': './model/chatglm-6b',
    'config_name': './config/config_small.json',
    'tokenizer_name': './model/chatglm-6b',
    'convert_onnx': False, # 转换onnx模型
    'do_train': True,
    'train_file':  [ './data/finetune_train_examples.json', './data/finetune_train_examples.json'],
    'max_epochs': 1,
    'max_steps': -1,
    'optimizer': 'lion', # one of adamw,adam,lamb,lion

    # 学习率调度器（scheduler），其中'scheduler_type'指定了使用哪种类型的scheduler，而'scheduler'参数指定了对应类型的scheduler所需的参数。
    # 带慢启动的余弦学习率调度器
    'scheduler_type': 'CAWR',
    # 'T_mult': 一个整数，控制多少个周期后重新开始下一次周期。默认值为1。
    # 'rewarm_epoch_num': 一个浮点数，表示一个完整周期的比例，以epoch为单位。默认值为0.5，即一半的周期用于重新热启动。
    # 'verbose': 表示是否打印详细信息。默认值为False。
    'scheduler':{'T_mult': 1, 'rewarm_epoch_num': 0.5, 'verbose': True},

    # linear线性方法来调整学习率。
    # 'scheduler_type': 'linear',# one of [linear,WarmupCosine,CAWR,CAL,Step,ReduceLROnPlateau
    # 'scheduler': None,

    # 切换scheduler类型
    # 'scheduler_type': 'WarmupCosine',
    # 'scheduler': None,

    # 该scheduler在验证集上监测到学习率的性能不佳时，将学习率降低一个因子。例如，如果在验证集上的准确率停滞不前，则将学习率除以10
    # 'scheduler_type': 'ReduceLROnPlateau',
    # 'scheduler': None,

    # 该scheduler在指定的epoch步长上按照固定的因子来降低学习率。需要以下参数：  # 'decay_rate': 一个浮点数，控制学习率下降的速率。默认值为0.999。   # 'decay_steps': 一个整数，表示在多少个epoch后降低学习率。默认值为100。
    # 'scheduler_type': 'Step',
    # 'scheduler':{ 'decay_rate': 0.999,'decay_steps': 100,'verbose': True},

    # 'scheduler_type': 'CAWR',
    # 'scheduler':{'T_mult': 1, 'rewarm_epoch_num': 2, 'verbose': True},

    # 'scheduler_type': 'CAL',
    # 'scheduler': {'rewarm_epoch_num': 2,'verbose': True},


    'optimizer_betas': (0.9, 0.999),
    'train_batch_size': 1,
    'eval_batch_size': 2,
    'test_batch_size': 2,
    'learning_rate': 2e-5,  # lr学习率
    'adam_epsilon': 1e-8,  # ε是Adam优化器的一个超参数，它是一个小的正数，通常取值为10的负8次方（1e-8），用于防止除零错误以及数值稳定性。
    'gradient_accumulation_steps': 1,  # 表示梯度积累的步数。
    # 'max_grad_norm'是模型训练过程中的一个超参数，它控制了梯度裁剪的大小。梯度裁剪是一种常用的防止梯度爆炸的方法，可以将梯度的范数限制在一个预设的范围内，以防止梯度过大，导致模型无法收敛。
    # 具体来说，'max_grad_norm'的值代表着允许的最大的梯度范数。在训练过程中，如果模型的梯度范数超过了这个阈值，那么就会将梯度进行裁剪，
    # 将其缩放到阈值以内。通常情况下，'max_grad_norm'的值被设置为1.0，表示不允许梯度的范数超过1.0。通过使用梯度裁剪，可以有效地控制模型的训练过程，防止模型过拟合或者无法收敛。
    'max_grad_norm': 1.0,
    'weight_decay': 0,
    'warmup_steps': 0,
    'output_dir': './output',
    'max_seq_length': 64, # 如果资源充足，推荐长度2048 与官方保持一致
    'max_target_length': 100,  # 预测最大长度, 保留字段
    'use_fast_tokenizer': False,  # 'use_fast_tokenizer'的值为False时，使用的是较慢但功能更全面的tokenizer，其速度较慢但可以进行更多的处理；当值为True时，使用的是速度更快但功能较少的tokenizer。
    'do_lower_case': False,  # 'do_lower_case'的值为False时，表示不进行小写处理，即保留原始文本中的大小写形式；当值为True时，表示将所有文本转化为小写形式进行处理。

    ##############  lora模块
    #注意lora和 ptuning-v2 禁止同时使用
    'with_lora': True,  # 是否启用lora模块
    'inference_mode': False, # 推理模型, 不需要手动设置
    'r': 8,  # 'r': 8是LoRa模块中的超参数，表示每个头被分为多少个局部相关性区域。
    'target_modules': ['query_key_value'],  # 表示在哪些模块上使用LoRa模块。在这里，LoRa模块只应用在查询（query）、键（key）和值（value）的操作上。
    'target_dtype': '16',  # '16'表示数据类型，这里为16位浮点数。
    'lora_alpha': 32,  # 是LoRa模块中的超参数，表示LoRa转换中$\alpha$的值。
    # 'enable_lora': [True],
    'lora_model_name_or_path':None,
    'enable_lora': None,
    'lora_dropout': 0.1,  # 是LoRa模块中的超参数，表示使用dropout正则化的概率。
    'bias': 'none',  # Bias type for Lora. Can be 'none', 'all' or 'lora_only'"
}

#lora 模式暂时不支持deepspeed
enable_deepspeed = False


data_conf = {
   'strategy': DataStrategy.truncation, # 数据策略选项
    DataStrategy.truncation: {
        'ensure_answer_min_length': 1,
    },
    DataStrategy.singlesliding: {
        'sliding_size': train_info_args['max_seq_length'] // 3 * 2, #prompt滑动窗口大小
        'p':1, # p < 0 , 随机选举prompt
    },
    DataStrategy.doublesliding: {
        'sliding_size': train_info_args['max_seq_length'] // 3 * 2, #双滑滑动窗口大小
        'p':1,# p < 0 , 随机选举prompt
    },
}



def get_deepspeed_config():
    # 是否开启deepspeed
    if not enable_deepspeed:
        return None
    with open('./deepspeed.json', mode='r', encoding='utf-8') as f:
        deepspeed_config = json.loads(f.read())
    return deepspeed_config

def preprocess(text):
  #text = text.replace("\n", "\\n").replace("\t", "\\t")
  return text

def postprocess(text):
  # return text.replace("\\n", "\n").replace("\\t", "\t")
  return text



class NN_DataHelper(DataHelper):
    index = 1
    def on_data_ready(self):
        self.index = -1

    # 切分词
    def on_data_process(self, data: typing.Any, mode: str):
        self.index += 1   # 增加self.index的值
        prompt = data[0]  # 从输入数据中提取问题
        answer = data[1]  # 从输入数据中提取答案

        # 获取tokenizer和config
        tokenizer: ChatGLMTokenizer
        config: ChatGLMConfig
        max_seq_length = self.max_seq_length_dict[mode]
        tokenizer = self.tokenizer
        config = self.config

        # 如果self.sptoken不存在，则将其初始化为[50256, 50256]，表示两个结束符
        if not hasattr(self, 'sptoken'):
            self.sptoken = tokenizer.encode(text="")[-2:]

        # 对问题和答案进行编码
        a_ids = tokenizer.encode(text=prompt, add_special_tokens=False)
        b_ids = tokenizer.encode(text=answer, add_special_tokens=False)

        # 获取数据处理策略
        strategy = data_conf['strategy']

        # 根据不同的处理策略选择不同的处理方法
        if strategy == DataStrategy.truncation:
            ds = TokenTruncation.process(tokenizer,config,a_ids, b_ids, max_seq_length, self.sptoken ,**data_conf[strategy])
        elif strategy == DataStrategy.singlesliding:
            ds = TokenSingleSliding.process(tokenizer,config, a_ids, b_ids, max_seq_length, self.sptoken, **data_conf[strategy])
        elif strategy == DataStrategy.doublesliding:
            ds = TokenDoubleSliding.process(tokenizer,config, a_ids, b_ids, max_seq_length, self.sptoken, **data_conf[strategy])
        else:
            raise ValueError('Invlid strategy',strategy)

        if not ds:
            return None

        if self.index < 3:
            print(ds[0])
        return ds

    # {
    #     "id": 0, "paragraph": [
    #     # 一轮会话
    #     {
    #         "q": "从南京到上海的路线",
    #         "a": [
    #             "你好，南京到上海的路线如下：",
    #             "1. 南京到上海，可以乘坐南京地铁1号线，在南京站乘坐轨道交通1号线。",
    #             "2. 南京到浦东机场，可以搭乘上海地铁1号，在陆家嘴站乘坐地铁1线，在浦东国际机场站乘坐机场快线，前往上海浦东国际机场。",
    #             "3. 上海到南京，可以换乘上海地铁2号线，从南京站换乘地铁2线，再从南京南站换乘地铁1路，然后到达上海站"
    #         ]
    #     }
    #     # 二轮....
    # ]
    # }

    # 读取文件 ,作用是读取文件列表中的文件，提取出每个文件中的对话数据， cfiles: typing.List并将其转化为模型需要的格式。 是文件路径的列表
    def on_get_corpus(self, files: typing.List, mode: str):
        D = []  # 返回值是一个列表 D，其中包含了所有对话的 prompt 和对应的 answer。对话的 prompt 是用于生成回答的问题描述，answer 是该问题的正确回答。
        for file in files:
            with open(file, mode='r', encoding='utf-8', newline='\n') as f:
                lines = f.readlines()  # 遍历 files 中的每个文件，并逐行读取文件中的内容。

            for line_id, line in enumerate(lines):
                jd = json.loads(line)  # 每一行都被解析成一个 JSON 对象 jd
                if not jd:
                    continue
                paragraph = jd['paragraph']  # 这个对象包含了一个 paragraph 字段，它是一个列表，表示对话中的每一个轮次。
                if line_id < 10:
                    print(paragraph)
                paragraph = [(preprocess(session['q']),preprocess('\n'.join(session['a']))) for session in paragraph]
                for sid,(q,a) in enumerate(paragraph):
                    assert len(a),ValueError('answer cannot empty')  # 提示答案不能为空。
                    if sid == 0:
                        D.append((q, a))  # 第一个问答直接存入D中。
                    else:
                        prompt_text = ''  # 第二次问答之后的prompt包含前几轮的问答信息
                        for j in range(sid + 1):
                            if j == sid:
                                prompt_text += "[Round {}]\n问：{}\n答：".format(sid, paragraph[j][0])
                            else:
                                prompt_text += "[Round {}]\n问：{}\n答：{}".format(j, paragraph[j][0], paragraph[j][1])
                        D.append((prompt_text,a))
        return D # 最终返回的D是[(q,a),(q,a),(q,a),(q,a)……]

    # 该函数的主要功能是将 batch 中的数据组合成一个 dict，dict 中的 key 为 input_ids、attention_mask、position_ids、labels 等，对应的 value 为 torch.Tensor。
    def collate_fn(self,batch):
        if not hasattr(self,'sptoken'):
            self.sptoken = self.tokenizer.encode(text="")[-2:]

        o = {}
        for i, b in enumerate(batch):
            if i == 0:
                for k in b:
                    o[k] = [torch.tensor(b[k])]
            else:
                for k in b:
                    o[k].append(torch.tensor(b[k]))
        for k in o:
            o[k] = torch.stack(o[k])


        max_len = torch.max(o.pop('seqlen')).tolist()
        b_input_ids = o['input_ids'][:, :max_len]
        ctxlens = o.pop('ctxlen')  # 兼容旧版本数据
        if ctxlens is None:
            ctxlens = [None] * len(b_input_ids)

        b_position_ids,b_attention_mask = [],[]
        for input_ids,context_length in zip(b_input_ids,ctxlens):
            context_length = context_length.squeeze(dim=-1)
            mask_position = context_length - 1
            position_ids = list(range(context_length)) + [mask_position] * (max_len - context_length)
            block_position_ids = [0] * context_length + list(range(1, max_len - context_length + 1))


            attention_mask = torch.ones((1, max_len, max_len))
            attention_mask = torch.tril(attention_mask)
            attention_mask[..., :context_length] = 1
            attention_mask = (attention_mask < 0.5)

            b_position_ids.append(torch.stack((torch.tensor(position_ids),torch.tensor(block_position_ids))))
            b_attention_mask.append(attention_mask)

        b_attention_mask = torch.stack(b_attention_mask, dim=0)
        b_position_ids = torch.stack(b_position_ids,dim=0)

        o['input_ids'] = b_input_ids.long()
        o['attention_mask'] = b_attention_mask.bool()
        o['position_ids'] = b_position_ids.long()
        o['labels'] = o['labels'][:, :max_len].long()
        return o


if __name__ == '__main__':
    parser = HfArgumentParser((ModelArguments, TrainingArguments, DataArguments, LoraArguments))
    model_args, training_args, data_args, lora_args = parser.parse_dict(train_info_args)

    dataHelper = NN_DataHelper(model_args, training_args, data_args)
    tokenizer, config, _,_ = dataHelper.load_tokenizer_and_config(tokenizer_class_name=ChatGLMTokenizer,config_class_name=ChatGLMConfig)
    config.eos_token_id = 130005



    # 缓存数据集
    # 检测是否存在 output/dataset_0-train.record ，不存在则制作数据集
    if data_args.do_train:
        # 生成训练数据集。其中 data_args.train_file 表示训练数据文件的路径，mixed_data=False 表示训练数据是否为多轮对话数据，shuffle=True 表示是否对训练数据进行随机排序，mode='train' 表示生成的是训练数据集。
        dataHelper.make_dataset_with_args(data_args.train_file,mixed_data=False,shuffle=True,mode='train')
    if data_args.do_eval:
        dataHelper.make_dataset_with_args(data_args.eval_file, shuffle=False,mode='eval')
    if data_args.do_test:
        dataHelper.make_dataset_with_args(data_args.test_file, shuffle=False,mode='test')


    # 这个函数的作用是将多个TFRecord格式的数据文件进行打乱顺序，并输出到一个新的TFRecord文件中。
    def shuffle_records(record_filenames, outfile, compression_type='GZIP'):
        print('shuffle_records record...')
        options = RECORD.TFRecordOptions(compression_type=compression_type)
        dataset_reader = Loader.RandomDataset(record_filenames, options=options, with_share_memory=True)
        data_size = len(dataset_reader)
        all_example = []
        for i in tqdm(range(data_size), desc='load records'):
            serialized = dataset_reader[i]
            all_example.append(serialized)
        dataset_reader.close()

        shuffle_idx = list(range(data_size))
        random.shuffle(shuffle_idx)
        writer = WriterObject(outfile, options=options)
        for i in tqdm(shuffle_idx, desc='shuffle record'):
            example = all_example[i]
            writer.write(example)
        writer.close()


    # 对每个record 再次打乱
    for filename in dataHelper.train_files:
        shuffle_records(filename, filename)