如何使用训练有素的BERT NER（命名实体识别）模型来预测新示例？

Question

我在这篇中等职位之后训练了自己的BERT NER：https://medium.com/@yingbiao/ner-with-bert-in-action-936ff275bc73

我将模型保存到光盘并成功加载。

model = BertForTokenClassification.from_pretrained(bert_out_address, num_labels=len(tag2idx))

model.eval（）有效：

model.eval()

我是BERT和Transformer lib的新手。我希望类似

model.predict('Hello I am an example sentence') 

会告诉我公认的实体。

我也尝试过：

input_ids = torch.tensor([tokenizer.encode("Here is some text to encode")])
output = model(input_ids)

其中的输出给了我一个大的张量，我不知道该怎么做。

现在如何使用模型来预测例句中的实体？我应该如何处理输出？

谢谢！

Answer 1

BertForTokenClassification的文档说，它会在softmax之前返回分数，即标记的非标准化概率。

您可以通过从分布中获取最大值（应为维度2）来解码标签。这将为您提供最可能的标签的索引。在训练模型之前，您已将标签转换为索引（使用tag2idx表），现在您需要执行相反的过程以及ID的标签。