import torch
from pytorch_pretrained_bert import BertTokenizer, BertModel, BertForMaskedLM

# Load pre-trained model tokenizer (vocabulary)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Tokenized input
text = "Who was <NAME> ? <NAME> was a puppeteer"
tokenized_text = tokenizer.tokenize(text)

# Mask a token that we will try to predict back with `BertForMaskedLM`
masked_index = 6
tokenized_text[masked_index] = '[MASK]'
assert tokenized_text == ['who', 'was', 'jim', 'henson', '?', 'jim', '[MASK]', 'was', 'a', 'puppet', '##eer']

# Convert token to vocabulary indices
indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
# Define sentence A and B indices associated to 1st and 2nd sentences (see paper)
segments_ids = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]

# Convert inputs to PyTorch tensors
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])

# ========================= BertForMaskedLM ==============================
# Load pre-trained model (weights)
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model.eval()

"""
predictions.size():
    torch.Size([1, 11, 30522])

predictions[0, masked_index]:
    tensor([-7.8384, -7.8162, -7.8893,  ..., -6.9924, -6.1897, -4.5417],
       grad_fn=<SelectBackward>)

predictions[0, masked_index].size():
    torch.Size([30522])
"""
# Predict all tokens
predictions = model(tokens_tensor, segments_tensors)

# 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'

import torch
from pytorch_pretrained_bert import BertTokenizer, BertModel, BertForMaskedLM

# Load pre-trained model tokenizer (vocabulary)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Tokenized input
text = "Who was <NAME> ? <NAME> was a puppeteer"
tokenized_text = tokenizer.tokenize(text)

# Mask a token that we will try to predict back with `BertForMaskedLM`
masked_index = 6
tokenized_text[masked_index] = '[MASK]'
assert tokenized_text == ['who', 'was', 'jim', 'henson', '?', 'jim', '[MASK]', 'was', 'a', 'puppet', '##eer']

# Convert token to vocabulary indices
indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
# Define sentence A and B indices associated to 1st and 2nd sentences (see paper)
segments_ids = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]

# Convert inputs to PyTorch tensors
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])

# ========================= BertForMaskedLM ==============================
# Load pre-trained model (weights)
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model.eval()

"""
predictions.size():
    torch.Size([1, 11, 30522])

predictions[0, masked_index]:
    tensor([-7.8384, -7.8162, -7.8893,  ..., -6.9924, -6.1897, -4.5417],
       grad_fn=<SelectBackward>)

predictions[0, masked_index].size():
    torch.Size([30522])
"""
# Predict all tokens
predictions = model(tokens_tensor, segments_tensors)

# 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'

# Load pre-trained model tokenizer (vocabulary) # Tokenized input # Mask a token that we will try to predict back with `BertForMaskedLM` #eer'] # Convert token to vocabulary indices # Define sentence A and B indices associated to 1st and 2nd sentences (see paper) # Convert inputs to PyTorch tensors # ========================= BertForMaskedLM ============================== # Load pre-trained model (weights) predictions.size():     torch.Size([1, 11, 30522])  predictions[0, masked_index]:     tensor([-7.8384, -7.8162, -7.8893,  ..., -6.9924, -6.1897, -4.5417],        grad_fn=<SelectBackward>)  predictions[0, masked_index].size():     torch.Size([30522]) # Predict all tokens # confirm we were able to predict 'henson'