Calin:

notebooks/token_classification.ipynb at master · huggingface/notebooks · GitHub


cc’ing @sgugger, these demo notebooks really need an inference part.
To do inference with NER, you need to load an AutoModelForTokenClassification rather than an AutoModel, like so:
from transformers import AutoTokenizer, AutoModelForTokenClassification

tokenizer = AutoTokenizer.from_pretrained("path_to_directory")
model = AutoModelForTokenClassification.from_pretrained("path_to_directory")

input_sentence = "John Nash is a great mathematician, he lives in France"
encoding = tokenizer([input_sentence], return_tensors='pt')

# forward pass
outputs = model(**encoding)
logits = outputs.logits

predictions = logits.argmax(-1)

The logits will be of shape (batch_size, seq_len, num_labels).