How to do NER predictions with Huggingface BERT transformer

Question

I am trying to do a prediction on a test data set without any labels for an NER problem.

Here is some background. I am doing named entity recognition using tensorflow and Keras. I am using huggingface transformers.

I have two datasets. A train dataset and a test dataset. The training set has labels, the tests does not. Below you will see what a tokenized sentence looks like, what it's labels look like, and what it looks like after encoding

['The', 'pope', "isn't", 'really', 'making', 'much', 'of', 'an', 'effort', '.', 'He', "'s", 'wearing', 'the', 'same', 'clothes', 'as', 'yesterday', '.']
['O', 'B-person', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O']
[101, 1109, 17460, 2762, 112, 189, 1541, 1543, 1277, 1104, 1126, 3098, 119, 1124, 112, 188, 3351, 1103, 1269, 3459, 1112, 8128, 119, 102, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]

Here is the code on how I tokenized my text and encoded my labels

from transformers import DistilBertTokenizerFast

tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-cased')
train_encodings = tokenizer(train_texts, is_split_into_words=True, return_offsets_mapping=True, padding=True, truncation=True)
val_encodings = tokenizer(val_texts, is_split_into_words=True, return_offsets_mapping=True, padding=True, truncation=True)

def encode_tags(tags, encodings):
    labels = [[tag2id[tag] for tag in doc] for doc in tags]
    encoded_labels = []
    for doc_labels, doc_offset in zip(labels, encodings.offset_mapping):
        # create an empty array of -100
        doc_enc_labels = np.ones(len(doc_offset),dtype=int) * -100
        arr_offset = np.array(doc_offset)

        # set labels whose first offset position is 0 and the second is not 0
        doc_enc_labels[(arr_offset[:,0] == 0) & (arr_offset[:,1] != 0)] = doc_labels
        encoded_labels.append(doc_enc_labels.tolist())

    return encoded_labels

train_labels = encode_tags(train_tags, train_encodings)
val_labels = encode_tags(val_tags, val_encodings)

I have gotten my model to train and work. I'm getting pretty goods numbers when validating. Here is how that was done

from transformers import TFDistilBertForTokenClassification, TFTrainer, TFTrainingArguments

training_args = TFTrainingArguments(
    output_dir='./results',
    num_train_epochs=5,              # total number of training epochs
    per_device_train_batch_size=16,  # batch size per device during training
    per_device_eval_batch_size=16,   # batch size for evaluation
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    evaluation_strategy = "epoch",
    learning_rate = 2e-5
)

with training_args.strategy.scope():
    model = TFDistilBertForTokenClassification.from_pretrained('distilbert-base-cased', num_labels=len(unique_tags))

trainer = TFTrainer(
    model=model,                         # the instantiated 🤗 Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=val_dataset,            # evaluation dataset
    compute_metrics=compute_metrics
)

trainer.train()

trainer.evaluate()

My main issue is that I don't know how to predict with this. I'm not familiar with the library and the documentation has not been helping much.

I can apparently use trainer.predict(*param*), but I can't figure out what to actually input as the param.

On the other hand, when I do model.predict(param) where the param is the encoded sentence example I show above, I get this result

TFTokenClassifierOutput(loss=None, logits=array([[[-0.3232851 ,  0.12578554, -0.47193137, ...,  0.16509804,
          0.19799986, -0.3560003 ]],

       [[-1.8808482 , -1.07631   , -0.49765658, ..., -0.7443374 ,
         -1.2379731 , -0.5022731 ]],

       [[-1.4291595 , -1.8587289 , -1.5842767 , ..., -1.1863587 ,
         -0.21151644, -0.52205306]],

       ...,

       [[-1.6405941 , -1.2474233 , -1.0701559 , ..., -1.1816512 ,
          0.323739  , -0.45317683]],

       [[-1.6405947 , -1.247423  , -1.0701554 , ..., -1.1816509 ,
          0.3237388 , -0.45317668]],

       [[-1.6405947 , -1.247423  , -1.0701554 , ..., -1.1816509 ,
          0.3237388 , -0.45317668]]], dtype=float32), hidden_states=None, attentions=None)

I don't know how I'm supposed to take that result and decode it back into labels. What am I supposed to do with the logits array? How am I supposed to predict this?

Answer 1

Once the training is completed, use your trained model instance in a NER pipeline, using the same tokenizer as before:

from transformers import pipeline
nlp = pipeline('ner', model=model, tokenizer=tokenizer)
s = "I am Joe and live in London"
print(nlp(s))

Don't forget to save your model after training using save_pretrained!