使用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]