支付宝红包
京东盲盒抽奖
幸运转盘
秒杀
自营热卖
支付宝红包

使用transformers调用bert进行token预测

银烛秋光冷画屏 1年前   阅读数 133 0

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

注意:本文归作者所有,未经作者允许,不得转载

全部评论: 0

    我有话说: