HuggingFace Bert语言模型抽取句子表征向量

1. 安装相应的包

pip install transformers

2. 具体代码

from transformers import BertTokenizer, BertModel

tokenizer = BertTokenizer.from_pretrained('bert_base_cased', cache_dir="dblp/bert_base_cased", local_files_only=True)
model = BertModel.from_pretrained('bert_base_cased', cache_dir="dblp/bert_base_cased", local_files_only=True).cuda().eval()

inputs = tokenizer(setence, return_tensors='pt', max_length=512, truncation=True, padding=True)['input_ids'].cuda()
# [1, squence, hidden_dim]
last_hidden_state = model(**inputs).last_hidden_state
cls_token = last_hidden_state[:, 0, :]

3. 使用预训练权重

step1. 在huggingface上下载想用的模型权重
step2. 创建模型的时候, 手动指定权重的位置, 如果不手动指定, 则默认会先通过网络下载模型权重到本地, 然后再加载模型权重。 第2步中就是加载了指定位置的模型权重。