-
Notifications
You must be signed in to change notification settings - Fork 2
/
ner_inference.py
executable file
·54 lines (45 loc) · 2.39 KB
/
ner_inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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']
"""