Python调用BERT模型的代码
时间: 2023-05-31 21:08:08 浏览: 69
以下是Python调用BERT模型的示例代码:
```python
import torch
from transformers import BertTokenizer, BertModel
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertModel.from_pretrained('bert-base-chinese', output_hidden_states=True)
text = "我是一名语言模型。"
encoded_input = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
output = model(**encoded_input)
last_hidden_states = output.last_hidden_state # 最后一层的hidden states
pooler_output = output.pooler_output # 经过全局平均池化之后的hidden states(用于任务分类)
hidden_states = output.hidden_states # 所有层的hidden states,第一层为input embeddings
```
此代码使用BERT-base中文模型(bert-base-chinese)对一句中文文本进行编码,得到最后一层的hidden states、经过全局平均池化之后的hidden states、所有层的hidden states。在使用中,可以根据具体需求选择使用哪种输出。