在BERT中的屏蔽遮蔽(Masked Language Model,MLM)任务中,我们需要预测一些被遮蔽的词语。下面是一个使用PyTorch实现的BERT MLM任务的代码示例:
import torch
import torch.nn as nn
from transformers import BertTokenizer, BertForMaskedLM
# 加载预训练的BERT模型和分词器
model_name = 'bert-base-uncased'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForMaskedLM.from_pretrained(model_name)
# 输入句子
sentence = "I love [MASK]."
# 分词
tokens = tokenizer.tokenize(sentence)
# 添加特殊标记 [CLS] 和 [SEP]
tokens = ['[CLS]'] + tokens + ['[SEP]']
# 将词语转换为索引
input_ids = tokenizer.convert_tokens_to_ids(tokens)
# 创建遮蔽掩码
mask_idx = tokens.index('[MASK]')
attention_mask = [1] * len(input_ids)
attention_mask[mask_idx] = 0
# 将输入转换为PyTorch张量
input_ids = torch.tensor([input_ids])
attention_mask = torch.tensor([attention_mask])
# 使用BERT模型进行预测
outputs = model(input_ids, attention_mask=attention_mask)
masked_lm_logits = outputs.logits
# 获取预测结果
masked_token_logits = masked_lm_logits[0, mask_idx]
predicted_token_id = torch.argmax(masked_token_logits).item()
predicted_token = tokenizer.convert_ids_to_tokens([predicted_token_id])[0]
print("Predicted token:", predicted_token)
在上面的示例中,我们使用了Hugging Face的transformers
库来加载预训练的BERT模型和分词器。首先,我们从预训练的BERT模型中加载了bert-base-uncased
模型和相应的分词器。然后,我们提供了一个句子作为输入,对其进行分词并添加特殊标记[CLS]
和[SEP]
。接下来,我们将分词后的输入转换为BERT模型所需的输入张量,并创建了一个遮蔽掩码,将遮蔽位置的注意力掩码设置为0。最后,我们使用BERT模型对输入进行预测,并获取了被遮蔽位置的预测结果。
请注意,上述代码示例中使用的是bert-base-uncased
模型,该模型适用于小写字母的英文文本。如果你需要处理其他语言或希望使用不同的BERT模型,可以根据需要更改model_name
变量,并下载相应的模型和分词器。