Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
# Load pre-trained model (weights)
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
# If you have a GPU, put everything on cuda
tokens_tensor ='cuda')
segments_tensors ='cuda')'cuda')
# Predict all tokens
with torch.no_grad():
outputs = model(tokens_tensor, token_type_ids=segments_tensors)
predictions = outputs[0]
# confirm we were able to predict 'henson'
predicted_index = torch.argmax(predictions[0, masked_index]).item()
predicted_token = tokenizer.convert_ids_to_tokens([predicted_index])[0]
assert predicted_token == 'henson'
print('Predicted token is:',predicted_token)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment