# \[HELP\] NER task single sentence/sample prediction

**URL:** <https://discuss.huggingface.co/t/help-ner-task-single-sentence-sample-prediction/9460>\
**Category:** 🤗Transformers\
**Created:** [August 25, 2021, 7:25am UTC](https://discuss.huggingface.co/t/help-ner-task-single-sentence-sample-prediction/9460 "2021-08-25T07:25:22Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![Calin](https://avatars.discourse-cdn.com/v4/letter/c/a87d85/32.png) [@Calin](https://discuss.huggingface.co/u/Calin)\
**Post date:** [August 25, 2021, 7:25am UTC](https://discuss.huggingface.co/t/help-ner-task-single-sentence-sample-prediction/9460/1 "2021-08-25T07:25:22Z")

</div>

Hello everybody.

I am trying to predict with the NER model, as in the tutorial from huggingface (it contains only the training+evaluation part).

I am following this exact tutorial here : [https://github.com/huggingface/notebooks/blob/master/examples/token\_classification.ipynb](https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb).

It works flawlessly, but the problems that I have begin when I try to predict on a simple sample.

```auto
loaded_model = AutoModel.from_pretrained('./my_model_own_custom_training.pth', from_tf=False)

input_sentence = "John Nash is a great mathematician, he lives in France"
tokenized_input_sentence = tokenizer([input_sentence],
                                     truncation=True, 
                                     is_split_into_words=False,
                                     return_tensors='pt')
predictions = loaded_model(tokenized_input_sentence["input_ids"])[0]

```

predictions is of shape` (1,13,768)`

How can I arrive at the final result of the form [JOHN ↔ ‘B-PER’, … France ↔ “B-LOC”], where B-PER and B-LOC are two ground truth labels, representing the tag for a person and location respectively?

The result of the prediction is:

`torch.Size([1, 13, 768])`

If I write:

```auto
print(predictions.shape)
print(predictions.argmax(axis=2))
tensor([613, 705, 244, 620, 206, 206, 206, 620, 620, 620, 477, 693, 308])

```

I get the tensor above.

However I would have expected to get the tensor representing the ground truth [0…8] labels from the ground truth annotations.

Summary when loading the model :

> loading configuration file ./my\_model\_own\_custom\_training.pth/config.json  
> Model config DistilBertConfig {  
> “_name\_or\_path": “distilbert-base-uncased”,  
> “activation”: “gelu”,  
> “architectures”: [  
> “DistilBertForTokenClassification”  
> ],  
> “attention\_dropout”: 0.1,  
> “dim”: 768,  
> “dropout”: 0.1,  
> “hidden\_dim”: 3072,  
> “id2label”: {  
> “0”: “LABEL\_0”,  
> “1”: “LABEL\_1”,  
> “2”: “LABEL\_2”,  
> “3”: “LABEL\_3”,  
> “4”: “LABEL\_4”,  
> “5”: “LABEL\_5”,  
> “6”: “LABEL\_6”,  
> “7”: “LABEL\_7”,  
> “8”: “LABEL\_8”  
> },  
> “initializer\_range”: 0.02,  
> “label2id”: {  
> “LABEL\_0”: 0,  
> “LABEL\_1”: 1,  
> “LABEL\_2”: 2,  
> “LABEL\_3”: 3,  
> “LABEL\_4”: 4,  
> “LABEL\_5”: 5,  
> “LABEL\_6”: 6,  
> “LABEL\_7”: 7,  
> “LABEL\_8”: 8  
> },  
> “max\_position\_embeddings”: 512,  
> “model\_type”: “distilbert”,  
> “n\_heads”: 12,  
> “n\_layers”: 6,  
> “pad\_token\_id”: 0,  
> “qa\_dropout”: 0.1,  
> “seq\_classif\_dropout”: 0.2,  
> “sinusoidal\_pos\_embds”: false,  
> "tie\_weights_”: true,  
> “transformers\_version”: “4.8.1”,  
> “vocab\_size”: 30522  
> }

---

<div class="post-metadata">

**Author:** ![nielsr](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/nielsr/32/39617_2.png) [@nielsr](https://discuss.huggingface.co/u/nielsr)\
**Post date:** [August 25, 2021, 8:18am UTC](https://discuss.huggingface.co/t/help-ner-task-single-sentence-sample-prediction/9460/2 "2021-08-25T08:18:37Z")

</div>

> [@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:

```auto
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)`.

---

<div class="post-metadata">

**Author:** ![Calin](https://avatars.discourse-cdn.com/v4/letter/c/a87d85/32.png) [@Calin](https://discuss.huggingface.co/u/Calin)\
**Post date:** [August 25, 2021, 8:23am UTC](https://discuss.huggingface.co/t/help-ner-task-single-sentence-sample-prediction/9460/3 "2021-08-25T08:23:04Z")

</div>

Wow, it works perfectly.

Thank you for the awesome response!!!