使用transformers调用bert进行token预测

大家好，我是亓官劼（qí guān jié ），在【亓官劼】公众号、CSDN、GitHub、B站、华为开发者论坛等平台分享一些技术博文，主要包括前端开发、python后端开发、小程序开发、数据结构与算法、docker、Linux常用运维、NLP等相关技术博文，时光荏苒，未来可期，加油~

如果喜欢博主的文章可以关注博主的个人公众号【亓官劼】(qí guān jié)，里面的文章更全更新更快。如果有需要找博主的话可以在公众号后台留言，我会尽快回复消息，其他平台私信回复较慢。

huggingface开源的transfomers中目前已有100多个模型，可以方便大家进行使用测试，github仓库为：huggingface/transformers

可以直接在安装transformers包进行使用：pip install transformers

具体的方法示例如下所示，代码中有相应注释：

"""
    使用transformers进行调用BERT，进行MLM任务，对输入句子MASK的内容进行还原。
    author: 亓官劼
    date: 2022-01-14
"""
from transformers import BertForMaskedLM, BertTokenizer
import torch

# model_name = 'bert-base-chinese'  # 中文bert
model_name = 'bert-base-uncased'  # 英文bert

# MLM任务
# BERT的输入引入了[CLS]、[SEP]、[MASK]标记
# [CLS]表示分类token，该token最后一个transformer层输出被用来聚集整个序列的表征信息。
# [SEP]用来表示句子结束
# [MASK]为引入的掩盖标签，表示该处的token被遮挡，需要还原
input_text = ['How are you ? [SEP] I am [MASK] , [MASK] you !', '[MASK] are you ? [SEP] I am [MASK] , [MASK] you !']

tokenizer = BertTokenizer.from_pretrained(model_name)
# 使用tokenizer()进行tokenization的话，会自动在每句话的前面加一个[CLS]标签，以及句子结尾加一个[SEP]标签
tokenized_text = tokenizer(input_text)
print("tokenized_text：", tokenized_text)
input_ids = torch.LongTensor(tokenized_text['input_ids'])  # 输入token的id,后续输入到模型需要使用tensor形式
print("input_ids: ", input_ids)
# cache_dir可以填下载好的模型地址，也可以不填这个属性
model = BertForMaskedLM.from_pretrained(model_name, cache_dir=f"./{
      
      model_name}/")
model.eval()

outputs = model(input_ids)
outputs_token = [torch.argmax(outputs[0], dim=-1)[i] for i in range(input_text.__len__())]
print("outputs_token： ", outputs_token)
ouput_text = [tokenizer.convert_ids_to_tokens(outputs_token[i]) for i in range(input_text.__len__())]
print("ouput_text:")
for item in ouput_text:
    print(" ".join(item[1:-1]))  # 去掉tokenizer自动添加的的[CLS]和[SEP]

使用transformers调用bert进行token预测

使用transformers调用bert进行token预测

Guess you like