如何使用BERT預訓練模型提取文本特征?
1 安裝第三方庫transformers
transformers使用指南
transformers安裝
pip install transformers
2 加載預訓練模型
2.1 在線加載
import torch
from transformers import BertTokenizer, BertModel
# 根據模型名稱加載
# 第一次會在線加載模型,並且保存至用戶子目錄"\.cache\torch\transformers\"
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
bert = BertModel.from_pretrained('bert-base-uncased')
2.2 本地加載
在線加載出問題時可選擇本地加載方式
2.2.1 手動下載模型
以模型Chinese-BERT-wwm為例:
- 打開上面的鏈接,找到要使用的模型並下載相應的版本,如BERT-wwm-ext, Chinese PyTorch版
- 接着解壓到某個目錄中,如D:\Program\pretrained_bert_models\chinese_wwm_ext_pytorch
- 將文件bert_config.json重命名為config.json,此時含有三個文件:config.json、pytorch_model.bin、vocab.txt
2.2.2 本地加載模型
from transformers import BertTokenizer, BertModel
# 根據模型文件目錄加載
bert_path = 'D:/Program/pretrained_bert_models/chinese_wwm_ext_pytorch'
tokenizer = BertTokenizer.from_pretrained(bert_path)
bert = BertModel.from_pretrained(bert_path, return_dict=True)
3 文本特征提取
3.1 一個句子的特征提取
inputs = tokenizer("中文", return_tensors="pt") # "pt"表示"pytorch"
outputs = bert(**inputs)
tokenizer會在文本開始和結尾分別添加[CLS]和[SEP]標記,返回給inputs的是一個字典:
key | value |
---|---|
'input_ids' | tensor([[ 101, 704, 3152, 102]]) |
'token_type_ids' | tensor([[0, 0, 0, 0]]) |
'attention_mask' | tensor([[1, 1, 1, 1]]) |
bert的輸出outputs包括last_hidden_state與pooler_output:
outputs.last_hidden_state.shape # torch.Size([1, 4, 768])
outputs.pooler_output.shape # torch.Size([1, 768])
# 若加載bert模型時return_dict=False,則輸出last_hidden_state與pooler_output組成的元組
其中last_hidden_state的形狀的含義為 (batch_size, sequence_length, hidden_size);pooler_output是last_hidden_state中token [CLS]對應的特征(即last_hidden_state[:, 0, :])經過全連接運算得到,一般可以認為其概括了整句話的信息,其形狀的含義為(batch_size, hidden_size)。
3.2 多個句子的特征提取
text_list = ["中文", "自然語言處理"]
inputs = tokenizer(text_list, return_tensors="pt", padding=True) # 長度不等時會進行填充
outputs = bert(**inputs)