Views
No views yet
1from transformers import AutoTokenizer, AutoModelForTokenClassification
2
3tokenizer = AutoTokenizer.from_pretrained("ageng-anugrah/indobert-large-p2-finetuned-ner")
4model = AutoModelForTokenClassification.from_pretrained("ageng-anugrah/indobert-large-p2-finetuned-ner")1import torch
2def predict(model, tokenizer, sentence):
3 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
4 inputs = tokenizer(sentence.split(),
5 is_split_into_words = True,
6 return_offsets_mapping=True,
7 return_tensors="pt",
8 padding='max_length',
9 truncation=True,
10 max_length=512)
11
12 model.to(device)
13 # move to gpu
14 ids = inputs["input_ids"].to(device)
15 mask = inputs["attention_mask"].to(device)
16
17 # forward pass
18 outputs = model(ids, attention_mask=mask)
19 logits = outputs[0]
20
21 active_logits = logits.view(-1, model.num_labels) # shape (batch_size * seq_len, num_labels)
22 flattened_predictions = torch.argmax(active_logits, axis=1) # shape (batch_size*seq_len,) - predictions at the token level
23
24 tokens = tokenizer.convert_ids_to_tokens(ids.squeeze().tolist())
25 token_predictions = [model.config.id2label[i] for i in flattened_predictions.cpu().numpy()]
26 wp_preds = list(zip(tokens, token_predictions)) # list of tuples. Each tuple = (wordpiece, prediction)
27
28 prediction = []
29 for token_pred, mapping in zip(wp_preds, inputs["offset_mapping"].squeeze().tolist()):
30 #only predictions on first word pieces are important
31 if mapping[0] == 0 and mapping[1] != 0:
32 prediction.append(token_pred[1])
33 else:
34 continue
35
36 return sentence.split(), prediction
37
38sentence = "BJ Habibie adalah Presiden Indonesia ke-3 yang lahir pada tanggl 25 Juni 1936"
39words, labels = predict(model, tokenizer, sentence)