ner_inference.py

import torch
import os
from models.BertSequence import load_bert_sequence_model
from transformers import BertTokenizerFast
from dataloaders.ner_conll2003 import get_labels

def inference_ner(sentence, path_checkpoint, num_labels, device='cuda'):
    model = load_bert_sequence_model(path_checkpoint, num_labels, device)
    tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
    ids_to_labels = {i: label for i, label in enumerate(get_labels())}
    
    inputs = tokenizer(sentence.split(),
                        is_split_into_words=True,
                        return_offsets_mapping=True,
                        padding='max_length',
                        truncation=True,
                        max_length=128,
                        return_tensors="pt"
                    )
    ids = inputs["input_ids"].to(device)
    mask = inputs["attention_mask"].to(device)

    outputs = model.module.predict(input_ids=ids, attention_mask=mask)

    logits = outputs[0]
    
    active_logits = logits.view(-1, model.module.num_labels) # shape (batch_size * seq_len, num_labels)
    flattened_predictions = torch.argmax(active_logits, axis=1) # shape (batch_size*seq_len,) - predictions at the token level

    tokens = tokenizer.convert_ids_to_tokens(ids.squeeze().tolist())
    token_predictions = [ids_to_labels[i] for i in flattened_predictions.cpu().numpy()]
    wp_preds = list(zip(tokens, token_predictions)) # list of tuples. Each tuple = (wordpiece, prediction)

    prediction = []
    for token_pred, mapping in zip(wp_preds, inputs["offset_mapping"].squeeze().tolist()):
        #only predictions on first word pieces are important
        if mapping[0] == 0 and mapping[1] != 0:
            prediction.append(token_pred[1])
        else:
            continue

    print(sentence.split())
    print(prediction)

if __name__ =='__main__':
    sentence = "@HuggingFace is a company based in New York, but is also has employees working in Paris"
    path_checkpoint = os.path.join('checkpoints', 'conll2003', 'SEED0_NER_CoNLL2003_clean_data', 'best.pt')
    num_labels = len(get_labels())
    inference_ner(sentence, path_checkpoint, num_labels)
    """
    Result:
    ['@HuggingFace', 'is', 'a', 'company', 'based', 'in', 'New', 'York,', 'but', 'is', 'also', 'has', 'employees', 'working', 'in', 'Paris']
    ['B-ORG', 'O', 'O', 'O', 'O', 'O', 'B-LOC', 'I-LOC', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'B-LOC']
    """