How could I do simple inference on my Fine-Tuned Transformers NER models? | by Martin Keywood | Towards Data Science

I knew what I wanted to do. I wanted to generate NER in a biomedical domain.
I had done it in the wonderful scispaCy package, and even in Transformers via the amazing Simple Transformers, but I wanted to do it in the raw HuggingFace Transformers package.
Why? I had it working in scispaCy out the box with their pre-trained models, and in Simple Transformers with only a couple of lines of code and some preprocessing of data, so wasn’t that enough?
Of course not. I wanted to understand ‘how’ it worked. How to do it myself in Transformers, and then push the boundaries of what I could do.
I found so so many examples of how to train and / or fine-tune a Transformer model.
By a lot of experimenting and dark magic I was able to make some sense of the preprocessing of the data and ‘thought’ I had it right. After all, I had a model that trained. The loss going down sensibly. Awesome. Done, right?
Sadly, not. I want to use that model. Easily and sensibly. So how do I do simple inference on the model?
In Simple Transformers I could just call something: model.predict(["the text I was to predict on"]) and get my predictions back. No preprocessing the text and just so, well, simple. I guess that's why it was called Simple Transformers :)
Most other ML libraries I had used in the past (sklearn, fast.ai, etc) all had something similar. a ‘predict’ method on the model.
And trying to figure out how to do this in raw Transformers was where my problems started :(
OK so firstly I needed some data.
I thought I’d start with a well understood Biomedical NER dataset, BC5CDR, and a little Pre-Processing to get the data into a format to make it easier to manipulate the sentences.
import pandas as pd with open('train.tsv', 'r') as f: data = f.readlines() sentence_id = 1sentence_id_array = []word_array = []label_array = [] for line in data: line = line.strip() if len(line):  word, label = line.split('\t')  sentence_id_array.append(sentence_id)  word_array.append(word)  label_array.append(label) else:  sentence_id += 1 df = pd.DataFrame({ 'sentence_id': sentence_id_array, 'words': word_array, 'labels': label_array}) df.to_csv('train_clean.csv', index=False)
Firstly make sure everything is installed, and define some defaults settings that will help. (note: a lot of the Train section here was inspired by the wonderful run_ner example in Transformers)
!pip install transformers!pip install datasets!pip install seqevalimport osimport sysimport pandas as pdimport numpy as npfrom seqeval.metrics import accuracy_score, f1_score, precision_score, recall_scoreimport torchimport transformersfrom transformers import (AutoConfig, AutoModelForTokenClassification, AutoTokenizer, Trainer, TrainingArguments, set_seed)task_name = "ner"model_name_or_path = 'bert-base-cased'output_dir = 'tmp'overwrite_output_dir = Truenum_train_epochs = 1train_file = "train_clean.csv"preprocessing_num_workers = Nonepadding = "max_length"label_all_tokens = Trueseed = 42set_seed(seed)training_args = TrainingArguments(output_dir=output_dir, do_train=True, do_eval=True, num_train_epochs=num_train_epochs)
OK so let’s load the data and create some simple lookups of the associated labels.
df_train = pd.read_csv(train_file)df_train[‘sentence_id’] = df_train[‘sentence_id’].astype(int)text_column_name = “words”label_column_name = “labels”label_list = list(set(df_train[label_column_name]))label_to_id = {l: i for i, l in enumerate(label_list)}id_to_label = {label_to_id[x]: x  for x in label_to_id}num_labels = len(label_list)num_labels, label_list, label_to_id, id_to_label
[‘B-Disease’, ‘I-Disease’, ‘I-Chemical’, ‘B-Chemical’, ‘O’]
{‘B-Chemical’: 3, ‘B-Disease’: 0, ‘I-Chemical’: 2, ‘I-Disease’: 1, ‘O’: 4}
{0: ‘B-Disease’, 1: ‘I-Disease’, 2: ‘I-Chemical’, 3: ‘B-Chemical’, 4: ‘O’}
Load the pre-trained model and tokenizer.
config = AutoConfig.from_pretrained(model_name_or_path, num_labels=num_labels, finetuning_task=task_name)tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)model = AutoModelForTokenClassification.from_pretrained( model_name_or_path, from_tf=bool(".ckpt" in model_name_or_path), config=config)
We are doing ‘max length’ padding and need a little helper function to Tokenize all texts and align the labels with them.
Why??
The words are already separated into ‘tokens’ aren’t they? One per line, with a label for each word? We can just use those by passing in is_split_into_words = True and it's all fine, right?
On the surface that seems true, however given the tokenizer is based on the model, and it may be a Byte Pair Encoding Tokenizer (extremely likely) then the words are potentially split further.
So ['Selegiline', '-', 'induced', 'postural', 'hypotension', 'in', 'Parkinson', "'", 's', 'disease', ':', 'a', 'longitudinal', 'study', 'on', 'the', 'effects', 'of', 'drug', 'withdrawal', '.']
becomes ['[CLS]', 'Se', '##leg', '##ili', '##ne', '-', 'induced', 'post', '##ural', 'h', '##y', '##pot', '##ens', '##ion', 'in', 'Parkinson', "'", 's', 'disease', ':', 'a', 'longitudinal', 'study', 'on', 'the', 'effects', 'of', 'drug', 'withdrawal', '.', '[SEP]']
So regardless of the special tokens, the number of tokens goes from 21 to 29 in this example. And so the labels need to be rejigged to continue to align with the base words.
# Tokenize all texts and align the labels with them.def tokenize_and_align_labels(df): inputs = [] labels = [] for sid in range(1, max(df['sentence_id'])+1):  examples = list(df.loc[df['sentence_id'] == sid, text_column_name])  try:   tokenized_inputs = tokenizer(examples, padding=padding, truncation=True, is_split_into_words=True)   sentence_labels = list(df.loc[df['sentence_id'] == sid, label_column_name])   label_ids = []   for word_idx in tokenized_inputs.word_ids():    # Special tokens have a word id that is None. We set the label to -100 so they are automatically ignored in the loss function.    if word_idx is None:      label_ids.append(-100)    # We set the label for the first token of each word.    elif word_idx != previous_word_idx:     label_ids.append(label_to_id[sentence_labels[word_idx]])    # For the other tokens in a word, we set the label to either the current label or -100, depending on the label_all_tokens flag.    else:     label_ids.append(label_to_id[sentence_labels[word_idx]] if label_all_tokens else -100)    previous_word_idx = word_idxinputs.append(tokenized_inputs)   labels.append(torch.tensor(label_ids))  except:   pass return inputs, labelsencodings, labels = tokenize_and_align_labels(df_train)
Use tried and tested methods to split into train / val sets, and use those to generate some Transformer Datasets, ensuring we make them all Tensors along the way.
class MakeDataset(torch.utils.data.Dataset): def __init__(self, encodings, labels): self.encodings = encodings self.labels = labelsdef __getitem__(self, idx): item = dict() item[‘input_ids’] = torch.tensor(self.encodings[idx][‘input_ids’]) item[‘token_type_ids’] = torch.tensor(self.encodings[idx][‘token_type_ids’]) item[‘attention_mask’] = torch.tensor(self.encodings[idx][‘attention_mask’]) item[‘labels’] = torch.tensor(self.labels[idx]) return itemdef __len__(self): return len(self.labels)train_dataset = MakeDataset(train_encodings, train_labels)val_dataset = MakeDataset(val_encodings, val_labels)
And define some helpful metrics.
def compute_metrics(p): predictions, labels = p predictions = np.argmax(predictions, axis=2)# Remove ignored index (special tokens) true_predictions = [    [label_list[p] for (p, l) in zip(prediction, label) if l != -100]    for prediction, label in zip(predictions, labels) ] true_labels = [  [label_list[l] for (p, l) in zip(prediction, label) if l != -100]  for prediction, label in zip(predictions, labels) ]return {   "accuracy_score": accuracy_score(true_labels, true_predictions),   "precision": precision_score(true_labels, true_predictions),   "recall": recall_score(true_labels, true_predictions),   "f1": f1_score(true_labels, true_predictions), }
OK let’s train.
# Initialize our Trainertrainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_metrics)# Trainingtrainer.train(model_path=model_name_or_path)trainer.save_model()  # Saves the tokenizer too for easy upload# Evaluationresults = trainer.evaluate()output_eval_file = os.path.join(training_args.output_dir, "eval_results_ner.txt")if trainer.is_world_process_zero(): with open(output_eval_file, "w") as writer:  for key, value in results.items():   writer.write(f"{key} = {value}\n")
eval_loss = 0.11932332068681717
eval_accuracy_score = 0.9635168195718654
eval_precision = 0.9027568922305764
eval_recall = 0.9226434426229508
eval_f1 = 0.9125918419052444
Not bad after only 1 epoch :)
And this was where my problems started.
I could find nothing, nada, zero, nowt, zilch that could explain to me ‘simply’ how to do it.
There was no predict method on the model. The predict method on the Trainer just made no sense to me and required DataSets creating and I wanted to do inference on just a single string: “heart disease, a new house, a dose of penicillin, and bowel cancer. The diagnosis of COPD, a flashy new car and a skin rash.”
Was it badly written Transformers libraries? Not at all. It’s amazing.
Was it me just me being daft? Not enough coffee? Yeah I think so. One thing that was clear was it was ‘user error’ and when I finally discovered the answer I kicked myself:
Firstly, tokenize the same way we did for the training data
Then put the model in eval mode and just call the model passing in the ‘input_ids’ and ‘attention_mask’ (note, I am explicitly pushing them to cuda to be on the GPU, where my model was trained, but obviously remove the cuda calls if all on CPU), taking the argmax of the results
And decode that back to the original tokens and labels
test_encodings = tokenizer(["heart disease, a new house, a dose of penicillin, and bowel cancer. The diagnosis of COPD, a flashy new car and a skin rash."], padding=padding, truncation=True)model.eval()res = model(torch.tensor(test_encodings["input_ids"]).cuda(), attention_mask=torch.tensor(test_encodings["attention_mask"]).cuda())[0].argmax(dim=2)for i, enc in enumerate(test_encodings['input_ids'][0]): if enc:  print(tokenizer.decode(enc), "\t", id_to_label[res[0][i].item()])
And the results:
[CLS]   Oheart   B-Diseasedisease I-Disease,       Oa       Onew     Ohouse   O,       Oa       Odose    Oof      Open     B-Chemical##ici   B-Chemical##llin  B-Chemical,       Oand     Obow     B-Disease##el    B-Diseasecancer  I-Disease.       OThe     Odiagnosis Oof      OCO      B-Disease##PD    B-Disease,       Oa       Oflash   O##y     Onew     Ocar     Oand     Oa       Oskin    B-Diseaser       I-Disease##ash   I-Disease.       O[SEP]   O
Well that is just amazing. I mean, mind blowingly ace :) It’s picked out the diseases and the medication (chemicals).
And it got better. I could misspell ‘cancer’ as ‘canrce’, or ‘bowel’ as ‘bowl’ and it still found it.
This was all really exciting for me and has opened up the door to me exploring a lot more in this area.
Who’s laughing now, eh!!
My hope is this will help someone else in a similar situation and a similar stage on their NLP journey, where they follow all the tutorials to train and fine-tune and then say ‘now what? how do I use it’ :) That’s certainly where I was.
And if you’re interested and like it, please let me know and I will be sure to write up future experiments too :)