bert实现完形填空简单案例

使用 bert 来实现一个完形填空的案例，使用预训练模型 bert-base-chinese ，这个模型下载到跟代码同目录下即可，下载可参考：bert预训练模型下载-CSDN博客

通过这个案例来了解一下怎么使用预训练模型来完成下游任务，算是对怎么使用 bert 的流程有一个初步的了解，代码都写注释了，直接看代码注释就好：

import torch
from transformers import BertTokenizer, BertForMaskedLM, BertConfig


#加载预训练模型 tokenizer (vocabulary)，词表文件
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')

#输入文本
# [CLS] 表示段落开始，只有一个
# [SEP] 表示句子的结束，可以有多个
text = "[CLS] 我是谁 ? [SEP] zjk原来是程序员 [SEP]"
# 转换输入文本，有点像分词，如词在词表中找不到，会以 # 开头
tokenized_text = tokenizer.tokenize(text)
print(tokenized_text)

# 将索引为 11 的字用 [MASK] 屏蔽，如果被屏蔽的是 # 那么预测不出来结果
masked_index = 11  # 掩码一个标记，用' BertForMaskedLM '预测回来
tokenized_text[masked_index] = '[MASK]'
print(tokenized_text)

# 将标记转换为词汇表索引
indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
print(indexed_tokens)

# 将输入转换为PyTorch张量
tokens_tensor = torch.tensor([indexed_tokens])
print(tokens_tensor)

#指定设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

# 加载预训练模型 (weights)
model = BertForMaskedLM.from_pretrained('bert-base-chinese')
# 为了确保模型在处理新样本时能提供稳定且具有代表性的预测结果，禁用dropout和batch normalization等在训练阶段启用但在预测时应关闭的功能。
model.eval()
model.to(device)

# 段标识索引，标识输入文本中的第一句，第2据，0对应属于第一个句子的，1代表对应属于第二个句子的
segments_ids = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]
segments_tensors = torch.tensor([segments_ids]).to(device)

tokens_tensor = tokens_tensor.to(device)  # 输入句子的张量，其实就是每个字在词表中的索引
# 预测所有的tokens，可以理解为标准固定写法
with torch.no_grad():  # 模型不会累积梯度，适合用于预测任务
    outputs = model(tokens_tensor, token_type_ids=segments_tensors)  # 传入了tokens张量和对应的句子类型张量
# BERT模型的输出是一个元组，其中第一个元素通常是 mask 对应于每个位置上所有可能token的概率分布，
# 形状与输入tokens张量相同（这里为 [1, 16, 21128]，表示batch大小为1、序列长度为16、词汇表大小为21128的三维张量）
# 的是模型在每次迭代更新参数时处理的样本数量，这里我们就处理了一个样本
predictions = outputs[0]  # [1, 16, 21128]
# 找到在mask位置上的概率最高的token索引。masked_index 是待预测 token 在 tokens 张量中的位置
predicted_index = torch.argmax(predictions[0, masked_index]).item()
predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])[0]  # 转成单词
print('Predicted token is:', predicted_token)

运行结果：