Views
No views yet
1from transformers import AutoTokenizer, AlbertForMaskedLM
2import torch
3from torch.nn.functional import softmax
4
5pretrained = 'voidful/albert_chinese_base'
6tokenizer = AutoTokenizer.from_pretrained(pretrained)
7model = AlbertForMaskedLM.from_pretrained(pretrained)
8
9inputtext = "今天[MASK]情很好"
10
11maskpos = tokenizer.encode(inputtext, add_special_tokens=True).index(103)
12
13input_ids = torch.tensor(tokenizer.encode(inputtext, add_special_tokens=True)).unsqueeze(0) # Batch size 1
14outputs = model(input_ids, labels=input_ids)
15loss, prediction_scores = outputs[:2]
16logit_prob = softmax(prediction_scores[0, maskpos],dim=-1).data.tolist()
17predicted_index = torch.argmax(prediction_scores[0, maskpos]).item()
18predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])[0]
19print(predicted_token, logit_prob[predicted_index])感 0.36333346366882324