首页 > 解决方案 > 如何使用经过训练的 BERT NER(命名实体识别)模型来预测新示例?

问题描述

我在这篇 Medium 帖子之后训练了自己的 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 和变压器库的新手。我希望有类似的东西

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

会向我展示公认的实体。

我也试过:

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

其中输出给了我一个很大的张量,我不知道如何处理它。

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

谢谢!

标签: named-entity-recognitionbert-language-modelhuggingface-transformers

解决方案


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

您可以通过从分布中获取最大值来解码标签(应该是维度 2)。这将为您提供最可能标签的索引。在训练模型之前,您将标签转换为索引(使用tag2idx表格),现在您需要执行相反的过程和 id 的标签。


推荐阅读