<reponame>waking95/easy-bert
import json
import math
import os
import random
import re

import numpy as np
import torch
import torch.cuda
from sklearn.utils import shuffle
from torch.cuda.amp import autocast, GradScaler
from torch.nn import DataParallel
from transformers import AdamW
from transformers import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup, \
    get_linear_schedule_with_warmup

from easy_bert import logger
from easy_bert.adversarial import FGM, PGD
from easy_bert.base.base_trainer import BaseTrainer
from easy_bert.bert4classification.classification_model import ClassificationModel
from easy_bert.vocab import Vocab


class ClassificationTrainer(BaseTrainer):
    def __init__(self, pretrained_model_dir, model_dir, learning_rate=5e-5, ckpt_name='bert_model.bin',
                 vocab_name='vocab.json', enable_parallel=False, adversarial=None, dropout_rate=0.5,
                 loss_type='cross_entropy_loss', focal_loss_gamma=2, focal_loss_alpha=None, random_seed=0,
                 warmup_type=None, warmup_step_num=10, enable_fp16=False,
                 add_on=None, rnn_hidden=256, rnn_lr=1e-3,
                 load_last_ckpt=False, layer_wise_lr_decay=False, lr_decay_rate=0.95):
        # 设置目录
        self.pretrained_model_dir, self.model_dir = pretrained_model_dir, model_dir
        self.ckpt_name, self.vocab_name = ckpt_name, vocab_name

        self.enable_parallel = enable_parallel

        self.load_last_ckpt = load_last_ckpt  # 是否加载上次模型

        # 设置随机种子
        self.random_seed = random_seed
        self._set_random_seed(random_seed)

        # 对抗训练配置
        self.adversarial = adversarial
        assert adversarial in (None, 'fgm', 'pgd')

        # loss配置
        self.loss_type = loss_type
        self.focal_loss_gamma = focal_loss_gamma
        self.focal_loss_alpha = focal_loss_alpha

        # 自动获取当前设备
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.dropout_rate = dropout_rate
        self.learning_rate = learning_rate
        self.batch_size = None
        self.epoch = None

        # 混合精度配置
        if enable_fp16:
            self.grad_scaler = GradScaler()  # 设置梯度缩放
        self.enable_fp16 = enable_fp16

        # warmup配置
        assert warmup_type in ('constant', 'cosine', 'linear', None)
        assert type(warmup_step_num) in (int, float)
        self.warmup_type = warmup_type
        self.warmup_step_num = warmup_step_num

        # 附加层配置（add-on）
        self.add_on = add_on
        self.rnn_hidden, self.rnn_lr = rnn_hidden, rnn_lr

        # llrd配置
        self.layer_wise_lr_decay = layer_wise_lr_decay
        self.lr_decay_rate = lr_decay_rate

        self.vocab = Vocab()

    def _set_random_seed(self, seed):
        """针对torch torch.cuda numpy random分别设定随机种子"""
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        np.random.seed(seed)
        random.seed(seed)

    def _build_model(self):
        """构建bert模型"""
        # 实例化ClassificationModel模型
        self.model = ClassificationModel(
            self.pretrained_model_dir, self.vocab.label_size, dropout_rate=self.dropout_rate,
            loss_type=self.loss_type, focal_loss_alpha=self.focal_loss_alpha, focal_loss_gamma=self.focal_loss_gamma,
            add_on=self.add_on, rnn_hidden=self.rnn_hidden
        )

        # 如果启用增量训练，预加载之前训练好的模型
        if self.load_last_ckpt and os.path.exists('{}/{}'.format(self.model_dir, self.ckpt_name)):
            self.model.load_state_dict(
                torch.load('{}/{}'.format(self.model_dir, self.ckpt_name), map_location=self.device)
            )
            logger.info('load last ckpt {}/{} success'.format(self.model_dir, self.ckpt_name))

        # 设置AdamW优化器
        no_decay = ["bias", "LayerNorm.weight"]  # bias和LayerNorm不使用正则化
        # 参数分组
        embedding_parameters, encoder_parameters = [], {}
        linear_parameters, rnn_parameters = [], []
        for name, param in self.model.named_parameters():
            if 'rnn' in name:
                rnn_parameters.append((name, param))
            elif 'embeddings' in name:
                embedding_parameters.append((name, param))
            elif 'encoder.layer' in name:
                layer_no = int(re.search('encoder\.layer\.(\d+)', name).group(1))
                if layer_no not in encoder_parameters:
                    encoder_parameters[layer_no] = []
                encoder_parameters[layer_no].append((name, param))
            else:
                linear_parameters.append((name, param))
        # 为不同参数组设置不同的lr
        num_hidden_layers = self.model.bert_model.config.num_hidden_layers
        optimizer_grouped_parameters = [
            {'params': [p for n, p in embedding_parameters if not any(nd in n for nd in no_decay)],
             'weight_decay': 0.01,
             'lr': self.learning_rate * self.lr_decay_rate ** num_hidden_layers if self.layer_wise_lr_decay else self.learning_rate},
            {'params': [p for n, p in embedding_parameters if any(nd in n for nd in no_decay)],
             'weight_decay': 0.0,
             'lr': self.learning_rate * self.lr_decay_rate ** num_hidden_layers if self.layer_wise_lr_decay else self.learning_rate},
        ]
        for layer_no in range(num_hidden_layers):
            optimizer_grouped_parameters.extend([
                {'params': [p for n, p in encoder_parameters[layer_no] if not any(nd in n for nd in no_decay)],
                 'weight_decay': 0.01,
                 'lr': self.learning_rate * self.lr_decay_rate ** (
                         num_hidden_layers - layer_no - 1) if self.layer_wise_lr_decay else self.learning_rate},
                {'params': [p for n, p in encoder_parameters[layer_no] if any(nd in n for nd in no_decay)],
                 'weight_decay': 0.0,
                 'lr': self.learning_rate * self.lr_decay_rate ** (
                         num_hidden_layers - layer_no - 1) if self.layer_wise_lr_decay else self.learning_rate},
            ])
        optimizer_grouped_parameters.extend([
            {'params': [p for n, p in linear_parameters], 'lr': self.learning_rate},
            {'params': [p for n, p in rnn_parameters], 'lr': self.rnn_lr}  # rnn layer设置自己的lr
        ])
        self.optimizer = AdamW(optimizer_grouped_parameters, lr=self.learning_rate)

        # 使用bert的vocab更新Vocab对象
        self.vocab.set_vocab2id(self.model.get_bert_tokenizer().vocab)
        self.vocab.set_id2vocab({_id: char for char, _id in self.vocab.vocab2id.items()})
        self.vocab.set_unk_vocab_id(self.vocab.vocab2id['[UNK]'])
        self.vocab.set_pad_vocab_id(self.vocab.vocab2id['[PAD]'])

        # 启用并行，使用DataParallel封装model
        if self.enable_parallel:
            self.model = DataParallel(self.model, device_ids=list(range(torch.cuda.device_count())))

        # 将模型拷贝到当前设备
        self.model.to(self.device)

    def _save_config(self):
        """保存训练参数配置为json文件"""
        config = {
            'vocab_size': self.vocab.vocab_size,
            'label_size': self.vocab.label_size,
            'learning_rate': self.learning_rate,
            'batch_size': self.batch_size,
            'dropout_rate': self.dropout_rate,
            'epoch': self.epoch,
            'ckpt_name': self.ckpt_name,
            'vocab_name': self.vocab_name,
            'enable_parallel': self.enable_parallel,
            'enable_fp16': self.enable_fp16,
            'adversarial': self.adversarial,
            'loss_type': self.loss_type,
            'focal_loss_gamma': self.focal_loss_gamma,
            'focal_loss_alpha': self.focal_loss_alpha,
            'pretrained_model': os.path.basename(self.pretrained_model_dir),
            'add_on': self.add_on,
            'rnn_hidden': self.rnn_hidden,
            'rnn_lr': self.rnn_lr,
            'random_seed': self.random_seed,
            'layer_wise_lr_decay': self.layer_wise_lr_decay,
            'lr_decay_rate': self.lr_decay_rate,
        }
        with open('{}/train_config.json'.format(self.model_dir), 'w') as f:
            f.write(json.dumps(config, indent=4))

    def _transform_batch(self, batch_texts, batch_labels, max_length=512):
        """将batch的文本及labels转换为bert的输入tensor形式"""
        batch_input_ids, batch_att_mask, batch_label_ids = [], [], []
        for text, label in zip(batch_texts, batch_labels):  # 迭代每条样本
            assert isinstance(text, str)
            # 根据是否并行，获取bert_tokenizer
            bert_tokenizer = self.model.bert_tokenizer if not self.enable_parallel else self.model.module.bert_tokenizer
            encoded_dict = bert_tokenizer.encode_plus(text, max_length=max_length, padding='max_length',
                                                      return_tensors='pt', truncation=True)
            batch_input_ids.append(encoded_dict['input_ids'])
            batch_att_mask.append(encoded_dict['attention_mask'])
            batch_label_ids.append(self.vocab.tag2id[label])
        batch_input_ids, batch_att_mask, batch_label_ids = \
            torch.cat(batch_input_ids), torch.cat(batch_att_mask), torch.LongTensor(batch_label_ids)

        # 将数据拷贝到当前设备
        batch_input_ids, batch_att_mask, batch_label_ids = \
            batch_input_ids.to(self.device), batch_att_mask.to(self.device), batch_label_ids.to(self.device)

        return batch_input_ids, batch_att_mask, batch_label_ids

    def train(self, train_texts, train_labels, validate_texts, validate_labels,
              batch_size=30, epoch=10, max_loss_num=10, warning_max_len=True):
        """训练
        Args:
            train_texts: list[str] 训练集样本
            train_labels: list[str] 训练集标签
            validate_texts: list[str] 验证集样本
            validate_labels: list[str] 验证集标签
            batch_size: int
            epoch: int
            max_loss_num: int 连续max_loss_num个validate loss平均，根据其保存模型
            warning_max_len: bool 是否警告max_len超过512（普通bert不能处理超过512，除了一些变体如longformer）
        """
        self.batch_size = batch_size
        self.epoch = epoch

        # 构建词库表（只构建label id映射，词库bert已有）、构建模型
        self.vocab.build_vocab(labels=train_labels, build_texts=False, with_build_in_tag_id=False)
        self._build_model()
        # 保存词库、保存训练config
        self.vocab.save_vocab('{}/{}'.format(self.model_dir, self.vocab_name))
        self._save_config()

        logger.info('train samples: {}, validate samples: {}'.format(len(train_texts), len(validate_texts)))

        best_loss = float("inf")  # 当前最小的loss
        loss_buff = []  # 保存最近max_loss_num的个valid loss

        step = 0

        # 根据对抗配置，实例化不同的对抗实例
        if self.adversarial:
            logger.info('enable adversarial training')
            adv_class = {'fgm': FGM, 'pgd': PGD}[self.adversarial]
            adv = adv_class(self.model, grad_scaler=self.grad_scaler if self.enable_fp16 else None)

        # 根据warmup配置，设置warmup
        if self.warmup_type:
            total_steps = len(train_texts) // batch_size * epoch
            num_warmup_steps = self.warmup_step_num if isinstance(self.warmup_step_num, int) else \
                int(total_steps * self.warmup_step_num)
            assert num_warmup_steps <= total_steps, \
                'num_warmup_steps {} is too large, more than total_steps {}'.format(num_warmup_steps, total_steps)
            if self.warmup_type == 'linear':
                warmup_scheduler = get_linear_schedule_with_warmup(self.optimizer, num_warmup_steps, total_steps)
            elif self.warmup_type == 'cosine':
                warmup_scheduler = get_cosine_schedule_with_warmup(self.optimizer, num_warmup_steps, total_steps)
            else:
                warmup_scheduler = get_constant_schedule_with_warmup(self.optimizer, num_warmup_steps)

        for epoch in range(epoch):
            for batch_idx in range(math.ceil(len(train_texts) / batch_size)):
                text_batch = train_texts[batch_size * batch_idx: batch_size * (batch_idx + 1)]
                labels_batch = train_labels[batch_size * batch_idx: batch_size * (batch_idx + 1)]

                step += 1
                self.model.train()  # 设置为train模式
                self.model.zero_grad()  # 清空梯度

                # 训练
                batch_max_len = max([len(text) for text in text_batch]) + 2  # 长度得加上[CLS]和[SEP]
                if warning_max_len and batch_max_len > 512:
                    logger.warning(
                        'current batch max_len is {}, > 512, which may not be processed by bert!'.format(batch_max_len)
                    )
                batch_input_ids, batch_att_mask, batch_label_ids = self._transform_batch(text_batch,
                                                                                         labels_batch,
                                                                                         max_length=batch_max_len)
                if self.enable_fp16:  # 如果启用混合精度训练，用autocast封装
                    with autocast():
                        best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)
                else:  # 不启用混合精度，正常训练
                    best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)

                # 如果启用并行，需将多张卡返回的sub-batch loss平均
                if self.enable_parallel:
                    loss = loss.mean()

                # 反向传播（混合精度需要放大loss再反传）
                loss.backward() if not self.enable_fp16 else self.grad_scaler.scale(loss).backward()

                # 对抗训练
                if self.adversarial:
                    adv.train(batch_input_ids, batch_att_mask, labels=batch_label_ids)

                # 更新参数
                if self.enable_fp16:
                    self.grad_scaler.step(self.optimizer)
                    self.grad_scaler.update()
                else:
                    self.optimizer.step()

                # 如果启用warmup，更新lr
                if self.warmup_type:
                    warmup_scheduler.step()

                # 计算train acc、valid acc
                train_acc = self._get_acc_one_step(best_paths, batch_label_ids)
                valid_acc, valid_loss = self.validate(validate_texts, validate_labels, batch_size, warning_max_len)

                # 计算连续max_loss_num个平均valid_loss
                loss_buff.append(valid_loss)
                if len(loss_buff) > max_loss_num:
                    loss_buff = loss_buff[-max_loss_num:]
                avg_loss = sum(loss_buff) / len(loss_buff) if len(loss_buff) == max_loss_num else None

                logger.info(
                    'epoch %d, step %d, train loss %.4f, train acc %.4f, valid loss %.4f valid acc %.4f, '
                    'last %d avg valid loss %s' % (
                        epoch, step, loss, train_acc, valid_loss, valid_acc, max_loss_num,
                        '%.4f' % avg_loss if avg_loss else avg_loss
                    )
                )

                # 如果avg_loss在降，保存模型
                if avg_loss and avg_loss < best_loss:
                    best_loss = avg_loss
                    # 根据是否启用并行，获得state_dict
                    state_dict = self.model.state_dict() if not self.enable_parallel else self.model.module.state_dict()
                    torch.save(state_dict, '{}/{}'.format(self.model_dir, self.ckpt_name))
                    logger.info("model saved")

        logger.info("finished")

    def validate(self, validate_texts, validate_labels, sample_size=100, warning_max_len=True):
        """验证
        Args:
            validate_texts: list[str]
            validate_labels: list[str]
            sample_size: int 采样大小(使用全量验证集较慢，这里每次随机采样sample_size个样本做验证)
            warning_max_len: bool 是否警告max_len超过512（普通bert不能处理超过512，除了一些变体如longformer）
        Returns:
            float 验证集上acc, loss
        """
        # 设置为evaluate模式
        self.model.eval()

        # 随机采样sample_size个样本
        batch_texts, batch_labels = [
            return_val[:sample_size] for return_val in shuffle(validate_texts, validate_labels)
        ]

        # 计算valid acc, valid loss
        batch_max_len = max([len(text) for text in batch_texts]) + 2
        if warning_max_len and batch_max_len > 512:
            logger.warning(
                'current batch max_len is {}, > 512, which may not be processed by bert!'.format(batch_max_len)
            )
        with torch.no_grad():  # eval时不计算梯度
            batch_input_ids, batch_att_mask, batch_label_ids = self._transform_batch(batch_texts,
                                                                                     batch_labels,
                                                                                     max_length=batch_max_len)
            # 如果启用混合精度训练，用autocast封装，并使用grad_scaler放大loss
            if self.enable_fp16:
                with autocast():
                    best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)
            else:
                best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)

            # 如果启用并行，需将多张卡返回的sub-batch loss平均
            if self.enable_parallel:
                loss = loss.mean()

            acc = self._get_acc_one_step(best_paths, batch_label_ids)
            return acc, loss

    def _get_acc_one_step(self, labels_predict_batch, labels_batch):
        """计算一个batch的acc, correct_label_num / total_label_num"""
        acc = (labels_predict_batch == labels_batch).sum().item() / labels_batch.shape[0]
        return float(acc)

import json
import math
import os
import random
import re

import numpy as np
import torch
import torch.cuda
from sklearn.utils import shuffle
from torch.cuda.amp import autocast, GradScaler
from torch.nn import DataParallel
from transformers import AdamW
from transformers import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup, \
    get_linear_schedule_with_warmup

from easy_bert import logger
from easy_bert.adversarial import FGM, PGD
from easy_bert.base.base_trainer import BaseTrainer
from easy_bert.bert4classification.classification_model import ClassificationModel
from easy_bert.vocab import Vocab


class ClassificationTrainer(BaseTrainer):
    def __init__(self, pretrained_model_dir, model_dir, learning_rate=5e-5, ckpt_name='bert_model.bin',
                 vocab_name='vocab.json', enable_parallel=False, adversarial=None, dropout_rate=0.5,
                 loss_type='cross_entropy_loss', focal_loss_gamma=2, focal_loss_alpha=None, random_seed=0,
                 warmup_type=None, warmup_step_num=10, enable_fp16=False,
                 add_on=None, rnn_hidden=256, rnn_lr=1e-3,
                 load_last_ckpt=False, layer_wise_lr_decay=False, lr_decay_rate=0.95):
        # 设置目录
        self.pretrained_model_dir, self.model_dir = pretrained_model_dir, model_dir
        self.ckpt_name, self.vocab_name = ckpt_name, vocab_name

        self.enable_parallel = enable_parallel

        self.load_last_ckpt = load_last_ckpt  # 是否加载上次模型

        # 设置随机种子
        self.random_seed = random_seed
        self._set_random_seed(random_seed)

        # 对抗训练配置
        self.adversarial = adversarial
        assert adversarial in (None, 'fgm', 'pgd')

        # loss配置
        self.loss_type = loss_type
        self.focal_loss_gamma = focal_loss_gamma
        self.focal_loss_alpha = focal_loss_alpha

        # 自动获取当前设备
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        self.dropout_rate = dropout_rate
        self.learning_rate = learning_rate
        self.batch_size = None
        self.epoch = None

        # 混合精度配置
        if enable_fp16:
            self.grad_scaler = GradScaler()  # 设置梯度缩放
        self.enable_fp16 = enable_fp16

        # warmup配置
        assert warmup_type in ('constant', 'cosine', 'linear', None)
        assert type(warmup_step_num) in (int, float)
        self.warmup_type = warmup_type
        self.warmup_step_num = warmup_step_num

        # 附加层配置（add-on）
        self.add_on = add_on
        self.rnn_hidden, self.rnn_lr = rnn_hidden, rnn_lr

        # llrd配置
        self.layer_wise_lr_decay = layer_wise_lr_decay
        self.lr_decay_rate = lr_decay_rate

        self.vocab = Vocab()

    def _set_random_seed(self, seed):
        """针对torch torch.cuda numpy random分别设定随机种子"""
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        np.random.seed(seed)
        random.seed(seed)

    def _build_model(self):
        """构建bert模型"""
        # 实例化ClassificationModel模型
        self.model = ClassificationModel(
            self.pretrained_model_dir, self.vocab.label_size, dropout_rate=self.dropout_rate,
            loss_type=self.loss_type, focal_loss_alpha=self.focal_loss_alpha, focal_loss_gamma=self.focal_loss_gamma,
            add_on=self.add_on, rnn_hidden=self.rnn_hidden
        )

        # 如果启用增量训练，预加载之前训练好的模型
        if self.load_last_ckpt and os.path.exists('{}/{}'.format(self.model_dir, self.ckpt_name)):
            self.model.load_state_dict(
                torch.load('{}/{}'.format(self.model_dir, self.ckpt_name), map_location=self.device)
            )
            logger.info('load last ckpt {}/{} success'.format(self.model_dir, self.ckpt_name))

        # 设置AdamW优化器
        no_decay = ["bias", "LayerNorm.weight"]  # bias和LayerNorm不使用正则化
        # 参数分组
        embedding_parameters, encoder_parameters = [], {}
        linear_parameters, rnn_parameters = [], []
        for name, param in self.model.named_parameters():
            if 'rnn' in name:
                rnn_parameters.append((name, param))
            elif 'embeddings' in name:
                embedding_parameters.append((name, param))
            elif 'encoder.layer' in name:
                layer_no = int(re.search('encoder\.layer\.(\d+)', name).group(1))
                if layer_no not in encoder_parameters:
                    encoder_parameters[layer_no] = []
                encoder_parameters[layer_no].append((name, param))
            else:
                linear_parameters.append((name, param))
        # 为不同参数组设置不同的lr
        num_hidden_layers = self.model.bert_model.config.num_hidden_layers
        optimizer_grouped_parameters = [
            {'params': [p for n, p in embedding_parameters if not any(nd in n for nd in no_decay)],
             'weight_decay': 0.01,
             'lr': self.learning_rate * self.lr_decay_rate ** num_hidden_layers if self.layer_wise_lr_decay else self.learning_rate},
            {'params': [p for n, p in embedding_parameters if any(nd in n for nd in no_decay)],
             'weight_decay': 0.0,
             'lr': self.learning_rate * self.lr_decay_rate ** num_hidden_layers if self.layer_wise_lr_decay else self.learning_rate},
        ]
        for layer_no in range(num_hidden_layers):
            optimizer_grouped_parameters.extend([
                {'params': [p for n, p in encoder_parameters[layer_no] if not any(nd in n for nd in no_decay)],
                 'weight_decay': 0.01,
                 'lr': self.learning_rate * self.lr_decay_rate ** (
                         num_hidden_layers - layer_no - 1) if self.layer_wise_lr_decay else self.learning_rate},
                {'params': [p for n, p in encoder_parameters[layer_no] if any(nd in n for nd in no_decay)],
                 'weight_decay': 0.0,
                 'lr': self.learning_rate * self.lr_decay_rate ** (
                         num_hidden_layers - layer_no - 1) if self.layer_wise_lr_decay else self.learning_rate},
            ])
        optimizer_grouped_parameters.extend([
            {'params': [p for n, p in linear_parameters], 'lr': self.learning_rate},
            {'params': [p for n, p in rnn_parameters], 'lr': self.rnn_lr}  # rnn layer设置自己的lr
        ])
        self.optimizer = AdamW(optimizer_grouped_parameters, lr=self.learning_rate)

        # 使用bert的vocab更新Vocab对象
        self.vocab.set_vocab2id(self.model.get_bert_tokenizer().vocab)
        self.vocab.set_id2vocab({_id: char for char, _id in self.vocab.vocab2id.items()})
        self.vocab.set_unk_vocab_id(self.vocab.vocab2id['[UNK]'])
        self.vocab.set_pad_vocab_id(self.vocab.vocab2id['[PAD]'])

        # 启用并行，使用DataParallel封装model
        if self.enable_parallel:
            self.model = DataParallel(self.model, device_ids=list(range(torch.cuda.device_count())))

        # 将模型拷贝到当前设备
        self.model.to(self.device)

    def _save_config(self):
        """保存训练参数配置为json文件"""
        config = {
            'vocab_size': self.vocab.vocab_size,
            'label_size': self.vocab.label_size,
            'learning_rate': self.learning_rate,
            'batch_size': self.batch_size,
            'dropout_rate': self.dropout_rate,
            'epoch': self.epoch,
            'ckpt_name': self.ckpt_name,
            'vocab_name': self.vocab_name,
            'enable_parallel': self.enable_parallel,
            'enable_fp16': self.enable_fp16,
            'adversarial': self.adversarial,
            'loss_type': self.loss_type,
            'focal_loss_gamma': self.focal_loss_gamma,
            'focal_loss_alpha': self.focal_loss_alpha,
            'pretrained_model': os.path.basename(self.pretrained_model_dir),
            'add_on': self.add_on,
            'rnn_hidden': self.rnn_hidden,
            'rnn_lr': self.rnn_lr,
            'random_seed': self.random_seed,
            'layer_wise_lr_decay': self.layer_wise_lr_decay,
            'lr_decay_rate': self.lr_decay_rate,
        }
        with open('{}/train_config.json'.format(self.model_dir), 'w') as f:
            f.write(json.dumps(config, indent=4))

    def _transform_batch(self, batch_texts, batch_labels, max_length=512):
        """将batch的文本及labels转换为bert的输入tensor形式"""
        batch_input_ids, batch_att_mask, batch_label_ids = [], [], []
        for text, label in zip(batch_texts, batch_labels):  # 迭代每条样本
            assert isinstance(text, str)
            # 根据是否并行，获取bert_tokenizer
            bert_tokenizer = self.model.bert_tokenizer if not self.enable_parallel else self.model.module.bert_tokenizer
            encoded_dict = bert_tokenizer.encode_plus(text, max_length=max_length, padding='max_length',
                                                      return_tensors='pt', truncation=True)
            batch_input_ids.append(encoded_dict['input_ids'])
            batch_att_mask.append(encoded_dict['attention_mask'])
            batch_label_ids.append(self.vocab.tag2id[label])
        batch_input_ids, batch_att_mask, batch_label_ids = \
            torch.cat(batch_input_ids), torch.cat(batch_att_mask), torch.LongTensor(batch_label_ids)

        # 将数据拷贝到当前设备
        batch_input_ids, batch_att_mask, batch_label_ids = \
            batch_input_ids.to(self.device), batch_att_mask.to(self.device), batch_label_ids.to(self.device)

        return batch_input_ids, batch_att_mask, batch_label_ids

    def train(self, train_texts, train_labels, validate_texts, validate_labels,
              batch_size=30, epoch=10, max_loss_num=10, warning_max_len=True):
        """训练
        Args:
            train_texts: list[str] 训练集样本
            train_labels: list[str] 训练集标签
            validate_texts: list[str] 验证集样本
            validate_labels: list[str] 验证集标签
            batch_size: int
            epoch: int
            max_loss_num: int 连续max_loss_num个validate loss平均，根据其保存模型
            warning_max_len: bool 是否警告max_len超过512（普通bert不能处理超过512，除了一些变体如longformer）
        """
        self.batch_size = batch_size
        self.epoch = epoch

        # 构建词库表（只构建label id映射，词库bert已有）、构建模型
        self.vocab.build_vocab(labels=train_labels, build_texts=False, with_build_in_tag_id=False)
        self._build_model()
        # 保存词库、保存训练config
        self.vocab.save_vocab('{}/{}'.format(self.model_dir, self.vocab_name))
        self._save_config()

        logger.info('train samples: {}, validate samples: {}'.format(len(train_texts), len(validate_texts)))

        best_loss = float("inf")  # 当前最小的loss
        loss_buff = []  # 保存最近max_loss_num的个valid loss

        step = 0

        # 根据对抗配置，实例化不同的对抗实例
        if self.adversarial:
            logger.info('enable adversarial training')
            adv_class = {'fgm': FGM, 'pgd': PGD}[self.adversarial]
            adv = adv_class(self.model, grad_scaler=self.grad_scaler if self.enable_fp16 else None)

        # 根据warmup配置，设置warmup
        if self.warmup_type:
            total_steps = len(train_texts) // batch_size * epoch
            num_warmup_steps = self.warmup_step_num if isinstance(self.warmup_step_num, int) else \
                int(total_steps * self.warmup_step_num)
            assert num_warmup_steps <= total_steps, \
                'num_warmup_steps {} is too large, more than total_steps {}'.format(num_warmup_steps, total_steps)
            if self.warmup_type == 'linear':
                warmup_scheduler = get_linear_schedule_with_warmup(self.optimizer, num_warmup_steps, total_steps)
            elif self.warmup_type == 'cosine':
                warmup_scheduler = get_cosine_schedule_with_warmup(self.optimizer, num_warmup_steps, total_steps)
            else:
                warmup_scheduler = get_constant_schedule_with_warmup(self.optimizer, num_warmup_steps)

        for epoch in range(epoch):
            for batch_idx in range(math.ceil(len(train_texts) / batch_size)):
                text_batch = train_texts[batch_size * batch_idx: batch_size * (batch_idx + 1)]
                labels_batch = train_labels[batch_size * batch_idx: batch_size * (batch_idx + 1)]

                step += 1
                self.model.train()  # 设置为train模式
                self.model.zero_grad()  # 清空梯度

                # 训练
                batch_max_len = max([len(text) for text in text_batch]) + 2  # 长度得加上[CLS]和[SEP]
                if warning_max_len and batch_max_len > 512:
                    logger.warning(
                        'current batch max_len is {}, > 512, which may not be processed by bert!'.format(batch_max_len)
                    )
                batch_input_ids, batch_att_mask, batch_label_ids = self._transform_batch(text_batch,
                                                                                         labels_batch,
                                                                                         max_length=batch_max_len)
                if self.enable_fp16:  # 如果启用混合精度训练，用autocast封装
                    with autocast():
                        best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)
                else:  # 不启用混合精度，正常训练
                    best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)

                # 如果启用并行，需将多张卡返回的sub-batch loss平均
                if self.enable_parallel:
                    loss = loss.mean()

                # 反向传播（混合精度需要放大loss再反传）
                loss.backward() if not self.enable_fp16 else self.grad_scaler.scale(loss).backward()

                # 对抗训练
                if self.adversarial:
                    adv.train(batch_input_ids, batch_att_mask, labels=batch_label_ids)

                # 更新参数
                if self.enable_fp16:
                    self.grad_scaler.step(self.optimizer)
                    self.grad_scaler.update()
                else:
                    self.optimizer.step()

                # 如果启用warmup，更新lr
                if self.warmup_type:
                    warmup_scheduler.step()

                # 计算train acc、valid acc
                train_acc = self._get_acc_one_step(best_paths, batch_label_ids)
                valid_acc, valid_loss = self.validate(validate_texts, validate_labels, batch_size, warning_max_len)

                # 计算连续max_loss_num个平均valid_loss
                loss_buff.append(valid_loss)
                if len(loss_buff) > max_loss_num:
                    loss_buff = loss_buff[-max_loss_num:]
                avg_loss = sum(loss_buff) / len(loss_buff) if len(loss_buff) == max_loss_num else None

                logger.info(
                    'epoch %d, step %d, train loss %.4f, train acc %.4f, valid loss %.4f valid acc %.4f, '
                    'last %d avg valid loss %s' % (
                        epoch, step, loss, train_acc, valid_loss, valid_acc, max_loss_num,
                        '%.4f' % avg_loss if avg_loss else avg_loss
                    )
                )

                # 如果avg_loss在降，保存模型
                if avg_loss and avg_loss < best_loss:
                    best_loss = avg_loss
                    # 根据是否启用并行，获得state_dict
                    state_dict = self.model.state_dict() if not self.enable_parallel else self.model.module.state_dict()
                    torch.save(state_dict, '{}/{}'.format(self.model_dir, self.ckpt_name))
                    logger.info("model saved")

        logger.info("finished")

    def validate(self, validate_texts, validate_labels, sample_size=100, warning_max_len=True):
        """验证
        Args:
            validate_texts: list[str]
            validate_labels: list[str]
            sample_size: int 采样大小(使用全量验证集较慢，这里每次随机采样sample_size个样本做验证)
            warning_max_len: bool 是否警告max_len超过512（普通bert不能处理超过512，除了一些变体如longformer）
        Returns:
            float 验证集上acc, loss
        """
        # 设置为evaluate模式
        self.model.eval()

        # 随机采样sample_size个样本
        batch_texts, batch_labels = [
            return_val[:sample_size] for return_val in shuffle(validate_texts, validate_labels)
        ]

        # 计算valid acc, valid loss
        batch_max_len = max([len(text) for text in batch_texts]) + 2
        if warning_max_len and batch_max_len > 512:
            logger.warning(
                'current batch max_len is {}, > 512, which may not be processed by bert!'.format(batch_max_len)
            )
        with torch.no_grad():  # eval时不计算梯度
            batch_input_ids, batch_att_mask, batch_label_ids = self._transform_batch(batch_texts,
                                                                                     batch_labels,
                                                                                     max_length=batch_max_len)
            # 如果启用混合精度训练，用autocast封装，并使用grad_scaler放大loss
            if self.enable_fp16:
                with autocast():
                    best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)
            else:
                best_paths, loss = self.model(batch_input_ids, batch_att_mask, labels=batch_label_ids)

            # 如果启用并行，需将多张卡返回的sub-batch loss平均
            if self.enable_parallel:
                loss = loss.mean()

            acc = self._get_acc_one_step(best_paths, batch_label_ids)
            return acc, loss

    def _get_acc_one_step(self, labels_predict_batch, labels_batch):
        """计算一个batch的acc, correct_label_num / total_label_num"""
        acc = (labels_predict_batch == labels_batch).sum().item() / labels_batch.shape[0]
        return float(acc)
# 设置目录 # 是否加载上次模型 # 设置随机种子 # 对抗训练配置 # loss配置 # 自动获取当前设备 # 混合精度配置 # 设置梯度缩放 # warmup配置 # 附加层配置（add-on） # llrd配置 针对torch torch.cuda numpy random分别设定随机种子 构建bert模型 # 实例化ClassificationModel模型 # 如果启用增量训练，预加载之前训练好的模型 # 设置AdamW优化器 # bias和LayerNorm不使用正则化 # 参数分组 # 为不同参数组设置不同的lr # rnn layer设置自己的lr # 使用bert的vocab更新Vocab对象 # 启用并行，使用DataParallel封装model # 将模型拷贝到当前设备 保存训练参数配置为json文件 将batch的文本及labels转换为bert的输入tensor形式 # 迭代每条样本 # 根据是否并行，获取bert_tokenizer # 将数据拷贝到当前设备 训练         Args:             train_texts: list[str] 训练集样本             train_labels: list[str] 训练集标签             validate_texts: list[str] 验证集样本             validate_labels: list[str] 验证集标签             batch_size: int             epoch: int             max_loss_num: int 连续max_loss_num个validate loss平均，根据其保存模型             warning_max_len: bool 是否警告max_len超过512（普通bert不能处理超过512，除了一些变体如longformer） # 构建词库表（只构建label id映射，词库bert已有）、构建模型 # 保存词库、保存训练config # 当前最小的loss # 保存最近max_loss_num的个valid loss # 根据对抗配置，实例化不同的对抗实例 # 根据warmup配置，设置warmup # 设置为train模式 # 清空梯度 # 训练 # 长度得加上[CLS]和[SEP] # 如果启用混合精度训练，用autocast封装 # 不启用混合精度，正常训练 # 如果启用并行，需将多张卡返回的sub-batch loss平均 # 反向传播（混合精度需要放大loss再反传） # 对抗训练 # 更新参数 # 如果启用warmup，更新lr # 计算train acc、valid acc # 计算连续max_loss_num个平均valid_loss # 如果avg_loss在降，保存模型 # 根据是否启用并行，获得state_dict 验证         Args:             validate_texts: list[str]             validate_labels: list[str]             sample_size: int 采样大小(使用全量验证集较慢，这里每次随机采样sample_size个样本做验证)             warning_max_len: bool 是否警告max_len超过512（普通bert不能处理超过512，除了一些变体如longformer）         Returns:             float 验证集上acc, loss # 设置为evaluate模式 # 随机采样sample_size个样本 # 计算valid acc, valid loss # eval时不计算梯度 # 如果启用混合精度训练，用autocast封装，并使用grad_scaler放大loss # 如果启用并行，需将多张卡返回的sub-batch loss平均 计算一个batch的acc, correct_label_num / total_label_num