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