# set seed to reproduce results. Feel free to change the seed though to get different results
from transformers import set_seed
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
set_seed(42)

torch_device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained("gpt2")

# add the EOS token as PAD token to avoid warnings
model = AutoModelForCausalLM.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id).to(torch_device)


# encode context the generation is conditioned on
model_inputs = tokenizer('I enjoy walking with my cute dog', return_tensors='pt').to(torch_device)

# # activate sampling and deactivate top_k by setting top_k sampling to 0
# sample_output = model.generate(
#     **model_inputs,
#     max_new_tokens=40,
#     do_sample=True,
#     top_k=0
# )

# print("Output:\n" + 100 * '-')
# print(tokenizer.decode(sample_output[0], skip_special_tokens=True))

# # use temperature to decrease the sensitivity to low probability candidates
# sample_output = model.generate(
#     **model_inputs,
#     max_new_tokens=40,
#     do_sample=True,
#     top_k=0,
#     temperature=0.6,
# )

# print("Output:\n" + 100 * '-')
# print(tokenizer.decode(sample_output[0], skip_special_tokens=True))

# # set seed to reproduce results. Feel free to change the seed though to get different results
# set_seed(42)

# # set top_k to 50
# sample_output = model.generate(
#     **model_inputs,
#     max_new_tokens=40,
#     do_sample=True,
#     top_k=50
# )

# print("Output:\n" + 100 * '-')
# print(tokenizer.decode(sample_output[0], skip_special_tokens=True))

# # set seed to reproduce results. Feel free to change the seed though to get different results
# set_seed(42)

# # set top_k to 50
# sample_output = model.generate(
#     **model_inputs,
#     max_new_tokens=40,
#     do_sample=True,
#     top_p=0.92,
#     top_k=0
# )

# print("Output:\n" + 100 * '-')
# print(tokenizer.decode(sample_output[0], skip_special_tokens=True))


# set seed to reproduce results. Feel free to change the seed though to get different results
set_seed(42)
input_sequence = "Rephrase the following sentence: \"I enjoy walking with my cute dog, Maxine, and my cat, Elene\""

# encode context the generation is conditioned on
model_inputs = tokenizer(input_sequence, return_tensors='pt').to(torch_device)

# set top_k = 50 and set top_p = 0.95 and num_return_sequences = 3
sample_outputs = model.generate(
    **model_inputs,
    max_new_tokens=20,
    do_sample=True,
    top_k=50,
    top_p=0.95,
    num_return_sequences=3,
)

print("Output:\n" + 100 * '-')
for i, sample_output in enumerate(sample_outputs):
  print("{}: {}".format(i, tokenizer.decode(sample_output, skip_special_tokens=True)))