# @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)
