training_args = {
    "output_dir": "/tmp/clm_gpt2_xl",
    "dataset_name": "wikitext",
    "dataset_config_name": "wikitext-2-raw-v1",
    "num_train_epochs": 1,
    "per_device_train_batch_size": 4,
    "per_device_eval_batch_size": 4,
    "gradient_checkpointing": True,
    "do_train": True,
    "do_eval": True,
    "overwrite_output_dir": True,
    "use_cache": False,
}pretraining = False
model_name = "gpt2-xl"

if pretraining:
    training_args["config_name"] = model_name
    training_args["tokenizer_name"] = model_name
else:
    training_args["model_name_or_path"] = model_nametraining_args["use_habana"] = True                  # Whether to use HPUs or not
training_args["use_lazy_mode"] = True               # Whether to use lazy or eager mode
training_args["gaudi_config_name"] = "Habana/gpt2"  # Gaudi configuration to use
training_args["throughput_warmup_steps"] = 3        # Remove the first N training iterations from throughput computationfrom optimum.habana.distributed import DistributedRunner

# Build the command to execute
training_args_command_line = " ".join(f"--{key} {value}" for key, value in training_args.items())
command = f"../examples/language-modeling/run_clm.py {training_args_command_line}"

# # Instantiate a distributed runner
# distributed_runner = DistributedRunner(
#     command_list=[command],  # The command(s) to execute
#     world_size=8,            # The number of HPUs
#     use_mpi=True,            # OpenMPI is used for multi-processing
# )

# # Launch training
# ret_code = distributed_runner.run()training_args["deepspeed"] = "configs/deepspeed_zero_2.json"# Build the command to execute
training_args_command_line = " ".join(f"--{key} {value}" for key, value in training_args.items())
command = f"../examples/language-modeling/run_clm.py {training_args_command_line}"

# Instantiate a distributed runner
distributed_runner = DistributedRunner(
    command_list=[command],  # The command(s) to execute
    world_size=8,            # The number of HPUs
    use_deepspeed=True,      # Enable DeepSpeed
)

# Launch training
ret_code = distributed_runner.run()# The sequence to complete
prompt_text = "This live AI webinar is organized by Habana Labs and Hugging Face and"

import torch

from transformers import GPT2LMHeadModel, GPT2Tokenizer
from habana_frameworks.torch.utils.library_loader import load_habana_module

load_habana_module()  # required to set torch.device("hpu")

path_to_model = training_args["output_dir"]  # the folder where everything related to our run was saved

device = torch.device("hpu")

# Load the tokenizer and the model
tokenizer = GPT2Tokenizer.from_pretrained(path_to_model)
model = GPT2LMHeadModel.from_pretrained(path_to_model)
model.to(device)

# Encode the prompt
encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
encoded_prompt = encoded_prompt.to(device)

# Generate the following of the prompt
output_sequences = model.generate(
    input_ids=encoded_prompt,
    max_length=16 + len(encoded_prompt[0]),
    do_sample=True,
    num_return_sequences=3,
)

# Remove the batch dimension when returning multiple sequences
if len(output_sequences.shape) > 2:
    output_sequences.squeeze_()

generated_sequences = []

for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
    print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===")
    generated_sequence = generated_sequence.tolist()

    # Decode text
    text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)

    # Remove all text after the stop token
    text = text[: text.find(".")]

    # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
    total_sequence = (
        prompt_text + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
    )

    generated_sequences.append(total_sequence)
    print(total_sequence)import numpy as np

gaudi_price_per_hour = 13.10904
v100_price_per_hour = 12.24

print(f"Gaudi    (dl1.24xlarge): training time = 630s, cost = {np.round(630 * gaudi_price_per_hour / 3600, 2)}$ ({gaudi_price_per_hour}$/hr)")
print(f"4 x V100 (p3.8xlarge)  : training time = 858s, cost = {np.round(858 * v100_price_per_hour / 3600, 2)}$ ({v100_price_per_hour}$/hr)")