import torch
from pytorch_transformers import GPT2Tokenizer, GPT2LMHeadModel

# Load pre-trained GPT-2 tokenizer model: 
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# encode the words in a sentence: 
text = "What is the fastest car in the"
indexed_tokens = tokenizer.encode(text)

# convert tokens to a PyTorch tensor:
tokens_tensor = torch.tensor([indexed_tokens])

# load pre-trained model (weights)
model = GPT2LMHeadModel.from_pretrained('gpt2')

# "eval" mode deactivates the DropOut modules:
model.eval()

# Predict each token:
with torch.no_grad():
  outputs = model(tokens_tensor)
  predictions = outputs[0]

print("=> list of predictions:")
print(predictions[0, -1, :])
print()

print("=> argmax of predictions:")
print(torch.argmax(predictions[0, -1, :]).item())
print()

# Get the predicted next sub-word
predicted_index = torch.argmax(predictions[0, -1, :]).item()
predicted_text = tokenizer.decode(indexed_tokens + [predicted_index])

# Print the predicted word
print("=> initial text:") 
print(text)
print()

print("=> Predicted next word:")
print(predicted_text)

