#!/usr/bin/env python
# coding=utf-8
"""
Fine-tuning the library models for sequence to sequence.
下面给出了一个命令行运行微调main.py的例子
#1.激活环境
#2.执行torchrun
#--standalone 这个标志指定使用“standalone”模式运行分布式训练。这意味着它不依赖于任何特定的集群管理器（如Slurm、Kubernetes等）。这通常是本地运行或简单的分布式设置的选择。
#--nnodes=1 这指定了参与分布式训练的节点（或机器）的数量。在这个例子中，它被设置为1，意味着只有一个节点（很可能是本地机器）参与训练
#--nproc-per-node=1:这指定了每个节点上要运行的进程数量。这通常与你想要使用的GPU数量相关。在这个例子中，它被设置为1，这意味着在指定的节点上只有一个进程运行。
#如果你在一台有4个GPU的机器上设置--nproc-per-node=4，那么会在该机器上启动4个进程，每个进程使用一个GPU
具体命令如下
source /public1/home/stu52275901023/anaconda3/bin/activate ptune
torchrun --standalone --nnodes=1 --nproc-per-node=1 main.py \
--do_train \
--train_file ../brain_data/brain_data_tra.json \
--validation_file ../brain_data/brain_data_val.json \
--preprocessing_num_workers 10 \
--prompt_column prompt \
--response_column response \
--history_column history \
--overwrite_cache \
--model_name_or_path ../model/chatglm2-6b \
--output_dir ../ptune_point_V1/ \
--overwrite_output_dir \
--max_source_length 512 \
--max_target_length 512 \
--per_device_train_batch_size 6 \
--per_device_eval_batch_size 1 \
--gradient_accumulation_steps 16 \
--predict_with_generate \
--max_steps 500 \
--logging_steps 10 \
--save_steps 100 \
--learning_rate 1e-2 \
--pre_seq_len 512 \
"""
# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
#基础库
import logging
import os
import sys
import json
import numpy as np
from datasets import load_dataset
import jieba
from rouge_chinese import Rouge
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
import torch
#transformer库
import transformers
from transformers import (
AutoConfig,
AutoModel,
AutoTokenizer,
DataCollatorForSeq2Seq,
HfArgumentParser,
Seq2SeqTrainingArguments,
set_seed,
)
from trainer_seq2seq import Seq2SeqTrainer
#注意，注意，从当前目录下的arguments.py导入ModelArguments, DataTrainingArguments数据参数类
from arguments import ModelArguments, DataTrainingArguments
logger = logging.getLogger(__name__)
def main():
#导入相关参数，√ 表示刚才命令行例子中用到了
"""
ModelArguments（自己定义的）:
这个类可能包含与模型配置和加载有关的参数。
√ model_name_or_path: 预训练模型的名称（如"bert-base-uncased"）或模型文件的路径。
config_name: 可选的配置名称或路径，如果没有设置，则默认使用model_name_or_path。
tokenizer_name: 可选的分词器名称或路径。
√ pre_seq_len:这个参数很重要，表示前缀编码器prefixencode层的第一层参数数量，同时设定了这个参数表示微调时只会修改prefixencode层参数
ptuning_checkpoint：可以设定这个值，该值是一个路径，如果有则会接着检查点继续训练
DataTrainingArguments（自己定义的）:
这个类通常定义了与训练数据相关的参数。
√ train_file: 训练数据的路径。
√ validation_file: 验证数据的路径。
√ max_source_length: 输入序列的最大长度,需要注意的是多轮对话下，这个长度应该大于当前轮的qurey+history。
√ max_target_length：目标序列的最大长度，这个长度应该大于任一一轮的response。
√ overwrite_cache: 如果设置为True，则在处理数据时覆盖现有的缓存。
√ preprocessing_num_workers ：用于预处理的进程数
√ prompt_column：表示训练的json文件中，prompt的关键字
√ response_column: 表示训练的json文件中，response的关键字
√ history_column: 表示训练的json文件中，history的关键字
Seq2SeqTrainingArguments（API在transformer下的trainer里：https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.gradient_accumulation_steps）:
这个类可能专为序列到序列模型定义了参数，如T5、Bart等。
√ output_dir：模型输出路径
√ overwrite_output_dir ：是否重写模型输出路径内容
√ learning_rate: 学习率。
num_train_epochs: 训练时代数。
√ max_steps : 最大训练步数，会覆盖num_train_epochs
√ predict_with_generate:是否使用generate计算指标
√ per_device_train_batch_size: 每个设备上的训练时批次大小。
√ per_device_eval_batch_size: 每个设备上的评估时批次大小。
√ gradient_accumulation_steps (int, optional, defaults to 1):更新梯度时累计的梯度次数
weight_decay: 权重衰减。
√ logging_steps: 多少步进行一次日志记录。
√ save_steps: 多少步保存一次模型。
√ do_train :bool 是否训练
√ do_eval :bool 是否评估
do_predict :bool 是否预测
"""
#Hf是Huggingface的缩写
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
#如果是json参数文件，则有两个参数，第二个参数是json的路径
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
# If we pass only one argument to the script and it's the path to a json file,
# let's parse it to get our arguments.
model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
else:
#否则直接命令行读取参数
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
# Setup logging
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
handlers=[logging.StreamHandler(sys.stdout)],
)
#看下训练参数里是否说要输出log
if training_args.should_log:
# The default of training_args.log_level is passive, so we set log level at info here to have that default.
transformers.utils.logging.set_verbosity_info()
log_level = training_args.get_process_log_level()
logger.setLevel(log_level)
# datasets.utils.logging.set_verbosity(log_level)
transformers.utils.logging.set_verbosity(log_level)
transformers.utils.logging.enable_default_handler()
transformers.utils.logging.enable_explicit_format()
# Log on each process the small summary:
logger.warning(
f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+ f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
)
logger.info(f"Training/evaluation parameters {training_args}")
# Set seed before initializing model.
set_seed(training_args.seed)
# Load dataset
data_files = {}
if data_args.train_file is not None:
data_files["train"] = data_args.train_file
extension = data_args.train_file.split(".")[-1]
if data_args.validation_file is not None:
data_files["validation"] = data_args.validation_file
extension = data_args.validation_file.split(".")[-1]
if data_args.test_file is not None:
data_files["test"] = data_args.test_file
extension = data_args.test_file.split(".")[-1]
#做成原始数据
raw_datasets = load_dataset(
extension,
data_files=data_files,
cache_dir=model_args.cache_dir,
use_auth_token=True if model_args.use_auth_token else None,
)
# Load pretrained model and tokenizer
config = AutoConfig.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)
#重新设定pre_seq_len的值
config.pre_seq_len = model_args.pre_seq_len
#重新设定prefix_projection的值，这个是个bool类型，true表示有前缀
config.prefix_projection = model_args.prefix_projection
tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)
#是否有额外前缀编码层状态
if model_args.ptuning_checkpoint is not None:
# Evaluation
# Loading extra state dict of prefix encoder
#在模型文件目录下有个config.json,里面完成了"AutoModel"到"modeling_chatglm.ChatGLMForConditionalGeneration"的映射
#所以下面的代码会导入modeling_chatglm.py文件里的ChatGLMForConditionalGeneration
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin"))
new_prefix_state_dict = {}
for k, v in prefix_state_dict.items():
if k.startswith("transformer.prefix_encoder."):
new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
else:
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
#是否量化
if model_args.quantization_bit is not None:
print(f"Quantized to {model_args.quantization_bit} bit")
model = model.quantize(model_args.quantization_bit)
if model_args.pre_seq_len is not None:
# P-tuning v2
model = model.half()
model.transformer.prefix_encoder.float()
else:
# Finetune
model = model.float()
#有前缀则取出前缀
prefix = data_args.source_prefix if data_args.source_prefix is not None else ""
# Preprocessing the datasets.
# We need to tokenize inputs and targets.
if training_args.do_train:
column_names = raw_datasets["train"].column_names
elif training_args.do_eval:
column_names = raw_datasets["validation"].column_names
elif training_args.do_predict:
column_names = raw_datasets["test"].column_names
else:
logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
return
# Get the column names for input/target.
prompt_column = data_args.prompt_column
response_column = data_args.response_column
history_column = data_args.history_column
# Temporarily set max_target_length for training.
max_target_length = data_args.max_target_length
def preprocess_function_eval（examples):
inputs, targets = [], []
for i in range(len(examples[prompt_column])):
if examples[prompt_column][i] and examples[response_column][i]:
query = examples[prompt_column][i]
history = examples[history_column][i] if history_column is not None else None
prompt = tokenizer.build_prompt(query, history)
inputs.append(prompt)
targets.append(examples[response_column][i])
inputs = [prefix + inp for inp in inputs]
model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, truncation=True, padding=True)
labels = tokenizer(text_target=targets, max_length=max_target_length, truncation=True)
if data_args.ignore_pad_token_for_loss:
labels["input_ids"] = [
[(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
]
model_inputs["labels"] = labels["input_ids"]
return model_inputs
def preprocess_function_train(examples):
#最大训练长度
max_seq_length = data_args.max_source_length + data_args.max_target_length + 1
#初始化训练时模型输入
model_inputs = {
"input_ids": [],#history+prompt
"labels": [],#response
}
'''
#遍历每一条数据,数据示例如下
{
"prompt": "翻译论语吾日三省吾身",
"response": "好的。\n 【译文】曾子说:\"我每天多次地反省自己:替别人办事是不是尽心竭力呢？跟朋友交往是不是诚实呢？老师传授的知识是不是复习过了呢？ \"",
"history": []
},
{
"prompt": "我想学习欧阳修的醉翁亭记",
"response": "好的。【原文】醉翁亭记 宋 欧阳修环滁皆山也。其西南诸峰，林壑(hè)尤美。望之蔚然而深秀者，琅琊也。山行六七里，渐 闻水声潺潺而泻出于两峰之间者，酿泉也。峰回路转，有亭翼然临于泉上者，醉翁亭也。作亭 者谁？山之僧智仙也。名之者谁？太守自谓也。太守与客来饮于此，饮少辄(zhé)醉，而年又 最高，故自号曰醉翁也。醉翁之意不在酒，在乎山水之间也。山水之乐，得之心而寓之酒也。若夫日出而林霏开，云归而岩穴暝，晦明变化者，山间之朝暮也。野芳发而幽香，佳木秀 而繁阴，风霜高洁，水落而石出者，山间之四时也。朝而往，暮而归，四时之景不同，而乐亦 无穷也。至于负者歌于途，行者休于树，前者呼，后者应，伛（yǔ）偻（lǚ）提携，往来而不绝者， 滁人游也。临溪而渔，溪深而鱼肥。酿泉为酒，泉香而酒洌；山肴野蔌（sù），杂然而前陈者，太守宴也。宴酣之乐，非丝非竹，射者中，弈者胜，觥（gōng）筹交错，起坐而喧哗者，众宾欢也。苍颜白发，颓然乎其间者，太守醉也。已而夕阳在山，人影散乱，太守归而宾客从也。树林阴翳（yì），鸣声上下，游人去而禽鸟乐也。然而禽鸟知山林之乐，而不知人之乐；人知从太守游而乐，而不知太守之乐其乐也。醉能同其乐，醒能述以文者，太守也。太守谓谁？庐陵欧阳修也。【译文】环绕着滁州城的都是山。它西南面的各个山峰，树林、山谷尤其优美，远远望去那树木茂盛，又幽深又秀丽的，是琅琊山。沿着山路走六七里，就渐渐听到潺潺的水声，从两座山峰中间飞泻出来的，是酿泉。山势回环，路也跟着拐弯，有亭子四角翘起，像鸟张开翅膀一样，坐落在泉水边上的，是醉翁亭。修建亭子的人是谁？是山中的和尚智仙。给它命名的人是谁？是太守用自己的号（醉翁）来命名的。太守和客人到这里来喝酒，喝一点就醉了，而且年龄又最大，所以给自己取个别号叫醉翁。醉翁的情趣不在于喝酒，而在于秀丽的山水之间。欣赏山水的乐趣，领会在心里，寄托在喝酒上。像那太阳出来，树林中的雾气消散，暮云回聚拢来，山里就昏暗了，或明或暗，变化不一，这是山间的早晨和傍晚。野花开放，散发清幽的香气，好的树木枝叶繁茂，形成浓郁的绿阴，秋高气爽,霜色洁白，水位低落，石头显露，这是山里的四季的景色。早晨上山，傍晚返回，四季的景色不同，因而乐趣也没有穷尽。至于背着东西的人路上唱歌，走路的人在树下休息，前面的人呼唤，后面的人答应，老人小孩络绎不绝，这是滁州人出游。到溪水捕鱼，溪水深，鱼儿肥，用酿泉的水酿酒，泉水香而酒色清澈，山中野味，野菜，交错地在前面摆着，这是太守举行的酒宴。酒宴上饮酒的乐趣，不在于音乐。投壶的人中了目标，下棋的人得胜了，酒杯和酒筹交互错杂，人们有时站立，有时坐着，大声喧哗，宾客们（尽情）欢乐。脸色苍老，头发花白，醉醺醺地在宾客们中间的是喝醉了的太守。不久夕阳落山，人影纵横散乱，太守回归而宾客跟从。这时树林里浓荫遮蔽，鸟儿到处鸣叫，游人离开后禽鸟就快乐了。然而禽鸟只知道山林的乐趣，却不知道人的乐趣，人们只知道跟随太守游玩的乐趣，却不知道太守因为他们的快乐而快乐。醉了能够和大家一起享受快乐，醒来能够用文章记述这乐事的人，是太守。太守是谁？是庐陵的欧阳修。",
"history": [
[
"我想学习一下杜牧的泊秦淮",
"好的。【原文】泊秦淮 唐 杜牧烟笼寒水月笼沙，夜泊秦淮近酒家。商女不知亡国恨，隔江犹唱后庭花。【译文】迷离月色和轻烟笼罩着寒水白沙，夜晚船只停泊在秦淮河边靠近岸上的酒家。卖唱的歌女不知道什么是亡国之恨，隔着江水仍在高唱着《玉树后庭花》。【注释】选自冯集梧《樊川文集》卷四。（上海古籍出版社1978年版）。秦淮，即秦淮河，发源于江苏句容大茅山与溧（lì）水东庐山两山间，经南京流入长江。相传为秦始皇南巡会稽时开凿的，用来疏通淮水，故称秦淮河。泊：停泊。商女：以卖唱为生的歌女。犹：副词。还，仍然。后庭花：歌曲《玉树后庭花》的简称。南朝陈皇帝陈叔宝（即陈后主）溺于声色，作此曲与后宫美女寻欢作乐，终致亡国，所以后世称此曲为“亡国之音”。"
],
[
"翻译一下诸葛亮的出师表",
"好的。【译文】臣诸葛亮上表进言：先帝开创大业未完成一半，竟中途去世。如今天下分成三国，我益州地区人力疲惫、 民生凋敝，这真是处在万分危急、存亡难料的时刻。但是，宫廷里侍奉守卫的臣子，不敢稍有懈怠；疆场上 忠诚有志的将士，舍身忘死的作战，这都是追念先帝的特殊恩遇，想报答给陛下的缘故。陛下确实应该广开 言路听取群臣意见，发扬光大先帝遗留下来的美德，振奋鼓舞志士们的勇气，绝不应随便看轻自己，说出无 道理的话，从而堵塞了忠诚进谏的道路。宫里身边的近臣和丞相府统领的官吏，本都是一个整体，升赏惩罚， 扬善除恶，不应标准不同。如有作坏事违犯法纪的，或尽忠心做善事的，应该一律交给主管部门加以惩办或 奖赏，以显示陛下在治理方面公允明察，切不应私心偏袒，使宫廷内外施法不同。 侍中、侍郎郭攸之、费、董允等，这都是些品德良善诚实、情志意念忠贞纯正的人，因而先帝才选留下 来辅佐陛下。我认为宫内的事情，事无论大小，都要征询他们的意见，然后再去施行。这样一定能够补正疏 失，增益实效。将军向宠，性情德行平和公正，了解通晓军事，当年试用，先帝曾加以称赞，说他能干，因 而经众人评议荐举任命为中部督。我认为军营里的事情，事情无论大小，都要征询他的意见，就一定能够使 军伍团结和睦，德才高低的人各有合适的安排。亲近贤臣，远避小人，这是汉朝前期所以能够兴盛的原因； 亲近小人，远避贤臣，这是汉朝后期所以衰败的原因。先帝在世的时候，每次跟我评论起这些事，对于桓帝、灵帝时代，没有不哀叹和憾恨的。侍中郭攸之、费，尚书陈震，长史张裔，参军蒋琬，这些都是忠贞、坦直，能以死报国的节义臣子，诚愿陛下亲近他们，信任他们，则汉王室的兴盛，就时间不远了。我本是个平民，在南阳郡务农亲耕，在乱世间只求保全性命，不希求诸侯知道我而获得显贵。先帝不介意我的卑贱，委屈地自我降低身份，接连三次到草庐来访看我，征询我对时局大事的意见，因此我深为感激，从而答应为先帝驱遣效力。后来正遇危亡关头，在战事失败的时候我接受了任命，在危机患难期间我受到委任，至今已有二十一年了。先帝深知我做事谨慎，所以临去世时把国家大事嘱托给我了。接受遗命以来，日夜担忧兴叹，只恐怕托付给我的大任不能完成，从而损害先帝的英明。所以我五月率兵南渡泸水，深入荒芜之境。如今南方已经平定，武库兵器充足，应当鼓励和统率全军，北伐平定中原地区，我希望竭尽自己低下的才能，消灭奸邪势力，复兴汉朝王室，迁归旧日国都。这是我用来报答先帝，并尽忠心于陛下的职责本分。至于掂量利弊得失，毫无保留地进献忠言，那就是郭攸之、费、董允的责任了。希望陛下责成我去讨伐奸贼并取得成效，如果不取得成效，那就惩治我失职的罪过，用来上告先帝的神灵。如果没有发扬圣德的言论，那就责备郭攸之、费、董允等人的怠慢，公布他们的罪责。陛下也应该自己思虑谋划，征询从善的道理，明察和接受正直的进言，远念先帝遗诏中的旨意，我就受恩、感激不尽了。如今正当离朝远征，流着泪写了这篇表文，激动得不知该说些什么话。"
],
[
"介绍一下王维的使至塞上",
"好的。【原文】使至塞上 唐 王维单车欲问边，属国过居延。 征蓬出汉塞，归雁入胡天。 大漠孤烟直，长河落日圆。 萧关逢候骑，都护在燕然。【译文】轻车简从将要去慰问边关，我要到远在西北边塞的居延。像随风而去的蓬草一样出临边塞，北归大雁正翱翔云天。浩瀚沙漠中孤烟直上云霄，黄河边上落日浑圆。到萧关时遇到侦察骑兵，得知主帅尚在前线未归。【注释】使至塞上：奉命出使边塞。使：出使。单车：一辆车，车辆少，这里形容轻车简从。问边：到边塞去察看，指慰问守卫边疆的官兵。属国：有几种解释：一指少数民族附属于汉族朝廷而存其国号者。汉、唐两朝均有一些属国。二指官名，秦汉时有一种官职名为典属国，苏武归汉后即授典属国官职。唐人有时以“属国”代称出使边陲的使臣。居延：地名，汉代称居延泽，唐代称居延海，在今内蒙古额济纳旗北境。又西汉张掖郡有居延县（参《汉书·地理志》），故城分布在今内蒙古自治区额济纳旗和甘肃省金塔县的境内。又东汉凉州刺史部有张掖居延属国，辖境在居延泽一带。此句一般注本均言王维路过居延。然而王维此次出使，实际上无需经过居延。因而林庚、冯沅君主编的《中国历代诗歌选》认为此句是写唐王朝“边塞的辽阔，附属国直到居延以外”。征蓬：随风飘飞的蓬草，此处为诗人自喻。归雁：雁是候鸟，春天北飞，秋天南行，这里是指大雁北飞。胡天：胡人的领空。这里是指唐军占领的北方地方。大漠：大沙漠，此处大约是指凉州之北的沙漠。孤烟：赵殿成注有二解：一云古代边防报警时燃狼粪，“其烟直而聚，虽风吹之不散”。二云塞外多旋风，“袅烟沙而直上”。据后人有到甘肃、新疆实地考察者证实，确有旋风如“孤烟直上”。又：孤烟也可能是唐代边防使用的平安火。长河：黄河。萧关：古关名，又名陇山关，故址在今宁夏固原东南。候骑：负责侦察、通讯的骑兵。王维出使河西并不经过萧关，此处大概是用何逊诗“候骑出萧关，追兵赴马邑”之意，非实写。都护：唐朝在西北边疆置安西、安北等六大都护府，其长官称都护，每府派大都护一人，副都护二人，负责辖区一切事务。这里指前线统帅。燕然：燕然山，即今蒙古国杭爱山。东汉窦宪北破匈奴，曾于此刻石记功。这里代指前线。“单车”两句：一作“衔命辞天阙，单车欲问边”。候骑：骑马的侦察兵。"
]
]
}
'''
for i in range(len(examples[prompt_column])):
#如果有prompt和response
if examples[prompt_column][i] and examples[response_column][i]:
query, answer = examples[prompt_column][i], examples[response_column][i]
#如果有history
history = examples[history_column][i] if history_column is not None else None
'''
将当前对话prompt/query和历史对话history制作成新的prompt,源码如下
prompt = ""
for i, (old_query, response) in enumerate(history):
prompt += "[Round {}]\n\n问：{}\n\n答：{}\n\n".format(i + 1, old_query, response)
prompt += "[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query)
'''
prompt = tokenizer.build_prompt(query, history)
#有前缀则加在前面
prompt = prefix + prompt
#将prompt变成id列表
a_ids = tokenizer.encode(text=prompt, add_special_tokens=True, truncation=True,
max_length=data_args.max_source_length)
#将response变成id列表
b_ids = tokenizer.encode(text=answer, add_special_tokens=False, truncation=True,
max_length=data_args.max_target_length)
#记录prompt长度
context_length = len(a_ids)
#prompt和response的ids组成新ids，最后加上结束符用以训练。
input_ids = a_ids + b_ids + [tokenizer.eos_token_id]
#tokenizer.pad_token_id 是Hugging Face transformers库中tokenizer的一个属性，它返回用于padding的token的ID。
#实际训练时这个值为100。
labels = [tokenizer.pad_token_id] * context_length + b_ids + [tokenizer.eos_token_id]
#input_ids和labels扩充到最大训练长度
pad_len = max_seq_length - len(input_ids)
input_ids = input_ids + [tokenizer.pad_token_id] * pad_len
labels = labels + [tokenizer.pad_token_id] * pad_len
if data_args.ignore_pad_token_for_loss:
labels = [(l if l != tokenizer.pad_token_id else -100) for l in labels]
model_inputs["input_ids"].append(input_ids)
model_inputs["labels"].append(labels)
return model_inputs
def print_dataset_example(example):
print("input_ids", example["input_ids"])
print("inputs", tokenizer.decode(example["input_ids"]))
print("label_ids", example["labels"])
print("labels", tokenizer.decode(example["labels"]))
#如果训练
if training_args.do_train:
if "train" not in raw_datasets:
raise ValueError("--do_train requires a train dataset")
train_dataset = raw_datasets["train"]
#是否设定了训练数据条数，如果有的话就选随机选择max_train_samples条数据
if data_args.max_train_samples is not None:
max_train_samples = min(len(train_dataset), data_args.max_train_samples)
train_dataset = train_dataset.select(range(max_train_samples))
#开会话调用刚才preprocess_function_train映射成数据集
with training_args.main_process_first(desc="train dataset map pre-processing"):
train_dataset = train_dataset.map(
preprocess_function_train,
batched=True,
num_proc=data_args.preprocessing_num_workers,
remove_columns=column_names,
load_from_cache_file=not data_args.overwrite_cache,
desc="Running tokenizer on train dataset",
)
#输出第一个训练实例
print_dataset_example(train_dataset[0])
if training_args.do_eval:
max_target_length = data_args.val_max_target_length
if "validation" not in raw_datasets:
raise ValueError("--do_eval requires a validation dataset")
eval_dataset = raw_datasets["validation"]
if data_args.max_eval_samples is not None:
max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
eval_dataset = eval_dataset.select(range(max_eval_samples))
with training_args.main_process_first(desc="validation dataset map pre-processing"):
eval_dataset = eval_dataset.map(
preprocess_function_eval,
batched=True,
num_proc=data_args.preprocessing_num_workers,
remove_columns=column_names,
load_from_cache_file=not data_args.overwrite_cache,
desc="Running tokenizer on validation dataset",
)
print_dataset_example(eval_dataset[0])
if training_args.do_predict:
max_target_length = data_args.val_max_target_length
if "test" not in raw_datasets:
raise ValueError("--do_predict requires a test dataset")
predict_dataset = raw_datasets["test"]
if data_args.max_predict_samples is not None:
max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
predict_dataset = predict_dataset.select(range(max_predict_samples))
with training_args.main_process_first(desc="prediction dataset map pre-processing"):
predict_dataset = predict_dataset.map(
preprocess_function_eval,
batched=True,
num_proc=data_args.preprocessing_num_workers,
remove_columns=column_names,
load_from_cache_file=not data_args.overwrite_cache,
desc="Running tokenizer on prediction dataset",
)
print_dataset_example(predict_dataset[0])
# Data collator
label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
#DataCollatorForSeq2Seq对象是用于处理和准备序列到序列任务的数据的数据整理器
data_collator = DataCollatorForSeq2Seq(
tokenizer,
model=model,
label_pad_token_id=label_pad_token_id,
pad_to_multiple_of=None,
padding=False
)
# Metric
def compute_metrics(eval_preds):
preds, labels = eval_preds
if isinstance(preds, tuple):
preds = preds[0]
decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
if data_args.ignore_pad_token_for_loss:
# Replace -100 in the labels as we can't decode them.
labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
score_dict = {
"rouge-1": [],
"rouge-2": [],
"rouge-l": [],
"bleu-4": []
}
for pred, label in zip(decoded_preds, decoded_labels):
hypothesis = list(jieba.cut(pred))
reference = list(jieba.cut(label))
rouge = Rouge()
scores = rouge.get_scores(' '.join(hypothesis) , ' '.join(reference))
result = scores[0]
for k, v in result.items():
score_dict[k].append(round(v["f"] * 100, 4))
bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
score_dict["bleu-4"].append(round(bleu_score * 100, 4))
for k, v in score_dict.items():
score_dict[k] = float(np.mean(v))
return score_dict
# Override the decoding parameters of Seq2SeqTrainer
#设置最大生成长度
training_args.generation_max_length = (
training_args.generation_max_length
if training_args.generation_max_length is not None
else data_args.val_max_target_length
)
#设置beam search数量
training_args.generation_num_beams = (
data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
)
# Initialize our Trainer
trainer = Seq2SeqTrainer(
model=model,
args=training_args,
train_dataset=train_dataset if training_args.do_train else None,
eval_dataset=eval_dataset if training_args.do_eval else None,
tokenizer=tokenizer,
data_collator=data_collator,
compute_metrics=compute_metrics if training_args.predict_with_generate else None,
save_changed=model_args.pre_seq_len is not None
)
# Training
if training_args.do_train:
checkpoint = None
#是否从历史检查点恢复
if training_args.resume_from_checkpoint is not None:
checkpoint = training_args.resume_from_checkpoint
# elif last_checkpoint is not None:
# checkpoint = last_checkpoint
#Gradient Checkpointing 是一种内存优化技巧，允许你用更少的GPU内存进行训练，尤其是对于大型模型和长序列
#它的原理是不保存每一层的所有中间激活，而是只保存一部分，并在反向传播时重新计算其他激活。这样可以节省内存，但可能会增加计算时间。
model.gradient_checkpointing_enable()
#tensor有一个属性叫requires_grad，它决定是否在反向传播时计算这个tensor的梯度。通常，模型的权重会设置requires_grad=True，因为我们需要它们的梯度来更新权重
model.enable_input_require_grads()
#训练
train_result = trainer.train(resume_from_checkpoint=checkpoint)
# trainer.save_model() # Saves the tokenizer too for easy upload
metrics = train_result.metrics
max_train_samples = (
data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset)
)
metrics["train_samples"] = min(max_train_samples, len(train_dataset))
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)
trainer.save_state()
# Evaluation
results = {}
max_seq_length = data_args.max_source_length + data_args.max_target_length + 1
if training_args.do_eval:
logger.info("*** Evaluate ***")
metrics = trainer.evaluate(metric_key_prefix="eval", do_sample=True, top_p=0.7, max_length=max_seq_length, temperature=0.95)
max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset)
metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset))
trainer.log_metrics("eval", metrics)
trainer.save_metrics("eval", metrics)
if training_args.do_predict:
logger.info("*** Predict ***")
predict_results = trainer.predict(predict_dataset, metric_key_prefix="predict", max_length=max_seq_length, do_sample=True, top_p=0.7, temperature=0.95)
metrics = predict_results.metrics
max_predict_samples = (
data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset)
)
metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset))
trainer.log_metrics("predict", metrics)
trainer.save_metrics("predict", metrics)
if trainer.is_world_process_zero():
if training_args.predict_with_generate:
predictions = tokenizer.batch_decode(
predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True
)
predictions = [pred.strip() for pred in predictions]
labels = tokenizer.batch_decode(
predict_results.label_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)
labels = [label.strip() for label in labels]
output_prediction_file = os.path.join(training_args.output_dir, "generated_predictions.txt")
with open(output_prediction_file, "w", encoding="utf-8") as writer:
for p, l in zip(predictions, labels):
res = json.dumps({"labels": l, "predict": p}, ensure_ascii=False)
writer.write(f"{res}\n")
return results
def _mp_fn(index):
# For xla_spawn (TPUs)
main()
if __name__ == "__main__":
main()