使用Bert预训练模型生成词嵌入

import torch
from torch import nn
from transformers import BertModel,BertConfig,BertTokenizer

# 预训练模型存储位置
pretrained_path = 'bert-base-chinese/'
config = BertConfig.from_pretrained(pretrained_path)
tokenizer = BertTokenizer.from_pretrained(pretrained_path)
model = BertModel.from_pretrained(pretrained_path,config=config)

batch = tokenizer.encode_plus('这是一个晴朗的早晨') # encode_plus 返回一个字典
input_ids = torch.tensor([batch['input_ids']])
token_type_ids = torch.tensor([batch['token_type_ids']])
attention_mask = torch.tensor([batch['attention_mask']])

embedding = model(input_ids,token_type_ids=token_type_ids)