CODEWITHSUNDEEP
named-entity-recognitionbert-language-modelhuggingface-transformers
quisi
quisi

Reputation: 21

How do I use my trained BERT NER (named entity recognition) model to predict a new example?

I trained my own BERT NER following this Medium post: https://medium.com/@yingbiao/ner-with-bert-in-action-936ff275bc73

I saved my model to the disc and successfully loaded it. 

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

model.eval() works:

model.eval()

I am new to BERT and the transformer lib. I would expect something similar to 

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

would show me recognised entities.

I also tried:

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

where output gives me a big tensor and I do not know what to do with it.

How can I use the model now to predict the entities in an example sentence? What am I supposed to do with the output?

Thanks!

Upvotes: 2

Views: 2985

Answers (1)

Jindřich
Jindřich

Reputation: 11213

The documentation of BertForTokenClassification says it returns scores before softmax, i.e., unnormalized probabilities of the tags.

You can decode the tags by taking the maximum from the distributions (should be dimension 2). This will give you indices of the most probable tags. Before you trained the model, you converted the tags into indices (using tag2idx table), now you need to do the reversed process and the tags for the ids.

Upvotes: 4

Related Questions