【案例】新零售评价决策系统 · P-Tuning 软模板

# -*- coding:utf-8 -*-
"""为什么偏偏挑 [unused1] 这种 token 当伪 token。

bert-base-chinese 的词表里留了一批 [unusedN] 占位符，
预训练阶段它们从未在真实语料里出现过，embedding 基本没被更新过——
相当于一块干净的画布。用「的」「了」这种高频字当伪 token，
梯度要先把它原有的语义抹掉，收敛更慢也更不稳。
"""


def unused_ids(n=6, base=1):
    """中文 BERT 词表里 [unused1] 的 id 就是 1，往后连号。"""
    return [base + i for i in range(n)]


if __name__ == '__main__':
    print('[unused1..6] 的 token id:', unused_ids(6))
    assert unused_ids(6) == [1, 2, 3, 4, 5, 6]

    # 词表结构：0=[PAD], 1..99=[unusedN], 100=[UNK], 101=[CLS], 102=[SEP], 103=[MASK]
    special = {0: '[PAD]', 100: '[UNK]', 101: '[CLS]', 102: '[SEP]', 103: '[MASK]'}
    print('\nid   token')
    for i in sorted(special):
        print('%-4d %s' % (i, special[i]))
    print('1-99 [unused1] ~ [unused99]  <- 可挪用的空位一共 99 个')

    # 越界检查：p_embedding_num 不能超过 99
    for p_num in (6, 32, 99):
        ids = unused_ids(p_num)
        assert max(ids) <= 99, 'p_embedding_num 最多 99'
        print('p_embedding_num=%-3d 占用 id %d..%d' % (p_num, ids[0], ids[-1]))

    try:
        p_num = 120
        ids = unused_ids(p_num)
        assert max(ids) <= 99
    except AssertionError:
        print('\np_embedding_num=120 会撞上 [UNK] 及之后的真实 token —— 被断言拦住')

    # 伪 token 占的是序列长度预算，不是免费的
    max_seq_len, max_label_len = 512, 2
    for p_num in (6, 32, 99):
        usable = max_seq_len - p_num - max_label_len - 2
        print('p_embedding_num=%-3d 正文还剩 %d 个 token' % (p_num, usable))
    assert 512 - 99 - 2 - 2 == 409
    print('\n用满 99 个伪 token，正文可用长度从 502 掉到 409 —— 长文本任务要算这笔账')

# -*- coding:utf-8 -*-
"""软提示到底是怎么拼进去的：纯标准库复刻四步拼接，可直接运行。

P-Tuning 的全部结构改动就这四步。写错的典型后果是
mask_positions 与实际 [MASK] 的位置对不上，损失算在了正文的字上。
"""

CLS, SEP, PAD, MASK = 101, 102, 0, 103
UNUSED_BASE = 1          # [unused1] 在中文 BERT 词表里的 id 就是 1


def build_input_ids(content_ids, p_num=6, mask_len=2, max_seq_len=24):
    """复刻 data_preprocess.convert_example 里的拼接逻辑。

    参数:
        content_ids: tokenizer 编码后的完整 input_ids，形如 [CLS] ... [SEP]
    """
    p_token_ids = [UNUSED_BASE + i for i in range(p_num)]      # [1..6]
    mask_ids = [MASK] * mask_len
    start_mask_position = 1                                     # 伪 token 插在 [CLS] 之后

    # ① 先把最后一个 token 摘掉（通常是 [SEP]），待会儿补回来
    tmp = content_ids[:-1]
    # ② 给伪 token、mask 和末尾的 [SEP] 腾出位置
    tmp = tmp[:max_seq_len - len(mask_ids) - len(p_token_ids) - 1]
    # ③ 真正的拼接：伪 token + [CLS] + mask + 正文
    tmp = p_token_ids + tmp[:1] + mask_ids + tmp[start_mask_position:]
    # ④ 补回结尾：原句里已有 [SEP] 就补 pad，否则补 [SEP]
    if SEP in tmp:
        input_ids = tmp + [PAD]
    else:
        input_ids = tmp + [SEP]

    mask_positions = [len(p_token_ids) + start_mask_position + i for i in range(mask_len)]
    attention_mask = [1 if t > 0 else 0 for t in input_ids]
    return input_ids, mask_positions, attention_mask


if __name__ == '__main__':
    # 模拟一句 10 个字的评论：[CLS] + 10 个字 + [SEP]
    content = [CLS] + [3000 + i for i in range(10)] + [SEP]
    ids, pos, att = build_input_ids(content, p_num=6, mask_len=2, max_seq_len=24)

    print('拼接后的 input_ids:', ids)
    print('长度:', len(ids))
    print('mask_positions   :', pos)

    # 断言一：前 6 个位置必须是伪 token
    assert ids[:6] == [1, 2, 3, 4, 5, 6], '伪 token 必须排在最前面'
    # 断言二：第 7 个位置是 [CLS]
    assert ids[6] == CLS
    # 断言三：mask 紧跟在 [CLS] 之后，落在下标 7、8
    assert ids[7] == MASK and ids[8] == MASK
    assert pos == [7, 8], 'mask_positions 由公式算出，不是数出来的'
    print('拼接顺序校验通过：[unused1-6] [CLS] [MASK][MASK] 正文 …')

    # 断言四：mask 位置与伪 token 个数直接挂钩，改一个另一个跟着变
    for p_num in (2, 4, 8):
        _, pos_n, _ = build_input_ids(content, p_num=p_num, mask_len=2, max_seq_len=32)
        assert pos_n == [p_num + 1, p_num + 2]
        print('  p_embedding_num=%d -> mask_positions=%s' % (p_num, pos_n))
    print('结论：改 p_embedding_num 必须同步重新生成数据，缓存的旧特征会算错位置')

    # 断言五：attention_mask 必须重算。原 tokenizer 给的那份长度和内容都不对了
    print('attention_mask   :', att)
    assert len(att) == len(ids)
    assert att[:9] == [1] * 9, '伪 token 与 mask 都要被看见'
    assert att[-1] == 0 or ids[-1] != PAD
    print('全部断言通过')

# -*- coding:utf-8 -*-
"""为什么 attention_mask 必须重算：纯标准库对照实验，可直接运行。

tokenizer 给出的 attention_mask 是针对「原始序列」的。
拼进 6 个伪 token 和 2 个 mask 之后，序列长度和内容都变了，
直接沿用那一份，模型会把伪 token 当成 padding 忽略掉——软提示等于没加。
"""

CLS, SEP, PAD, MASK = 101, 102, 0, 103


def recompute(input_ids):
    """真实代码里的 np.where(np.array(ids) > 0, 1, 0)。"""
    return [1 if t > 0 else 0 for t in input_ids]


if __name__ == '__main__':
    # tokenizer 对一句 5 个字的评论、padding 到 12 的输出
    orig_ids = [CLS, 3001, 3002, 3003, 3004, 3005, SEP] + [PAD] * 5
    orig_att = [1] * 7 + [0] * 5
    print('原始 input_ids     :', orig_ids)
    print('原始 attention_mask:', orig_att)
    assert len(orig_ids) == len(orig_att) == 12

    # 拼接后的序列：6 个伪 token + [CLS] + 2 个 mask + 正文 + [SEP]
    new_ids = [1, 2, 3, 4, 5, 6, CLS, MASK, MASK, 3001, 3002, 3003, 3004, 3005, SEP, PAD]
    print('\n拼接后 input_ids   :', new_ids)
    print('长度从 %d 变成 %d' % (len(orig_ids), len(new_ids)))

    # 错误做法：沿用旧的 attention_mask
    wrong = orig_att + [0] * (len(new_ids) - len(orig_att))
    print('\n沿用旧 mask（错误）:', wrong)
    ignored = [new_ids[i] for i in range(len(new_ids)) if wrong[i] == 0 and new_ids[i] != PAD]
    print('被错误忽略的 token :', ignored)
    assert 3004 in ignored and 3005 in ignored, '正文尾部被当成 padding 丢掉了'
    assert SEP in ignored

    # 正确做法：按新的 input_ids 重算
    right = recompute(new_ids)
    print('\n重算的 mask（正确）:', right)
    assert len(right) == len(new_ids)
    assert right[:6] == [1] * 6, '伪 token 必须可见，否则软提示白拼'
    assert right[-1] == 0, 'padding 必须被屏蔽'
    assert sum(right) == len(new_ids) - 1
    print('可见 token 数:', sum(right), '/', len(new_ids))

    # 这个写法有个前提：pad_token_id 必须是 0
    print('\n前提：pad_token_id == 0，才能用 “>0” 判可见性')
    print('若换一个 pad_token_id 不为 0 的分词器，这行必须改成与 pad_token_id 比较')
    print('全部断言通过')

# -*- coding:utf-8 -*-
"""软提示到底是怎么拼进去的：纯标准库复刻四步拼接，可直接运行。

P-Tuning 的全部结构改动就这四步。写错的典型后果是
mask_positions 与实际 [MASK] 的位置对不上，损失算在了正文的字上。
"""

CLS, SEP, PAD, MASK = 101, 102, 0, 103
UNUSED_BASE = 1          # [unused1] 在中文 BERT 词表里的 id 就是 1


def build_input_ids(content_ids, p_num=6, mask_len=2, max_seq_len=24):
    """复刻 data_preprocess.convert_example 里的拼接逻辑。

    参数:
        content_ids: tokenizer 编码后的完整 input_ids，形如 [CLS] ... [SEP]
    """
    p_token_ids = [UNUSED_BASE + i for i in range(p_num)]      # [1..6]
    mask_ids = [MASK] * mask_len
    start_mask_position = 1                                     # 伪 token 插在 [CLS] 之后

    # ① 先把最后一个 token 摘掉（通常是 [SEP]），待会儿补回来
    tmp = content_ids[:-1]
    # ② 给伪 token、mask 和末尾的 [SEP] 腾出位置
    tmp = tmp[:max_seq_len - len(mask_ids) - len(p_token_ids) - 1]
    # ③ 真正的拼接：伪 token + [CLS] + mask + 正文
    tmp = p_token_ids + tmp[:1] + mask_ids + tmp[start_mask_position:]
    # ④ 补回结尾：原句里已有 [SEP] 就补 pad，否则补 [SEP]
    if SEP in tmp:
        input_ids = tmp + [PAD]
    else:
        input_ids = tmp + [SEP]

    mask_positions = [len(p_token_ids) + start_mask_position + i for i in range(mask_len)]
    attention_mask = [1 if t > 0 else 0 for t in input_ids]
    return input_ids, mask_positions, attention_mask


if __name__ == '__main__':
    # 模拟一句 10 个字的评论：[CLS] + 10 个字 + [SEP]
    content = [CLS] + [3000 + i for i in range(10)] + [SEP]
    ids, pos, att = build_input_ids(content, p_num=6, mask_len=2, max_seq_len=24)

    print('拼接后的 input_ids:', ids)
    print('长度:', len(ids))
    print('mask_positions   :', pos)

    # 断言一：前 6 个位置必须是伪 token
    assert ids[:6] == [1, 2, 3, 4, 5, 6], '伪 token 必须排在最前面'
    # 断言二：第 7 个位置是 [CLS]
    assert ids[6] == CLS
    # 断言三：mask 紧跟在 [CLS] 之后，落在下标 7、8
    assert ids[7] == MASK and ids[8] == MASK
    assert pos == [7, 8], 'mask_positions 由公式算出，不是数出来的'
    print('拼接顺序校验通过：[unused1-6] [CLS] [MASK][MASK] 正文 …')

    # 断言四：mask 位置与伪 token 个数直接挂钩，改一个另一个跟着变
    for p_num in (2, 4, 8):
        _, pos_n, _ = build_input_ids(content, p_num=p_num, mask_len=2, max_seq_len=32)
        assert pos_n == [p_num + 1, p_num + 2]
        print('  p_embedding_num=%d -> mask_positions=%s' % (p_num, pos_n))
    print('结论：改 p_embedding_num 必须同步重新生成数据，缓存的旧特征会算错位置')

    # 断言五：attention_mask 必须重算。原 tokenizer 给的那份长度和内容都不对了
    print('attention_mask   :', att)
    assert len(att) == len(ids)
    assert att[:9] == [1] * 9, '伪 token 与 mask 都要被看见'
    assert att[-1] == 0 or ids[-1] != PAD
    print('全部断言通过')

# -*- coding:utf-8 -*-
"""P-Tuning 的样本编码：把「标签\t正文」变成模型能吃的张量。

和 PET 版的唯一差别在 convert_example 内部——
PET 是「填模板」，这里是「在句首插伪 token」。外层接口完全一致，
所以 data_loader、train、inference 几乎可以原样复用。
"""
import numpy as np


def convert_example(examples: dict,
                    tokenizer,
                    max_seq_len: int,
                    max_label_len: int,
                    p_embedding_num=6,
                    train_mode=True,
                    return_tensor=False) -> dict:
    """
    参数:
        examples: {'text': ['体育	世界杯开赛了', ...]}，train_mode=False 时没有前面的标签列
        p_embedding_num: 伪 token 个数，必须与训练/推理两侧一致
    返回:
        input_ids / token_type_ids / attention_mask / mask_positions /
        mask_labels（仅训练）
    """
    tokenized_output = {
        'input_ids': [],
        'token_type_ids': [],
        'attention_mask': [],
        'mask_positions': [],
        'mask_labels': [],
    }

    for i, example in enumerate(examples['text']):
        try:
            if train_mode:
                label, content = example.strip().split('\t', 1)
            else:
                content, label = example.strip(), None
        except ValueError:
            print('[跳过] 第 %d 条格式不对（缺少 \\t）: %s' % (i, example[:30]))
            continue

        encoded_inputs = tokenizer(
            text=content,
            truncation=True,
            max_length=max_seq_len,
            padding='max_length')

        input_ids = encoded_inputs['input_ids']
        # ① 摘掉末尾 token（一般是 [SEP]），拼完再补回去
        input_ids = input_ids[:-1]
        # ② 给伪 token、mask、[SEP] 腾位置
        input_ids = input_ids[:max_seq_len - len(input_ids) * 0 - max_label_len - p_embedding_num - 1]

        p_token_ids = tokenizer.convert_tokens_to_ids(
            ['[unused%s]' % (i + 1) for i in range(p_embedding_num)])
        mask_ids = [tokenizer.mask_token_id] * max_label_len

        # ③ 伪 token + [CLS] + [MASK]*n + 正文
        input_ids = p_token_ids + input_ids[:1] + mask_ids + input_ids[1:]
        # ④ 补回结尾
        if tokenizer.sep_token_id in input_ids:
            input_ids = input_ids + [tokenizer.pad_token_id]
        else:
            input_ids = input_ids + [tokenizer.sep_token_id]

        # mask 位置由公式算出，不要去序列里 index() 查找：
        # 正文里若本来就含 [MASK] 字面，index() 会找错
        mask_positions = [len(p_token_ids) + 1 + i for i in range(max_label_len)]

        token_type_ids = [0] * len(input_ids)
        # attention_mask 必须按新序列重算，沿用 tokenizer 那份会把伪 token 当 padding
        attention_mask = np.where(np.array(input_ids) > 0, 1, 0).tolist()

        tokenized_output['input_ids'].append(input_ids)
        tokenized_output['token_type_ids'].append(token_type_ids)
        tokenized_output['attention_mask'].append(attention_mask)
        tokenized_output['mask_positions'].append(mask_positions)

        if train_mode:
            mask_labels = tokenizer(text=label)['input_ids'][1:-1]  # 去掉 [CLS] [SEP]
            mask_labels = mask_labels[:max_label_len]
            mask_labels += [tokenizer.pad_token_id] * (max_label_len - len(mask_labels))
            tokenized_output['mask_labels'].append(mask_labels)

    for k, v in tokenized_output.items():
        if not v:
            continue
        if return_tensor:
            import torch
            tokenized_output[k] = torch.LongTensor(v)
        else:
            tokenized_output[k] = np.array(v)
    return {k: v for k, v in tokenized_output.items() if len(v)}


if __name__ == '__main__':
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained('./bert-base-chinese')
    out = convert_example(
        {'text': ['电脑\t这个笔记本电脑散热做得不错，风扇噪音也小']},
        tokenizer, max_seq_len=32, max_label_len=2, p_embedding_num=6)
    print('input_ids  :', out['input_ids'][0][:14], '...')
    print('前 14 个token:', tokenizer.convert_ids_to_tokens(out['input_ids'][0][:14]))
    print('mask 位置  :', out['mask_positions'][0])
    print('标签 ids   :', out['mask_labels'][0])

电脑	电脑
水果	水果,苹果,香蕉,榴莲,西瓜
平板	平板
衣服	衣服
酒店	酒店
洗浴	洗浴
书籍	书籍
蒙牛	蒙牛
手机	手机
电器	电器

# coding:utf-8
"""P-Tuning 路线的配置。与 PET 路线只差两处，其余照搬。

差别一：多了 p_embedding_num —— 伪 token 的个数。
差别二：没有 prompt_file —— 模板不再是人写的文字，没有文件可读。
"""
import os

import torch


class ProjectConfig(object):
    def __init__(self):
        self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        self.root = os.environ.get(
            'PTUNE_PROJECT_ROOT',
            os.path.dirname(os.path.abspath(__file__)),
        )
        self.pre_model = os.path.join(self.root, 'bert-base-chinese')

        # 只要三个数据文件：train / dev / verbalizer。模板没有文件。
        self.train_path = os.path.join(self.root, 'data', 'train.txt')
        self.dev_path = os.path.join(self.root, 'data', 'dev.txt')
        self.verbalizer = os.path.join(self.root, 'data', 'verbalizer.txt')

        # 伪 token 个数：拼在句首的 [unused1] ~ [unused6]。
        # 调大 -> 模板容量更大但更难训、也更占序列长度；调小 -> 表达力不够。
        # 这是 P-Tuning 唯一真正需要调的结构超参数。
        self.p_embedding_num = 6

        self.max_seq_len = 512
        self.batch_size = 8
        self.learning_rate = 5e-5
        self.weight_decay = 0
        self.warmup_ratio = 0.06
        self.max_label_len = 2
        self.epochs = 10
        self.logging_steps = 5
        self.valid_steps = 20
        self.save_dir = os.path.join(self.root, 'checkpoints')


if __name__ == '__main__':
    pc = ProjectConfig()
    print('伪 token 个数:', pc.p_embedding_num)
    print('伪 token 字面:', ['[unused%d]' % (i + 1) for i in range(pc.p_embedding_num)])
    print('mask 个数    :', pc.max_label_len)
    # 这两个数决定了正文能用多长：总长 - 伪 token - mask - [CLS] - [SEP]
    print('正文可用长度 :', pc.max_seq_len - pc.p_embedding_num - pc.max_label_len - 2)

# -*- coding:utf-8 -*-
"""P-Tuning 的 DataLoader。与 PET 版几乎一致，只是 partial 里换了参数。"""
from functools import partial

from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, default_data_collator

from data_preprocess import convert_example
from ptune_config import ProjectConfig

pc = ProjectConfig()
tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)


def get_data():
    dataset = load_dataset('text', data_files={'train': pc.train_path,
                                               'dev': pc.dev_path})
    # 注意这里没有 hard_template：模板不再是文字，而是词表里的 6 个伪 token
    new_func = partial(convert_example,
                       tokenizer=tokenizer,
                       max_seq_len=pc.max_seq_len,
                       max_label_len=pc.max_label_len,
                       p_embedding_num=pc.p_embedding_num)

    dataset = dataset.map(new_func, batched=True)

    train_dataset = dataset['train']
    dev_dataset = dataset['dev']

    train_dataloader = DataLoader(train_dataset,
                                  shuffle=True,
                                  collate_fn=default_data_collator,
                                  batch_size=pc.batch_size)
    dev_dataloader = DataLoader(dev_dataset,
                                collate_fn=default_data_collator,
                                batch_size=pc.batch_size)
    return train_dataloader, dev_dataloader


if __name__ == '__main__':
    train_dataloader, dev_dataloader = get_data()
    print('train batch 数:', len(train_dataloader))
    print('dev   batch 数:', len(dev_dataloader))
    for batch in train_dataloader:
        print({k: tuple(v.shape) for k, v in batch.items()})
        break

# -*- coding:utf-8 -*-
"""P-Tuning 训练主脚本。

⚠️ 与 PET 的 train.py 对照着看：训练循环、优化器、调度器、存盘逻辑
一行都没改。变的只有 import 的 data_loader 和 config——
这就是 P-Tuning 工程上最大的优点：换模板方式不用重写训练代码。
"""
import os
import time

import torch
from tqdm import tqdm
from transformers import AutoModelForMaskedLM, AutoTokenizer, get_scheduler

from data_loader import get_data
from ptune_config import ProjectConfig
from utils.common_utils import convert_logits_to_ids, mlm_loss
from utils.metirc_utils import ClassEvaluator
from utils.verbalizer import Verbalizer

pc = ProjectConfig()


def evaluate_model(model, metric, data_loader, tokenizer, verbalizer):
    model.eval()
    metric.reset()
    with torch.no_grad():
        for batch in tqdm(data_loader, desc='eval'):
            logits = model(input_ids=batch['input_ids'].to(pc.device),
                           token_type_ids=batch['token_type_ids'].to(pc.device),
                           attention_mask=batch['attention_mask'].to(pc.device)).logits
            mask_labels = batch['mask_labels'].numpy().tolist()
            for i in range(len(mask_labels)):
                while tokenizer.pad_token_id in mask_labels[i]:
                    mask_labels[i].remove(tokenizer.pad_token_id)
            mask_labels = [''.join(tokenizer.convert_ids_to_tokens(t)) for t in mask_labels]
            predictions = convert_logits_to_ids(
                logits, batch['mask_positions']).cpu().numpy().tolist()
            predictions = verbalizer.batch_find_main_label(predictions)
            predictions = [ele['label'] for ele in predictions]
            metric.add_batch(pred_batch=predictions, gold_batch=mask_labels)
    eval_metric = metric.compute()
    model.train()
    return (eval_metric['accuracy'], eval_metric['precision'],
            eval_metric['recall'], eval_metric['f1'], eval_metric['class_metrics'])


def model2train():
    model = AutoModelForMaskedLM.from_pretrained(pc.pre_model)
    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model)
    verbalizer = Verbalizer(verbalizer_file=pc.verbalizer,
                            tokenizer=tokenizer,
                            max_label_len=pc.max_label_len)

    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
         'weight_decay': pc.weight_decay},
        {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
         'weight_decay': 0.0},
    ]
    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=pc.learning_rate)
    model.to(pc.device)

    train_dataloader, dev_dataloader = get_data()
    num_update_steps_per_epoch = len(train_dataloader)
    max_train_steps = pc.epochs * num_update_steps_per_epoch
    warm_steps = int(pc.warmup_ratio * max_train_steps)
    lr_scheduler = get_scheduler(name='linear',
                                 optimizer=optimizer,
                                 num_warmup_steps=warm_steps,
                                 num_training_steps=max_train_steps)

    loss_list = []
    tic_train = time.time()
    metric = ClassEvaluator()
    criterion = torch.nn.CrossEntropyLoss()
    global_step, best_f1 = 0, 0

    for epoch in range(pc.epochs):
        for batch in tqdm(train_dataloader, desc='epoch %d' % epoch):
            logits = model(input_ids=batch['input_ids'].to(pc.device),
                           token_type_ids=batch['token_type_ids'].to(pc.device),
                           attention_mask=batch['attention_mask'].to(pc.device)).logits
            mask_labels = batch['mask_labels'].numpy().tolist()
            sub_labels = verbalizer.batch_find_sub_labels(mask_labels)
            sub_labels = [ele['token_ids'] for ele in sub_labels]

            loss = mlm_loss(logits,
                            batch['mask_positions'].to(pc.device),
                            sub_labels,
                            criterion,
                            pc.device,
                            1.0)
            loss_list.append(float(loss.cpu().detach()))
            global_step += 1

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            lr_scheduler.step()

            if global_step % pc.logging_steps == 0:
                time_diff = time.time() - tic_train
                print('global step %d, epoch: %d, loss: %.5f, speed: %.2f step/s'
                      % (global_step, epoch, sum(loss_list) / len(loss_list),
                         pc.logging_steps / time_diff))
                tic_train = time.time()

            if global_step % pc.valid_steps == 0:
                acc, precision, recall, f1, class_metrics = evaluate_model(
                    model, metric, dev_dataloader, tokenizer, verbalizer)
                print('Evaluation precision: %.5f, recall: %.5f, F1: %.5f'
                      % (precision, recall, f1))
                if f1 > best_f1:
                    print('best F1 updated: %.5f --> %.5f' % (best_f1, f1))
                    best_f1 = f1
                    cur_save_dir = os.path.join(pc.save_dir, 'model_best')
                    os.makedirs(cur_save_dir, exist_ok=True)
                    model.save_pretrained(cur_save_dir)
                    tokenizer.save_pretrained(cur_save_dir)
                tic_train = time.time()

    print('训练结束，最好 F1: %.5f' % best_f1)


if __name__ == '__main__':
    model2train()

# -*- coding:utf-8 -*-
"""P-Tuning 推理。

⚠️ 最要命的一条：p_embedding_num 必须与训练时一模一样。
训练用 6、推理用 8，不会报错，只会让 mask 位置整体后移两格，
模型在正文的字上取答案，准确率断崖式下降却查不出原因。
"""
import os
import time
from typing import List

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

from data_preprocess import convert_example
from ptune_config import ProjectConfig
from utils.common_utils import convert_logits_to_ids
from utils.verbalizer import Verbalizer

pc = ProjectConfig()
model_path = os.path.join(pc.save_dir, 'model_best')

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForMaskedLM.from_pretrained(model_path)
model.to(pc.device).eval()

verbalizer = Verbalizer(verbalizer_file=pc.verbalizer,
                        tokenizer=tokenizer,
                        max_label_len=pc.max_label_len)


def inference(contents: List[str]):
    with torch.no_grad():
        start_time = time.time()
        tokenized_output = convert_example(
            {'text': contents},
            tokenizer,
            max_seq_len=128,
            max_label_len=pc.max_label_len,
            p_embedding_num=pc.p_embedding_num,   # 必须和训练一致
            train_mode=False,
            return_tensor=True,
        )
        logits = model(input_ids=tokenized_output['input_ids'].to(pc.device),
                       token_type_ids=tokenized_output['token_type_ids'].to(pc.device),
                       attention_mask=tokenized_output['attention_mask'].to(pc.device)).logits
        predictions = convert_logits_to_ids(
            logits, tokenized_output['mask_positions']).cpu().numpy().tolist()
        predictions = verbalizer.batch_find_main_label(predictions, hard_mapping=False)
        predictions = [ele['label'] for ele in predictions]
        print('用时 %.3fs' % (time.time() - start_time))
        return predictions


if __name__ == '__main__':
    contents = [
        '天台很好看，躺在躺椅上很悠闲，适合一家出行，下次有机会肯定还会再来的',
        '这个笔记本的键盘手感一般，但屏幕素质不错，办公够用了',
        '榴莲闻着臭吃着香，这次买的很新鲜，果肉厚实',
    ]
    for content, label in zip(contents, inference(contents)):
        print('%s  ->  %s' % (label, content[:22]))

# -*- coding:utf-8 -*-
"""P-Tuning 骨架模板：换任务时只改这四处。

拿它当起点比拿 train.py 快，因为「一定要改的」已经被 TODO 标出来了。
"""
import os

# TODO-1: 底座。必须是带 [unusedN] 占位符的词表，否则换一种挪用方案
PRE_MODEL = os.environ.get('PTUNE_PRE_MODEL', './bert-base-chinese')
# TODO-2: 伪 token 个数。改了它 => 必须重新生成数据、重新训练、推理同步改
P_EMBEDDING_NUM = 6
# TODO-3: 标签最长几个字。标签有三字词就要调到 3，否则被截断
MAX_LABEL_LEN = 2
# TODO-4: verbalizer.txt 路径。主标签 \t 子标签1,子标签2,...
VERBALIZER = './data/verbalizer.txt'

MAX_SEQ_LEN = 512


def check_config(label_texts):
    """开训前的四项自检，任何一项不过就别开始跑。"""
    problems = []

    # ① 标签长度够不够
    too_long = [t for t in label_texts if len(t) > MAX_LABEL_LEN]
    if too_long:
        problems.append('这些标签会被截断: %s（把 MAX_LABEL_LEN 调到 %d）'
                        % (too_long, max(len(t) for t in label_texts)))

    # ② 截断后有没有撞车：「笔记本」和「笔记」截到 2 字都是「笔记」
    truncated = {}
    for t in label_texts:
        truncated.setdefault(t[:MAX_LABEL_LEN], []).append(t)
    clashes = {k: v for k, v in truncated.items() if len(v) > 1}
    if clashes:
        problems.append('截断后标签撞车: %s' % clashes)

    # ③ 伪 token 数量不能超过词表里的空位
    if P_EMBEDDING_NUM > 99:
        problems.append('P_EMBEDDING_NUM=%d 超过 99 个 [unused] 空位' % P_EMBEDDING_NUM)

    # ④ 正文还剩多少位置
    usable = MAX_SEQ_LEN - P_EMBEDDING_NUM - MAX_LABEL_LEN - 2
    if usable < 64:
        problems.append('正文只剩 %d 个 token，长评论会被截没' % usable)

    return problems, usable


if __name__ == '__main__':
    labels = ['水果', '书籍', '平板', '衣服', '酒店', '洗浴', '母婴', '电器', '笔记本', '蒙牛']
    problems, usable = check_config(labels)
    print('正文可用长度:', usable)
    if problems:
        print('自检未通过:')
        for p in problems:
            print('  -', p)
    else:
        print('自检通过，可以开训')