java使用bert训练好的onnx模型 bert模型输入

背景

通常我们使用bert进行模型fine-tune时，大多是使用Transformer这个包，官方地址：https://huggingface.co/. 如果想使用Bert获取任务对应的动态词向量，那么这时我们就需要了解Bert模型需要的输入与输出了。如果对bert模型不了解的可以看看我之前的文章：【NLP】BERT(BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding)阅读与总结.

BertModel对象构建

想要使用Bert模型获取对应内容的动态词向量的第一步则是先构建Bert模型，现如今Bert模型也多种多样，在不同的任务中选择合适的预训练模型。对应的预训练模型可以在官网上离线下载，也可以在线下载。当然，毕竟模型那么大，我们通常使用的离线下载的模型。以bert-base-chinese为例。我们首先进入模型列表：

然后根据下面操作就可以进入bert-base-chinese预训练模型的地址了：

然后选择下面三个文件下载：

注意文件的名称要保持不变。下载的文件顾名思义就是对应模型文件，配置文件，词表文件。

那么模型的加载如下：

from transformers import BertModel

bert = BertModel.from_pretrained("pre_model/bert-base-chinese")
print(bert)

模型如下：

模型的输入

我们可以从官方文档中：https://huggingface.co/docs/transformers/model_doc/bert#transformers.BertModel 可以看到模型的输入：

( input_ids: typing.Optional[torch.Tensor] = None,
attention_mask: typing.Optional[torch.Tensor] = None,
token_type_ids: typing.Optional[torch.Tensor] = None,
position_ids: typing.Optional[torch.Tensor] = None,
head_mask: typing.Optional[torch.Tensor] = None,
inputs_embeds: typing.Optional[torch.Tensor] = None,
labels: typing.Optional[torch.Tensor] = None,
next_sentence_label: typing.Optional[torch.Tensor] = None,
output_attentions: typing.Optional[bool] = None,
output_hidden_states: typing.Optional[bool] = None,
return_dict: typing.Optional[bool] = None )

下面介绍一个核心的参数，更多详细内容可以查看官方文档：

• input_ids (torch.LongTensor of shape (batch_size, sequence_length)) — 词汇表中输入序列标记的索引
• attention_mask (torch.FloatTensor of shape (batch_size, sequence_length), optional) — 对输入数据进行mask，解决pad问题.在 [0, 1] 中选择的掩码值:1 表示未屏蔽的标记，0 表示已屏蔽的标记
• token_type_ids (torch.LongTensor of shape (batch_size, sequence_length), optional) — 分段标记索引以指示输入的第一和第二部分。在 [0, 1] 中选择索引:0对应一个句子A的token，1对应一个句子B的token。

模型输入我们可以自己构建，也可以利用Transformers中的分词器对象构建，操作如下：

from transformers import BertModel, BertTokenizer

bert = BertModel.from_pretrained("pre_model/bert-base-chinese")
tokenizer = BertTokenizer.from_pretrained("pre_model/bert-base-chinese")

test_sentence = "我在测试bert"
tokens = tokenizer.encode_plus(text=test_sentence)
print(tokens)

输出结果如下：

{'input_ids': [101, 2769, 1762, 3844, 6407, 8815, 8716, 102], 
'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0], 
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1]}

该方法就已经对输入的句子进行了分词，需要提醒的是：这里的tokenizer分词不同以往的分词方法，使用subword算法，并且为语句的句首和句首分别添加了[CLS]、[SEG]符号，我们可以瞅瞅：

tokenizer = BertTokenizer.from_pretrained("pre_model/bert-base-chinese")
test_sentence = "我在测试bert"
tokens = tokenizer.encode_plus(text=test_sentence)
print(tokens)
print(tokenizer.convert_ids_to_tokens(tokens['input_ids']))

输入的结果如下：

{'input_ids': [101, 2769, 1762, 3844, 6407, 8815, 8716, 102], 
'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0],
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1]}
['[CLS]', '我', '在', '测', '试', 'be', '##rt', '[SEP]']

其中将bert就分开了。

模型的输出

模型默认的输出是BaseModelOutputWithPoolingAndCrossAttentions，官方地址：https://huggingface.co/docs/transformers/main_classes/output#transformers.modeling_outputs.BaseModelOutputWithPoolingAndCrossAttentions，如下：

( last_hidden_state: FloatTensor = None,
pooler_output: FloatTensor = None,
hidden_states: typing.Optional[typing.Tuple[torch.FloatTensor]] = None,
past_key_values: typing.Optional[typing.Tuple[typing.Tuple[torch.FloatTensor]]] = None,
attentions: typing.Optional[typing.Tuple[torch.FloatTensor]] = None,
cross_attentions: typing.Optional[typing.Tuple[torch.FloatTensor]] = None )

我们经常使用的则是：

• last_hidden_state (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size)) — 模型最后一层输出的隐藏状态序列.
• hidden_states (tuple(torch.FloatTensor), optional, returned when output_hidden_states=True is passed or when config.output_hidden_states=True) — 形状为（batch_size、sequence_length、hidden_size）的torch.FloatTensor 的元组（一个用于嵌入的输出，如果模型有嵌入层，+ 一个用于每一层的输出）

使用案例如下：

from transformers import BertModel, BertTokenizer

bert = BertModel.from_pretrained("pre_model/bert-base-chinese")

tokenizer = BertTokenizer.from_pretrained("pre_model/bert-base-chinese")
test_sentence = "我在测试bert"
# 指定返回的数据是pytorch中的tensor数据类型
tokens = tokenizer.encode_plus(text=test_sentence, return_tensors='pt')
model_out = bert(**tokens)
print(model_out)

结果内容较多，读者可以运行上面的代码进行尝试。

总结

除了上面比较原始的预训练模型加载的类之外，针对不同类型的任务会有不同类型的Bert类，如：针对文本分类的有AutoModelForSequenceClassification，针对NER的有：AutoModelForTokenClassification等。当然具体该如何使用还需要参考官方文档。