大模型中的提示学习——情感预测示例项目

开始对评论数据进行情感分析，输入CTRL + C，则退出
输入的评论数据为：这家酒店是我在携程定的酒店里面是有点问题的。
情感极性为：负向
输入的评论数据为：好吧，就这样吧。
情感极性为：负向
输入的评论数据为：好吧
情感极性为：负向
输入的评论数据为：好
情感极性为：负向
输入的评论数据为：豪酒店啊！
情感极性为：负向
输入的评论数据为：好酒店啊！
情感极性为：正向
输入的评论数据为：豪华酒店！
情感极性为：正向
输入的评论数据为：华酒店！
情感极性为：负向
输入的评论数据为：华丽酒店！
情感极性为：正向

# -*- coding:utf-8 -*-
from torch.nn import CrossEntropyLoss
import torch.nn as nn
import torch
from transformers.models.bert.modeling_bert import BertModel, BertOnlyMLMHead, BertPreTrainedModel


class PromptModel(BertPreTrainedModel):
    """Prompt分类模型"""

    def __init__(self, config):
        super().__init__(config)
        """
        初始化函数
        Args:
            config: 配置参数
        """
        self.bert = BertModel(config, add_pooling_layer=False)
        self.cls = BertOnlyMLMHead(config)

    def forward(self, input_ids, attention_mask, mask_index, token_handler, words_ids, words_ids_mask,
                label=None):
        """
        前向函数，计算Prompt模型预测结果
        Args:
            input_ids:
            attention_mask:
            mask_index:
            token_handler:
            words_ids:
            words_ids_mask:
            label:

        Returns:

        """
        # 获取BERT模型的输出结果
        sequence_output = self.bert(input_ids=input_ids, attention_mask=attention_mask)[0]
        # 经过一个全连接层，获取隐层节点状态中的每一个位置的词表
        logits = self.cls(sequence_output)
        # 获取批次数据中每个样本内容对应mask位置标记
        logits_shapes = logits.shape
        mask = mask_index + torch.range(0, logits_shapes[0] - 1, dtype=torch.long, device=logits.device) * \
               logits_shapes[1]
        mask = mask.reshape([-1, 1]).repeat([1, logits_shapes[2]])
        # 获取每个mask标记对应的logits内容
        mask_logits = logits.reshape([-1, logits_shapes[2]]).gather(0, mask).reshape(-1, logits_shapes[2])
        # 获取答案空间映射的标签向量
        label_words_logits = self.process_logits(mask_logits, token_handler, words_ids, words_ids_mask)
        # 将其进行归一化以及获取对应标签
        score = torch.nn.functional.softmax(label_words_logits, dim=-1)
        pre_label = torch.argmax(label_words_logits, dim=1)
        outputs = (score, pre_label)
        # 当label不为空时，计算损失值
        if label is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(label_words_logits, label)
            outputs = (loss,) + outputs
        return outputs

    def process_logits(self, mask_logits, token_handler, words_ids, words_ids_mask):
        """
        获取答案空间映射的标签向量，用于分类判断
        Args:
            mask_logits: mask位置信息
            token_handler: 多token操作策略，包含first、mask和mean
            words_ids: 标签词id矩阵
            words_ids_mask: 标签词id掩码矩阵

        Returns:

        """
        # 获取标签词id及掩码矩阵
        label_words_ids = nn.Parameter(words_ids, requires_grad=False)
        label_words_mask = nn.Parameter(torch.clamp(words_ids_mask.sum(dim=-1), max=1), requires_grad=False)
        # 获取mask位置上标签词向量
        label_words_logits = mask_logits[:, label_words_ids]
        # 根据多token操作策略进行标签词向量构建
        if token_handler == "first":
            label_words_logits = label_words_logits.select(dim=-1, index=0)
        elif token_handler == "max":
            label_words_logits = label_words_logits - 1000 * (1 - words_ids_mask.unsqueeze(0))
            label_words_logits = label_words_logits.max(dim=-1).values
        elif token_handler == "mean":
            label_words_logits = (label_words_logits * words_ids_mask.unsqueeze(0)).sum(dim=-1) / (
                    words_ids_mask.unsqueeze(0).sum(dim=-1) + 1e-15)
        # 将填充的位置进行掩码
        label_words_logits -= 10000 * (1 - label_words_mask)
        # 最终获取mask标记对应的答案空间映射向量
        label_words_logits = (label_words_logits * label_words_mask).sum(-1) / label_words_mask.sum(-1)
        return label_words_logits

def train(model, device, train_data, test_data, args, tokenizer):
    """
    训练模型
    Args:
        model: 模型
        device: 设备信息
        train_data: 训练数据类
        test_data: 测试数据类
        args: 训练参数配置信息
        tokenizer: 分词器
    Returns:
    """
    tb_write = SummaryWriter()
    if args.gradient_accumulation_steps < 1:
        raise ValueError("gradient_accumulation_steps参数无效，必须大于等于1")
    # 计算真实的训练batch_size大小
    train_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps)
    train_sampler = RandomSampler(train_data)
    # 构造训练所需的data_loader
    train_data_loader = DataLoader(train_data, sampler=train_sampler,
                                   batch_size=train_batch_size, collate_fn=collate_func)
    total_steps = int(len(train_data_loader) * args.num_train_epochs / args.gradient_accumulation_steps)
    logger.info("总训练步数为:{}".format(total_steps))
    model.to(device)
    # 获取模型所有参数
    param_optimizer = list(model.named_parameters())
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(
            nd in n for nd in args.requires_grad_params)], 'weight_decay': 0.01},
        {'params': [p for n, p in param_optimizer if any(
            nd in n for nd in args.requires_grad_params)], 'weight_decay': 0.0}
    ]

    # 冻结不训练的参数
    for name, param in model.named_parameters():
        if not any(r_name in name for r_name in args.requires_grad_params):
            param.requires_grad = False

    # 验证是否冻结成功
    requires_grad_params = []
    for name, param in model.named_parameters():
        if param.requires_grad:
            requires_grad_params.append(name)
            print("需要训练参数为{}，大小为{}".format(name, param.size()))
    # 设置优化器
    optimizer = AdamW(optimizer_grouped_parameters,
                      lr=args.learning_rate, eps=args.adam_epsilon)
    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=int(args.warmup_proportion * total_steps),
                                                num_training_steps=total_steps)
    # 清空cuda缓存
    torch.cuda.empty_cache()
    # 将模型调至训练状态
    model.train()
    tr_loss, logging_loss, min_loss = 0.0, 0.0, 0.0
    global_step = 0
    words_ids = train_data.words_ids.to(device)
    words_ids_mask = train_data.words_ids_mask.to(device)
    # 开始训练模型
    for iepoch in trange(0, int(args.num_train_epochs), desc="Epoch", disable=False):
        iter_bar = tqdm(train_data_loader, desc="Iter (loss=X.XXX)", disable=False)
        for step, batch in enumerate(iter_bar):
            # 获取模型训练每个批次所需的输入内容，并放到对应设备上
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            mask_index = batch["mask_index"].to(device)
            label = batch["label"].to(device)

            # 获取训练结果
            outputs = model.forward(input_ids=input_ids, attention_mask=attention_mask, mask_index=mask_index,
                                    token_handler=args.token_handler,
                                    words_ids=words_ids, words_ids_mask=words_ids_mask,
                                    label=label)
            loss = outputs[0]
            tr_loss += loss.item()
            # 将损失值放到Iter中，方便观察
            iter_bar.set_description("Iter (loss=%5.3f)" % loss.item())
            # 判断是否进行梯度累积，如果进行，则将损失值除以累积步数
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps
            # 损失进行回传
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
            # 当训练步数整除累积步数时，进行参数优化
            if (step + 1) % args.gradient_accumulation_steps == 0:
                optimizer.step()
                scheduler.step()
                optimizer.zero_grad()
                global_step += 1
                # 如果步数整除logging_steps，则记录学习率和训练集损失值
                if args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    tb_write.add_scalar("lr", scheduler.get_lr()[0], global_step)
                    tb_write.add_scalar("train_loss", (tr_loss - logging_loss) /
                                        (args.logging_steps * args.gradient_accumulation_steps), global_step)
                    logging_loss = tr_loss

        # 每个Epoch对模型进行一次测试，记录测试集的损失
        eval_loss, eval_acc = evaluate(model, device, test_data, args)
        tb_write.add_scalar("test_loss", eval_loss, global_step)
        tb_write.add_scalar("test_acc", eval_acc, global_step)
        print("test_loss: {}, test_acc:{}".format(eval_loss, eval_acc))
        model.train()
        # 每个epoch进行完，则保存模型
        output_dir = os.path.join(args.output_dir, "checkpoint-{}".format(global_step))
        model_to_save = model.module if hasattr(model, "module") else model
        model_to_save.save_pretrained(output_dir)
        tokenizer.save_pretrained(output_dir)
        # 清空cuda缓存
        torch.cuda.empty_cache()