【完形填空实战】BERT中文MLM微调：数据构造→训练→评估，有完整可运行代码

1、任务介绍

# 输入一句话，MASK一个字，训练模型进行填空
[CLS] 选 择 珠 江 花 园 的 原 因 就 是 方 便 ， 有 [MASK] 动 扶 梯 直 接 到 达 海 边 ， 周 围 餐 馆 [SEP]
# 本句MASK的字为“电”

2、数据介绍

• 数据文件有三个train.csv，test.csv，validation.csv，数据样式都是一样的。

label,text
1,选择珠江花园的原因就是方便，有电动扶梯直接到达海边，周围餐馆、食廊、商场、超市、摊位一应俱全。酒店装修一般，但还算整洁。 泳池在大堂的屋顶，因此很小，不过女儿倒是喜欢。 包的早餐是西式的，还算丰富。 服务吗，一般
1,15.4寸笔记本的键盘确实爽，基本跟台式机差不多了，蛮喜欢数字小键盘，输数字特方便，样子也很美观，做工也相当不错
0,房间太小。其他的都一般。。。。。。。。。
0,"1.接电源没有几分钟,电源适配器热的不行. 2.摄像头用不起来. 3.机盖的钢琴漆，手不能摸，一摸一个印. 4.硬盘分区不好办."
1,"今天才知道这书还有第6卷,真有点郁闷:为什么同一套书有两种版本呢?当当网是不是该跟出版社商量商量,单独出个第6卷,让我们的孩子不会有所遗憾。"

• 导入工具包和辅助工具实例化对象

import torch
from torch.utils.data import DataLoader
from datasets import load_dataset
from transformers import BertTokenizer, BertModel
from transformers import AdamW
import time

# 加载字典和分词工具
my_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

# 加载预训练模型
my_model_pretrained = BertModel.from_pretrained('bert-base-chinese')

3、数据预处理

3.1 完形填空任务的真实值

完形填空（MLM）任务的真实值：藏在输入里的标签

1. 核心认知：标签就在句子内部

与文本分类任务（需要独立的 label 字段）不同，完形填空任务不需要额外的人工标注标签。它的真实值（Ground Truth）就是被挖掉的那个词原本的 ID。

你可以把训练数据想象成自带答案的填空题试卷：

• 分类任务数据：{'text': '...', 'label': 1} → 答案在外面。
• MLM 任务数据：{'text': '今天天气[MASK]不错'} → 答案就是 [MASK] 对应的词（例如“真”）。

2. 真实值是如何生成的？（数据处理流程）

模型看到的输入句子里带有 [MASK]，但计算 Loss 时，我们需要告诉模型正确答案是什么。这个过程通常在 collate_fn 或 Dataset 中完成。

构造逻辑图解：

关键规则：

1. labels 形状与 input_ids 完全一致。
2. 只在被遮盖的位置填入原始词 ID。
3. 其余位置全部填入 -100（PyTorch 的 CrossEntropyLoss 默认 ignore_index=-100，即忽略该位置的损失计算）。

【补充说明】：-100 是 PyTorch 损失函数的特殊值，你也可以通过 ignore_index 参数自定义，但 -100 是事实标准。

3. 代码实现对比（核心区别一目了然）

文本分类任务的 collate_fn（例子）：

def collate_fn_cls(batch):
    # batch 里的 label 是外部给定的，直接取出来堆叠即可
    texts = [item['text'] for item in batch]
    labels = [item['label'] for item in batch]  # 外部标签
    
    inputs = tokenizer(texts, padding=True, return_tensors='pt')
    return inputs, torch.tensor(labels)

完形填空任务的 collate_fn（真实值构造逻辑）： （仅为示例）

def collate_fn_mlm(batch):
    # batch 里通常只有纯文本 'text'，没有 'label'
    texts = [item['text'] for item in batch]
    
    # 1. 先正常编码（此时还没有 mask）
    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
    input_ids = inputs['input_ids']
    attention_mask = inputs['attention_mask']  # 【补充】需要获取 mask
    
    # 2. 创建初始标签，默认全为 -100 (忽略计算)
    labels = torch.full_like(input_ids, -100)
    
    # 3. 生成随机掩码矩阵 (15% 的概率被选中)
    probability_matrix = torch.full(input_ids.shape, 0.15)
    
    # 【修正】必须排除特殊 token：不能遮盖 [CLS], [SEP], [PAD]
    special_tokens_mask = tokenizer.get_special_tokens_mask(
        input_ids.tolist(), already_has_special_tokens=True
    )
    special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
    probability_matrix.masked_fill_(special_tokens_mask, value=0.0)  # 特殊 token 遮盖概率置 0
    
    masked_indices = torch.bernoulli(probability_matrix).bool()
    
    # 4. 【核心】：在 labels 中记录被遮盖位置的原始 ID
    labels[masked_indices] = input_ids[masked_indices]
    
    # 5. 修改 input_ids，实现 80-10-10 策略（BERT 论文标准）
    #    80% 替换为 [MASK]，10% 替换为随机词，10% 保持不变
    indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
    input_ids[indices_replaced] = tokenizer.mask_token_id
    
    indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
    random_words = torch.randint(tokenizer.vocab_size, labels.shape, dtype=torch.long)
    input_ids[indices_random] = random_words[indices_random]
    
    # 其余 10% 保持原词不变（input_ids 中已保留原值）
    
    return {
        'input_ids': input_ids,
        'attention_mask': attention_mask,
        'labels': labels
    }

【重要说明】：上述代码中的 80-10-10 策略是 BERT 论文的标准做法。前半部分简化为“全替换为 [MASK]”仅用于帮助理解标签构造的核心思想，实际训练时应当采用 80-10-10 策略，否则模型在推理时会遇到从未见过的 [MASK] 导致性能下降。

4. 为什么 Loss 计算时只关注被遮盖的词？

如果不把非遮盖位置设为 -100，模型就需要预测句子里的每一个词（包括“的”、“了”这种没有信息量的词），这会稀释掉模型对语义填空能力的学习。

通过将 labels 中 85% 的位置设为 -100：

• Loss 仅来自于被遮盖位置的预测偏差（注意：包括被随机替换或保持原样的词，只要在 masked_indices 中，它们的 labels 仍记录原词 ID）。
• 模型被迫只关注“根据上下文推理缺失词”的任务，这正是 BERT 理解能力的来源。

5. 一句话总结

3.2 张量视图 与 torch.clone

PyTorch 张量索引的引用机制：a[0] 不是值，是“窗口”

1. 核心结论

在 PyTorch 中，通过基本索引（单个整数、切片 :）获取的张量元素，返回的是原张量的视图（View）或引用，而不是独立的数据副本。

• 修改原张量 → 索引出的张量值会同步变化。
• 修改索引出的张量 → 原张量的对应位置也会同步变化。

这与 Python 列表 list 的行为截然不同（list[0] 返回的是独立的值拷贝）。

【补充】：a[0] 返回的是一个0维张量（标量张量），而非 Python 原生数值。它依然是一个拥有元数据的张量对象，其内部存储指向原张量的内存区域。

2. 现象复现（你已验证）

import torch

a = torch.tensor([1, 2, 3])
b = a[0]               # b 是一个 0 维张量（标量张量），但它引用 a 的内存

print(f"a: {a}")       # tensor([1, 2, 3])
print(f"b: {b}")       # tensor(1)

a[0] = 666             # 修改原张量
print(f"b after modifying a: {b}")  # tensor(666) —— b 跟着变了！

反向验证（修改 b 会影响 a）：

a = torch.tensor([1, 2, 3])
b = a[0]
b.fill_(999)           # 通过 b 修改值
print(a)               # tensor([999, 2, 3]) —— a 也被修改了！

2.5 重要辨析：变量名重新绑定 ≠ 视图修改

你可能会遇到以下看似矛盾的代码：

import torch

a = torch.tensor(1)
b = a                 # b 和 a 指向同一个张量对象

a = torch.tensor(666) # 将变量名 a 重新绑定到一个新张量
print(b)              # tensor(1) —— 为什么 b 没变？

原因解释：

类比理解：

• 修改内部状态（如 a[0]=...、b.fill_()）就像两个人用同一把钥匙开同一个保险箱，一人往里面放东西，另一人打开看到的就是变化后的内容。
• 变量重新赋值（a = ...）就像其中一个人把钥匙扔了，去配了一把新钥匙开新箱子，另一个人手里的旧钥匙依然能打开旧箱子。

结论：

a = torch.tensor(666) 改变的是“变量名 a 指谁”，而不是“原张量对象的内容”。而之前讨论的张量索引视图（如 a[0]）修改的是“张量对象内部的数值”，这是两个完全不同的概念。

3. 为什么会这样？

PyTorch 的张量在底层由两部分组成：

• 元数据：形状（shape）、步长（stride）、数据类型等。
• 存储区：一块连续的物理内存，存放实际数值。

基本索引操作不会复制数据，它只是创建了一个新的张量对象，但这个新对象的存储区指针仍然指向原张量的同一块内存。

a = torch.tensor([1, 2, 3])
b = a[0]

内存视角：
a 的存储区: [1, 2, 3]
b 的存储区: 指向 a 存储区的第一个元素

【补充】如何验证两个张量是否共享内存？
使用 .data_ptr() 查看张量底层存储的首地址：

a = torch.tensor([1, 2, 3])
b = a[0]
print(a.data_ptr())  # 例如 140332160000000
print(b.data_ptr())  # 相同地址，证明共享存储

c = a.clone()
print(c.data_ptr())  # 不同地址，独立存储

4. 哪些操作返回“视图/引用”？哪些返回“副本”？

判断口诀：用整数、切片、... 取出来的通常是视图（View）；用列表、张量、布尔条件取出来的是副本（Copy）。

5. 如何强制获取独立副本？

如果你需要断开与原张量的连接，使用 .clone() 方法：

a = torch.tensor([1, 2, 3])
b = a[0].clone()    # 深拷贝，独立存储

a[0] = 666
print(b)            # tensor(1) —— 不再受 a 的影响

6. 对比：Python 列表 vs PyTorch 张量

7. 一句话总结

PyTorch 的基本索引返回的是“窗口”而非“快照”。通过 a[0] 拿到的标量张量，依然是原张量在内存中的直接映射，修改一方必然影响另一方。但请注意区分：a = new_tensor 只是让变量名指向新对象，不会影响之前从它引出的其他变量。若要独立副本，请使用 .clone()。

3.3 [MASK] 掉 token 后，不用改 attention_mask

inputs = my_tokenizer(
    text=samples,
    padding=True,
    truncation=True,
    max_length=32,           # 设定的 [MASK] 为第 16 个 token, 只要设定的[MASK]位置比样本长度小即可
    return_tensors='pt'
)

把 inputs['input_ids'] 里指定的 token 替换为 mask_token_id 后，不需要把 inputs['attention_mask'] 里对应的 [MASk]

📝 为什么 Mask 掉 Token 后，不用改 Attention Mask？

简单直接的回答是：不需要改变 attention_mask。

在完形填空（MLM）任务代码中，attention_mask 的职责和 [MASK] 标记的职责是完全分开的。为了让你彻底明白，我们需要区分这两种“遮盖”。

1. Attention Mask 的真实职责：防“水”不防“词”

attention_mask 的设计初衷非常单一：它只负责屏蔽填充符（Padding），防止模型关注无意义的空白区域。

• 它的逻辑：
  • 1 = 这里是有效内容。无论是真实的字，还是被 [MASK] 遮盖的字，对模型来说都是“有东西”的位置，都需要参与计算。
  • 0 = 这里是凑数的（Padding）。模型请直接忽略，不要分配任何注意力权重。

2. 为什么 [MASK] 位置必须保留为 1？

在 MLM 任务中，我们将第 16 个词替换为 [MASK] ID，目的是让模型去预测这个位置原本是什么。

• 形象类比：
  • Input IDs 是试卷。[MASK] 就像是试卷上的填空题横线 ____。
  • Attention Mask 是告诉学生哪里需要答题。
  • 如果你把 [MASK] 位置的 mask 改为 0，就等于告诉模型：“这道题不用做，跳过”。那模型自然就不会去预测这个位置了。
• 结论：模型需要“看到”这个 [MASK] 的位置（即 attention_mask 为 1），并结合上下文来计算这个位置应该填什么。

3. 图解对比：两种 Mask 的区别

这是新手最容易混淆的地方，我们需要区分 “为了让模型预测而遮盖” 和 “为了让模型忽略而遮盖”。

4. 结合代码分析

让我们看看代码执行后的内存状态（假设 batch size=1, max_length=32，且句子较短，后面有大量 Padding）：

# 原始句子: "今天天气真好" (假设只有前几个词，后面补了 PAD)

# 1. inputs['attention_mask'] 的状态 (保持不变)
# [1,   1,   1,    1,   ...  1,   1,  1,   0,     0,      0] 
#  ^^^^^^^^^^真实内容(含Mask位)^^^^^^^^^^   ^^^Padding部分^^^
# 注意：即使第16位变成了[MASK]，它在 attention_mask 里依然是 1，代表它是有效位置

# 2. input_ids 的变化过程
# 初始: [101, 791, ...,   567(真)   , ..., 0, 0, 0]  <-- 0 是 PAD 的 ID
# 修改: [101, 791, ...,  103(MASK)  , ..., 0, 0, 0] 
#                           ^ 这里变成了 103，但它依然是句子的有效组成部分

5. 总结

在 collate_fn 中：

1. labels: 记录第 16 位原本是谁（比如 ID 567），这是标准答案。
2. input_ids: 把第 16 位改成 [MASK]（ID 103），这是出题。
3. attention_mask: 保持不变。它只需要诚实地反映“哪些位置是真实文本（含考题），哪些是 Padding”。对于模型来说，[MASK] 就是一个正常的、有意义的 Token，理应被分配注意力权重。

所以，代码逻辑是完全正确的，不需要额外修改 attention_mask。

3.4 最终代码

def collate_fn_(batch):
    # print(batch)
    # [
    #     {
    #         'label': 1,
    #         'text': '酒店周边环境比较乱，哈尔滨......会入住此酒店。'
    #     },
    #     {
    #         'label': 0,
    #         'text': '非常不满这酒店，配不上5星。第一,客......铁丝…睡得很不舒服'
    #     },
    #     {
    #         'label': 1,
    #         'text': '酒店的服务不错，我去是预定的......题。欢迎您的下次光临！'
    #     }
    # ]

    # 完形填空任务 的 真实值在句子内部, 不需要文本分类的 label 这种真实值
    samples = [label_text['text'] for label_text in batch]

    inputs = my_tokenizer(
        text=samples,
        padding=True,
        truncation=True,
        max_length=32,        # 设定的 [MASK] 为第 mask_index 个 token, 只要设定的[MASK]位置比样本长度小即可
        return_tensors='pt'
    )

    # 指定每个样本的第 mask_index 个token为 [MASK], 获取 [MASK] 对应的真实值
    input_ids = inputs['input_ids']
    # print(input_ids.shape)    # torch.Size([8, 32])

    labels = input_ids[: , mask_index].clone()      # 为什么要clone，在笔记《张量视图 与 `torch.clone`》有详情
    labels = labels.to(dtype=torch.long)
    # print(labels.shape)     # torch.Size([8])

    # 把 input_ids 里被指定为 [MASK] 的 token 变成 mask_token_id
    # 基本索引（单个整数、切片 `:`）获取的张量元素，返回的是原张量的视图（View）或引用，而不是独立的数据副本。
    # 在笔记《张量视图 与 `torch.clone`》有详情
    input_ids[: , mask_index] = my_tokenizer.mask_token_id     # 会自动广播, inputs['input_ids']会被修改

    # 不需要把 inputs['attention_mask'] 第 mask_index个 进行 [MASK], 详情笔记《[MASK] 掉 `token` 后，不用改 `attention_mask`》有详情

    return inputs, labels


# 测试 Dataset、DataLoader
def taste_dataloader():
    my_dataset = load_dataset(
        path='csv',
        data_files={'train': './data/train.csv'},
        split='train'
    )
    print(my_dataset)
    # Dataset({
    #     features: ['label', 'text'],
    #     num_rows: 9600
    # })

    # 只保留文本长度大于 32 的样本（随便多长就行, 案例而已）
    # 下面匿名函数里的 x 相当于 my_dataset[0] 取到的结果, 是个字典
    my_dataset = my_dataset.filter(lambda sample: len(sample['text']) > 32)
    print(my_dataset)
    # Dataset({
    #     features: ['label', 'text'],
    #     num_rows: 9035
    # })

    my_dataloader = DataLoader(
        dataset=my_dataset,
        batch_size=8,
        shuffle=True,
        collate_fn=collate_fn_
    )

    for inputs, labels in my_dataloader:
        print(inputs)
        print(labels)
        break

4、自定义下游任务网络模型

4.1 加载预训练模型 & 分词器

device = 'cuda' if torch.cuda.is_available() else 'cpu'
mask_index = 10    # [MASK] 掉 inputs['input_ids'] 第10个token

# 加载 bert 分词器
my_tokenizer = BertTokenizer.from_pretrained(r'./model/bert-base-chinese')

# 加载模型
bert_model = BertModel.from_pretrained(r'./model/bert-base-chinese')
bert_model = bert_model.to(device=device)

4.2 创建下游任务模型

# 自定义模型
class MyMask(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(
            in_features=bert_model.config.hidden_size,
            out_features=my_tokenizer.vocab_size
        )

    def forward(self, inputs):
        # '预训练模型' 不需要更新参数, 所以不需要计算梯度
        with torch.no_grad():
            output = bert_model(**inputs)

        # print(output)
        # 比如:
        # BaseModelOutputWithPoolingAndCrossAttentions(
        #     last_hidden_state=tensor([[[ 0.1772, -0.0687, -1.0350,  ...,  0.0109, -0.0692, -0.6392],
        #                                ...,
        #                                [ 0.3647, -0.0271, -1.0330,  ...,  0.3621,  0.3096, -0.8126]],
        #                    `           ...
        #                               [[-0.2139,  0.5916, -1.2234,  ...,  0.4955, -0.1464, -0.4860],
        #                                ...,
        #                                [ 0.6937, -0.1311, -0.3905,  ...,  0.1585, -0.1104, -0.4002]]]),
        #     pooler_output=tensor([[ 0.9990,  1.0000,  0.9704,  ..., -0.9735, -0.9996,  0.7838],
        #                           ...
        #                           [ 0.9998,  0.9994,  0.9999,  ..., -0.9993, -0.9936,  0.1820]]),
        #     hidden_states=None,
        #     past_key_values=None,
        #     attentions=None,
        #     cross_attentions=None
        # )

        # 只需要拿到 被[MASK] 掉的 token 的 logits 即可
        # 这里要使用 last_hidden_state, pooler_output是[CLS]的表示
        mask_logits = output.last_hidden_state[ : ,mask_index]
        # print(f'mask_logits = {mask_logits.shape}')    # torch.Size([8, 768])

        logits = self.linear(mask_logits)
        # print(f'logits = {logits.shape}')     # torch.Size([8, 21128])

        return logits

4.3 测试模型

# 测试模型
def taste_model():
    my_dataset = load_dataset(
        path='csv',
        data_files={'train': './data/train.csv'},
        split='train'
    )
    my_dataloader = DataLoader(
        dataset=my_dataset,
        batch_size=8,
        shuffle=True,
        collate_fn=collate_fn_
    )

    my_model = MyMask().to(device=device)

    for inputs, labels in my_dataloader:
        # inputs 是 my_tokenizer 得到的结果, 是 BatchEncoding类, 可以直接使用 .to(device)
        # 而不是 load_dataset() 返回的 DatasetDict 或 Dataset 对象本身调用 .to()。
        # print(type(inputs))     # BatchEncoding

        inputs = inputs.to(device=device)
        labels = labels.to(device=device)

        logits = my_model(inputs)
        break

5、模型训练

# 训练模型
def train_model():
    my_dataset = load_dataset(
        path='csv',
        data_files={'train': './data/train.csv'},
        split='train'
    ).filter(lambda sample: len(sample['text']) > 32)
    my_dataloader = DataLoader(
        dataset=my_dataset,
        batch_size=8,
        shuffle=True,
        collate_fn=collate_fn_
    )

    # 在 MyMask模型的 forward(() 里已经使用了 with torch.no_grad(), 这里再加一重保险
    # 冻结掉 bert_model 的参数
    for param in bert_model.parameters():
        param.requires_grad_(False)  # 带下划线的是原地修改

    my_model = MyMask().to(device=device)
    my_model.train()

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(params=my_model.parameters())

    epochs = 3
    for epoch in range(epochs):
        for i, (inputs, labels) in enumerate(tqdm(my_dataloader)):
            inputs = inputs.to(device=device)
            labels = labels.to(device=device)

            logits = my_model(inputs)
            loss = criterion(logits, labels)
            optimizer.zero_grad()       # 梯度清零
            loss.backward()             # 计算梯度
            optimizer.step()            # 更新参数

            if i % 30 == 0:
                # 取每个 batch 的第0个样本看看被 [MASK] 的原句, 及其 label
                # print(f'inputs.input_ids[0] = {inputs.input_ids[0]}')
                # tensor([ 101,  679, 6389, 3221,  868,  711, 2207, 6432, 8024, 6820,  103, 5466,
                #         1767, 3136, 4906,  741, 8024, 1091, 4638, 2523, 5125, 2506, 8039,  679,
                #         6389, 3221,  818, 5466, 1912,  821,  510,  102], device='cuda:0')

                input_ids_0 = inputs.input_ids[0]
                sentence_0 = my_tokenizer.decode(
                    token_ids=input_ids_0,
                    skip_special_tokens=False   # 要在句子中显示[MASK], 所以绝对不能跳过特殊 token
                )
                mask_token_label = my_tokenizer.decode(token_ids=labels[0])     # [MASK} 真实值

                right_id = torch.argmax(logits, dim=-1)
                mask_token_predict = my_tokenizer.decode(token_ids=right_id[0])    # [MASK] 预测值

                right_count = (right_id == labels).sum().item()
                print(f'\nloss = {loss}, 正确率 = {right_count / labels.size(-1)}, 真实[MASK]: {mask_token_label}   预测[MASK]: {mask_token_predict}   原句子: {sentence_0}')
                # loss = 0.6091822385787964, 正确率 = 0.875, 真实[MASK]: 较   预测[MASK]: 较   原句子: [CLS] 酒 店 比 价 破 ， 房 间 比 [MASK] 老 ， 地 毯 上 还 有 以 前 客 人 留 下 的 瓜 子 壳 ， 服 务 [SEP]

    # 保存模型
    torch.save(my_model.state_dict(), r'./model/my_fill_mask_model.pth')

6、模型评估

# 模型评估
def eval_model():
    with torch.no_grad():
        my_dataset = load_dataset(
            path='csv',
            data_files={'test': './data/test.csv'},
            split='test'
        ).filter(lambda sample: len(sample['text']) > 32)
        my_dataloader = DataLoader(
            dataset=my_dataset,
            batch_size=8,
            shuffle=True,
            collate_fn=collate_fn_
        )
        my_model = MyMask()
        my_model.load_state_dict(torch.load(r'./model/my_fill_mask_model.pth', map_location=device))
        my_model = my_model.to(device=device)
        my_model.eval()

        right_count = 0     # 正确个数
        total = 0           # 总个数
        for i, (inputs, labels) in enumerate(my_dataloader):
            inputs = inputs.to(device=device)
            labels = labels.to(device=device)

            logits = my_model(inputs)
            right_id = torch.argmax(logits, dim=-1)
            right_count += (right_id == labels).sum().item()
            total += labels.size(0)

            if i % 30 == 0:
                print(f'正确率 = {right_count / total}, 正确个数 : 总个数 = {right_count} : {total}')

        print(f'总正确率 = {right_count / total}, 正确个数 : 总个数 = {right_count} : {total}')
        # 总正确率 = 0.7173144876325088, 正确个数 : 总个数 = 812 : 1132

7、整体代码

#coding: utf-8
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from transformers import BertModel, BertTokenizer
from datasets import load_dataset
import torch.optim as optim
from tqdm import tqdm


device = 'cuda' if torch.cuda.is_available() else 'cpu'
mask_index = 10    # [MASK] 掉 inputs['input_ids'] 第10个token

# 加载 bert 分词器
my_tokenizer = BertTokenizer.from_pretrained(r'./model/bert-base-chinese')

# 加载模型
bert_model = BertModel.from_pretrained(r'./model/bert-base-chinese')
bert_model = bert_model.to(device=device)


def collate_fn_(batch):
    # print(batch)
    # [
    #     {
    #         'label': 1,
    #         'text': '酒店周边环境比较乱，哈尔滨......会入住此酒店。'
    #     },
    #     {
    #         'label': 0,
    #         'text': '非常不满这酒店，配不上5星。第一,客......铁丝…睡得很不舒服'
    #     },
    #     {
    #         'label': 1,
    #         'text': '酒店的服务不错，我去是预定的......题。欢迎您的下次光临！'
    #     }
    # ]

    # 完形填空任务 的 真实值在句子内部, 不需要文本分类的 label 这种真实值
    samples = [label_text['text'] for label_text in batch]

    inputs = my_tokenizer(
        text=samples,
        padding=True,
        truncation=True,
        max_length=32,        # 设定的 [MASK] 为第 mask_index 个 token, 只要设定的[MASK]位置比样本长度小即可
        return_tensors='pt'
    )

    # 指定每个样本的第 mask_index 个token为 [MASK], 获取 [MASK] 对应的真实值
    input_ids = inputs['input_ids']
    # print(input_ids.shape)    # torch.Size([8, 32])

    labels = input_ids[: , mask_index].clone()      # 为什么要clone，在笔记《张量视图 与 `torch.clone`》有详情
    labels = labels.to(dtype=torch.long)
    # print(labels.shape)     # torch.Size([8])

    # 把 input_ids 里被指定为 [MASK] 的 token 变成 mask_token_id
    # 基本索引（单个整数、切片 `:`）获取的张量元素，返回的是原张量的视图（View）或引用，而不是独立的数据副本。
    # 在笔记《张量视图 与 `torch.clone`》有详情
    input_ids[: , mask_index] = my_tokenizer.mask_token_id     # 会自动广播, inputs['input_ids']会被修改

    # 不需要把 inputs['attention_mask'] 第 mask_index个 进行 [MASK], 详情笔记《[MASK] 掉 `token` 后，不用改 `attention_mask`》有详情

    return inputs, labels


# 测试 Dataset、DataLoader
def taste_dataloader():
    my_dataset = load_dataset(
        path='csv',
        data_files={'train': './data/train.csv'},
        split='train'
    )
    print(my_dataset)
    # Dataset({
    #     features: ['label', 'text'],
    #     num_rows: 9600
    # })

    # 只保留文本长度大于 32 的样本（随便多长就行, 案例而已）
    # 下面匿名函数里的 x 相当于 my_dataset[0] 取到的结果, 是个字典
    my_dataset = my_dataset.filter(lambda sample: len(sample['text']) > 32)
    print(my_dataset)
    # Dataset({
    #     features: ['label', 'text'],
    #     num_rows: 9035
    # })

    my_dataloader = DataLoader(
        dataset=my_dataset,
        batch_size=8,
        shuffle=True,
        collate_fn=collate_fn_
    )

    for inputs, labels in my_dataloader:
        print(inputs)
        print(labels)
        break


# 自定义模型
class MyMask(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(
            in_features=bert_model.config.hidden_size,
            out_features=my_tokenizer.vocab_size
        )

    def forward(self, inputs):
        # '预训练模型' 不需要更新参数, 所以不需要计算梯度
        with torch.no_grad():
            output = bert_model(**inputs)

        # print(output)
        # 比如:
        # BaseModelOutputWithPoolingAndCrossAttentions(
        #     last_hidden_state=tensor([[[ 0.1772, -0.0687, -1.0350,  ...,  0.0109, -0.0692, -0.6392],
        #                                ...,
        #                                [ 0.3647, -0.0271, -1.0330,  ...,  0.3621,  0.3096, -0.8126]],
        #                    `           ...
        #                               [[-0.2139,  0.5916, -1.2234,  ...,  0.4955, -0.1464, -0.4860],
        #                                ...,
        #                                [ 0.6937, -0.1311, -0.3905,  ...,  0.1585, -0.1104, -0.4002]]]),
        #     pooler_output=tensor([[ 0.9990,  1.0000,  0.9704,  ..., -0.9735, -0.9996,  0.7838],
        #                           ...
        #                           [ 0.9998,  0.9994,  0.9999,  ..., -0.9993, -0.9936,  0.1820]]),
        #     hidden_states=None,
        #     past_key_values=None,
        #     attentions=None,
        #     cross_attentions=None
        # )

        # 只需要拿到 被[MASK] 掉的 token 的 logits 即可
        # 这里要使用 last_hidden_state, pooler_output是[CLS]的表示
        mask_logits = output.last_hidden_state[ : ,mask_index]
        # print(f'mask_logits = {mask_logits.shape}')    # torch.Size([8, 768])

        logits = self.linear(mask_logits)
        # print(f'logits = {logits.shape}')     # torch.Size([8, 21128])

        return logits


# 测试模型
def taste_model():
    my_dataset = load_dataset(
        path='csv',
        data_files={'train': './data/train.csv'},
        split='train'
    )
    my_dataloader = DataLoader(
        dataset=my_dataset,
        batch_size=8,
        shuffle=True,
        collate_fn=collate_fn_
    )

    my_model = MyMask().to(device=device)

    for inputs, labels in my_dataloader:
        # inputs 是 my_tokenizer 得到的结果, 是 BatchEncoding类, 可以直接使用 .to(device)
        # 而不是 load_dataset() 返回的 DatasetDict 或 Dataset 对象本身调用 .to()。
        # print(type(inputs))     # BatchEncoding

        inputs = inputs.to(device=device)
        labels = labels.to(device=device)

        logits = my_model(inputs)
        break


# 训练模型
def train_model():
    my_dataset = load_dataset(
        path='csv',
        data_files={'train': './data/train.csv'},
        split='train'
    ).filter(lambda sample: len(sample['text']) > 32)
    my_dataloader = DataLoader(
        dataset=my_dataset,
        batch_size=8,
        shuffle=True,
        collate_fn=collate_fn_
    )

    # 在 MyMask模型的 forward(() 里已经使用了 with torch.no_grad(), 这里再加一重保险
    # 冻结掉 bert_model 的参数
    for param in bert_model.parameters():
        param.requires_grad_(False)  # 带下划线的是原地修改

    my_model = MyMask().to(device=device)
    my_model.train()

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(params=my_model.parameters())

    epochs = 3
    for epoch in range(epochs):
        for i, (inputs, labels) in enumerate(tqdm(my_dataloader)):
            inputs = inputs.to(device=device)
            labels = labels.to(device=device)

            logits = my_model(inputs)
            loss = criterion(logits, labels)
            optimizer.zero_grad()       # 梯度清零
            loss.backward()             # 计算梯度
            optimizer.step()            # 更新参数

            if i % 30 == 0:
                # 取每个 batch 的第0个样本看看被 [MASK] 的原句, 及其 label
                # print(f'inputs.input_ids[0] = {inputs.input_ids[0]}')
                # tensor([ 101,  679, 6389, 3221,  868,  711, 2207, 6432, 8024, 6820,  103, 5466,
                #         1767, 3136, 4906,  741, 8024, 1091, 4638, 2523, 5125, 2506, 8039,  679,
                #         6389, 3221,  818, 5466, 1912,  821,  510,  102], device='cuda:0')

                input_ids_0 = inputs.input_ids[0]
                sentence_0 = my_tokenizer.decode(
                    token_ids=input_ids_0,
                    skip_special_tokens=False   # 要在句子中显示[MASK], 所以绝对不能跳过特殊 token
                )
                mask_token_label = my_tokenizer.decode(token_ids=labels[0])     # [MASK} 真实值

                right_id = torch.argmax(logits, dim=-1)
                mask_token_predict = my_tokenizer.decode(token_ids=right_id[0])    # [MASK] 预测值

                right_count = (right_id == labels).sum().item()
                print(f'\nloss = {loss}, 正确率 = {right_count / labels.size(-1)}, 真实[MASK]: {mask_token_label}   预测[MASK]: {mask_token_predict}   原句子: {sentence_0}')
                # loss = 0.6091822385787964, 正确率 = 0.875, 真实[MASK]: 较   预测[MASK]: 较   原句子: [CLS] 酒 店 比 价 破 ， 房 间 比 [MASK] 老 ， 地 毯 上 还 有 以 前 客 人 留 下 的 瓜 子 壳 ， 服 务 [SEP]

    # 保存模型
    torch.save(my_model.state_dict(), r'./model/my_fill_mask_model.pth')


# 模型评估
def eval_model():
    with torch.no_grad():
        my_dataset = load_dataset(
            path='csv',
            data_files={'test': './data/test.csv'},
            split='test'
        ).filter(lambda sample: len(sample['text']) > 32)
        my_dataloader = DataLoader(
            dataset=my_dataset,
            batch_size=8,
            shuffle=True,
            collate_fn=collate_fn_
        )
        my_model = MyMask()
        my_model.load_state_dict(torch.load(r'./model/my_fill_mask_model.pth', map_location=device))
        my_model = my_model.to(device=device)
        my_model.eval()

        right_count = 0     # 正确个数
        total = 0           # 总个数
        for i, (inputs, labels) in enumerate(my_dataloader):
            inputs = inputs.to(device=device)
            labels = labels.to(device=device)

            logits = my_model(inputs)
            right_id = torch.argmax(logits, dim=-1)
            right_count += (right_id == labels).sum().item()
            total += labels.size(0)

            if i % 30 == 0:
                print(f'正确率 = {right_count / total}, 正确个数 : 总个数 = {right_count} : {total}')

        print(f'总正确率 = {right_count / total}, 正确个数 : 总个数 = {right_count} : {total}')
        # 总正确率 = 0.7173144876325088, 正确个数 : 总个数 = 812 : 1132


if __name__ == '__main__':
    # taste_dataloader()
    # taste_model()
    # train_model()
    eval_model()