17d16
< """ Conditional text generation with the auto-regressive models of the library (GPT/GPT-2/CTRL/Transformer-XL/XLNet)
19c18,19
< 
---
> Conditional text generation on Habana Gaudi/Gaudi2.
> """
23c23,25
< from typing import Tuple
---
> import os
> import tempfile
> import time
25d26
< import numpy as np
27,43c28,31
< 
< from transformers import (
<     CTRLLMHeadModel,
<     CTRLTokenizer,
<     GenerationMixin,
<     GPT2LMHeadModel,
<     GPT2Tokenizer,
<     OpenAIGPTLMHeadModel,
<     OpenAIGPTTokenizer,
<     TransfoXLLMHeadModel,
<     TransfoXLTokenizer,
<     XLMTokenizer,
<     XLMWithLMHeadModel,
<     XLNetLMHeadModel,
<     XLNetTokenizer,
< )
< from transformers.modeling_outputs import CausalLMOutputWithPast
---
> import torch.nn.functional as F
> from checkpoint_utils import model_is_bloom, write_checkpoints_json
> from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
> from transformers.generation import GenerationConfig
53,280d40
< MAX_LENGTH = int(10000)  # Hardcoded max length to avoid infinite loop
< 
< MODEL_CLASSES = {
<     "gpt2": (GPT2LMHeadModel, GPT2Tokenizer),
<     "ctrl": (CTRLLMHeadModel, CTRLTokenizer),
<     "openai-gpt": (OpenAIGPTLMHeadModel, OpenAIGPTTokenizer),
<     "xlnet": (XLNetLMHeadModel, XLNetTokenizer),
<     "transfo-xl": (TransfoXLLMHeadModel, TransfoXLTokenizer),
<     "xlm": (XLMWithLMHeadModel, XLMTokenizer),
< }
< 
< # Padding text to help Transformer-XL and XLNet with short prompts as proposed by Aman Rusia
< # in https://github.com/rusiaaman/XLNet-gen#methodology
< # and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
< PREFIX = """In 1991, the remains of Russian Tsar Nicholas II and his family
< (except for Alexei and Maria) are discovered.
< The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the
< remainder of the story. 1883 Western Siberia,
< a young Grigori Rasputin is asked by his father and a group of men to perform magic.
< Rasputin has a vision and denounces one of the men as a horse thief. Although his
< father initially slaps him for making such an accusation, Rasputin watches as the
< man is chased outside and beaten. Twenty years later, Rasputin sees a vision of
< the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous,
< with people, even a bishop, begging for his blessing. <eod> </s> <eos>"""
< 
< 
< def set_seed(args):
<     np.random.seed(args.seed)
<     torch.manual_seed(args.seed)
<     if args.n_gpu > 0:
<         torch.cuda.manual_seed_all(args.seed)
< 
< 
< #
< # Functions to prepare models' input
< #
< 
< 
< def prepare_ctrl_input(args, _, tokenizer, prompt_text):
<     if args.temperature > 0.7:
<         logger.info("CTRL typically works better with lower temperatures (and lower top_k).")
< 
<     encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False)
<     if not any(encoded_prompt[0] == x for x in tokenizer.control_codes.values()):
<         logger.info("WARNING! You are not starting your generation from a control code so you won't get good results")
<     return prompt_text
< 
< 
< def prepare_xlm_input(args, model, tokenizer, prompt_text):
<     # kwargs = {"language": None, "mask_token_id": None}
< 
<     # Set the language
<     use_lang_emb = hasattr(model.config, "use_lang_emb") and model.config.use_lang_emb
<     if hasattr(model.config, "lang2id") and use_lang_emb:
<         available_languages = model.config.lang2id.keys()
<         if args.xlm_language in available_languages:
<             language = args.xlm_language
<         else:
<             language = None
<             while language not in available_languages:
<                 language = input("Using XLM. Select language in " + str(list(available_languages)) + " >>> ")
< 
<         model.config.lang_id = model.config.lang2id[language]
<         # kwargs["language"] = tokenizer.lang2id[language]
< 
<     # TODO fix mask_token_id setup when configurations will be synchronized between models and tokenizers
<     # XLM masked-language modeling (MLM) models need masked token
<     # is_xlm_mlm = "mlm" in args.model_name_or_path
<     # if is_xlm_mlm:
<     #     kwargs["mask_token_id"] = tokenizer.mask_token_id
< 
<     return prompt_text
< 
< 
< def prepare_xlnet_input(args, _, tokenizer, prompt_text):
<     prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
<     prompt_text = prefix + prompt_text
<     return prompt_text
< 
< 
< def prepare_transfoxl_input(args, _, tokenizer, prompt_text):
<     prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
<     prompt_text = prefix + prompt_text
<     return prompt_text
< 
< 
< PREPROCESSING_FUNCTIONS = {
<     "ctrl": prepare_ctrl_input,
<     "xlm": prepare_xlm_input,
<     "xlnet": prepare_xlnet_input,
<     "transfo-xl": prepare_transfoxl_input,
< }
< 
< 
< def adjust_length_to_model(length, max_sequence_length):
<     if length < 0 and max_sequence_length > 0:
<         length = max_sequence_length
<     elif 0 < max_sequence_length < length:
<         length = max_sequence_length  # No generation bigger than model size
<     elif length < 0:
<         length = MAX_LENGTH  # avoid infinite loop
<     return length
< 
< 
< def sparse_model_config(model_config):
<     embedding_size = None
<     if hasattr(model_config, "hidden_size"):
<         embedding_size = model_config.hidden_size
<     elif hasattr(model_config, "n_embed"):
<         embedding_size = model_config.n_embed
<     elif hasattr(model_config, "n_embd"):
<         embedding_size = model_config.n_embd
< 
<     num_head = None
<     if hasattr(model_config, "num_attention_heads"):
<         num_head = model_config.num_attention_heads
<     elif hasattr(model_config, "n_head"):
<         num_head = model_config.n_head
< 
<     if embedding_size is None or num_head is None or num_head == 0:
<         raise ValueError("Check the model config")
< 
<     num_embedding_size_per_head = int(embedding_size / num_head)
<     num_layer = model_config.n_layer
< 
<     return num_layer, num_head, num_embedding_size_per_head
< 
< 
< def prepare_jit_inputs(inputs, model, tokenizer):
<     num_batch = len(inputs)
<     dummy_input = tokenizer.batch_encode_plus(inputs, return_tensors="pt", padding=True)
<     num_block_layers, num_attention_heads, num_embedding_size_per_head = sparse_model_config(model.config)
<     if model.config.model_type == "bloom":
<         past_key_values = tuple(
<             (
<                 torch.zeros(int(num_attention_heads * num_batch), num_embedding_size_per_head, 1)
<                 .to(model.config.torch_dtype)
<                 .to(model.device),
<                 torch.zeros(int(num_attention_heads * num_batch), 1, num_embedding_size_per_head)
<                 .to(model.config.torch_dtype)
<                 .to(model.device),
<             )
<             for _ in range(num_block_layers)
<         )
<     else:
<         past_key_values = tuple(
<             (
<                 torch.zeros(num_batch, num_attention_heads, 1, num_embedding_size_per_head)
<                 .to(model.config.torch_dtype)
<                 .to(model.device),
<                 torch.zeros(num_batch, num_attention_heads, 1, num_embedding_size_per_head)
<                 .to(model.config.torch_dtype)
<                 .to(model.device),
<             )
<             for _ in range(num_block_layers)
<         )
< 
<     dummy_input["attention_mask"] = torch.cat(
<         [
<             torch.zeros(dummy_input["attention_mask"].shape[0], 1).to(dummy_input["attention_mask"].dtype),
<             dummy_input["attention_mask"],
<         ],
<         -1,
<     )
< 
<     if model.config.use_cache:
<         jit_inputs = (
<             dummy_input["input_ids"].to(model.device),
<             past_key_values,
<             dummy_input["attention_mask"].to(model.device),
<         )
<     else:
<         jit_inputs = (
<             dummy_input["input_ids"].to(model.device),
<             dummy_input["attention_mask"].to(model.device),
<         )
< 
<     return jit_inputs
< 
< 
< class _ModelFallbackWrapper(GenerationMixin):
<     __slots__ = ("_optimized", "_default")
< 
<     def __init__(self, optimized, default):
<         self._optimized = optimized
<         self._default = default
< 
<     def __call__(self, *args, **kwargs):
<         if kwargs["past_key_values"] is None:
<             return self._default(*args, **kwargs)
<         trace_graph_inputs = []
<         kwargs.pop("position_ids", None)
<         for k, v in kwargs.items():
<             if v is not None and not isinstance(v, bool):
<                 trace_graph_inputs.append(v)
<         trace_graph_inputs = tuple(trace_graph_inputs)
<         outputs = self._optimized(*trace_graph_inputs)
<         lm_logits = outputs[0]
<         past_key_values = outputs[1]
<         fixed_output = CausalLMOutputWithPast(
<             loss=None,
<             logits=lm_logits,
<             past_key_values=past_key_values,
<             hidden_states=None,
<             attentions=None,
<         )
<         return fixed_output
< 
<     def __getattr__(self, item):
<         return getattr(self._default, item)
< 
<     def prepare_inputs_for_generation(
<         self, input_ids, past_key_values=None, inputs_embeds=None, use_cache=None, **kwargs
<     ):
<         return self._default.prepare_inputs_for_generation(
<             input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, **kwargs
<         )
< 
<     def _reorder_cache(
<         self, past_key_values: Tuple[Tuple[torch.Tensor]], beam_idx: torch.Tensor
<     ) -> Tuple[Tuple[torch.Tensor]]:
<         """
<         This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
<         [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
<         beam_idx at every generation step.
<         """
<         return self._default._reorder_cache(past_key_values, beam_idx)
< 
282a43
>     # Arguments management
285c46
<         "--model_type",
---
>         "--model_name_or_path",
289c50
<         help="Model type selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
---
>         help="Path to pre-trained model (on the HF Hub or locally).",
292c53
<         "--model_name_or_path",
---
>         "--gaudi_config_name_or_path",
295,307c56
<         required=True,
<         help="Path to pre-trained model or shortcut name selected in the list: " + ", ".join(MODEL_CLASSES.keys()),
<     )
< 
<     parser.add_argument("--prompt", type=str, default="")
<     parser.add_argument("--length", type=int, default=20)
<     parser.add_argument("--stop_token", type=str, default=None, help="Token at which text generation is stopped")
< 
<     parser.add_argument(
<         "--temperature",
<         type=float,
<         default=1.0,
<         help="temperature of 1.0 has no effect, lower tend toward greedy sampling",
---
>         help="Path to Gaudi configuration (on the HF Hub or locally).",
308a58,61
>     parser.add_argument("--max_new_tokens", type=int, default=100)
>     parser.add_argument("--batch_size", type=int, default=1, help="Input batch size.")
>     parser.add_argument("--n_iterations", type=int, default=5, help="Number of inference iterations.")
>     parser.add_argument("--local_rank", type=int, default=-1, metavar="N", help="Local process rank.")
310c63,65
<         "--repetition_penalty", type=float, default=1.0, help="primarily useful for CTRL model; in that case, use 1.2"
---
>         "--use_kv_cache",
>         action="store_true",
>         help="Whether to use the key/value cache for decoding. It should speed up generation.",
312,321c67
<     parser.add_argument("--k", type=int, default=0)
<     parser.add_argument("--p", type=float, default=0.9)
< 
<     parser.add_argument("--prefix", type=str, default="", help="Text added prior to input.")
<     parser.add_argument("--padding_text", type=str, default="", help="Deprecated, the use of `--prefix` is preferred.")
<     parser.add_argument("--xlm_language", type=str, default="", help="Optional language when used with the XLM model.")
< 
<     parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
<     parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
<     parser.add_argument("--num_return_sequences", type=int, default=1, help="The number of samples to generate.")
---
>     parser.add_argument("--use_hpu_graphs", action="store_true", help="Whether to use HPU graphs or not.")
323,325c69,72
<         "--fp16",
<         action="store_true",
<         help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
---
>         "--dataset_name",
>         default=None,
>         type=str,
>         help="Optional argument if you want to assess your model on a given dataset of the HF Hub.",
328c75,78
<         "--jit", type=bool, default=False, help="Whether or not to use jit trace to accelerate inference"
---
>         "--column_name",
>         default=None,
>         type=str,
>         help="Optional argument if you want to assess your model on a given dataset of the HF Hub, this will be the name of the column to use as prompts for generation.",
330,333d79
<     args = parser.parse_args()
< 
<     args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
<     args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
335,350c81
<     logger.warning(f"device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fp16}")
< 
<     set_seed(args)
< 
<     # Initialize the model and tokenizer
<     try:
<         args.model_type = args.model_type.lower()
<         model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
<     except KeyError:
<         raise KeyError("the model {} you specified is not supported. You are welcome to add it and open a PR :)")
< 
<     tokenizer = tokenizer_class.from_pretrained(args.model_name_or_path)
<     if tokenizer.pad_token is None:
<         tokenizer.pad_token = tokenizer.eos_token
<     model = model_class.from_pretrained(args.model_name_or_path)
<     model.to(args.device)
---
>     args = parser.parse_args()
352,353c83,110
<     if args.fp16:
<         model.half()
---
>     # If the DeepSpeed launcher is used, the env variable _ will be equal to /usr/local/bin/deepspeed
>     # For multi node, the value of the env variable WORLD_SIZE should be larger than 8
>     use_deepspeed = "deepspeed" in os.environ["_"] or (
>         "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 8
>     )
>     if use_deepspeed:
>         # Set necessary env variables
>         os.environ.setdefault("PT_HPU_LAZY_ACC_PAR_MODE", "0")
>         os.environ.setdefault("PT_HPU_ENABLE_LAZY_COLLECTIVES", "true")
> 
>     # Device is HPU
>     args.device = "hpu"
>     import habana_frameworks.torch.hpu as torch_hpu
> 
>     # Get world size, rank and local rank
>     from habana_frameworks.torch.distributed.hccl import initialize_distributed_hpu
> 
>     world_size, rank, args.local_rank = initialize_distributed_hpu()
> 
>     if use_deepspeed:
>         # Check if DeepSpeed is installed
>         from transformers.deepspeed import is_deepspeed_available
> 
>         if not is_deepspeed_available():
>             raise ImportError(
>                 "This script requires deepspeed: `pip install" " git+https://github.com/HabanaAI/DeepSpeed.git@1.9.0`."
>             )
>         import deepspeed
355,356c112,122
<     args.length = adjust_length_to_model(args.length, max_sequence_length=model.config.max_position_embeddings)
<     logger.info(args)
---
>         # Initialize process(es) for DeepSpeed
>         deepspeed.init_distributed(dist_backend="hccl")
>         logger.info("DeepSpeed is enabled.")
>     else:
>         if args.gaudi_config_name_or_path is None:
>             gaudi_config = None
>             logger.warning(
>                 "`--gaudi_config_name_or_path` was not specified so not using Habana Mixed Precision for this run."
>             )
>         else:
>             from optimum.habana import GaudiConfig
358c124
<     prompt_text = args.prompt if args.prompt else input("Model prompt >>> ")
---
>             gaudi_config = GaudiConfig.from_pretrained(args.gaudi_config_name_or_path)
360,364c126,127
<     # Different models need different input formatting and/or extra arguments
<     requires_preprocessing = args.model_type in PREPROCESSING_FUNCTIONS.keys()
<     if requires_preprocessing:
<         prepare_input = PREPROCESSING_FUNCTIONS.get(args.model_type)
<         preprocessed_prompt_text = prepare_input(args, model, tokenizer, prompt_text)
---
>             if gaudi_config.use_habana_mixed_precision:
>                 from habana_frameworks.torch.hpex import hmp
366,367c129,160
<         if model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
<             tokenizer_kwargs = {"add_space_before_punct_symbol": True}
---
>                 # Open temporary files to mixed-precision write ops
>                 with tempfile.NamedTemporaryFile() as hmp_bf16_file:
>                     with tempfile.NamedTemporaryFile() as hmp_fp32_file:
>                         # hmp.convert needs ops to be written in text files
>                         gaudi_config.write_bf16_fp32_ops_to_text_files(
>                             hmp_bf16_file.name,
>                             hmp_fp32_file.name,
>                         )
>                         hmp.convert(
>                             opt_level=gaudi_config.hmp_opt_level,
>                             bf16_file_path=hmp_bf16_file.name,
>                             fp32_file_path=hmp_fp32_file.name,
>                             isVerbose=gaudi_config.hmp_is_verbose,
>                         )
>         logger.info("Single-device run.")
> 
>     # Tweak generation so that it runs faster on Gaudi
>     from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi
> 
>     adapt_transformers_to_gaudi()
> 
>     tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
> 
>     if use_deepspeed:
>         config = AutoConfig.from_pretrained(args.model_name_or_path)
>         args.dtype = torch.bfloat16
>         is_bloom = model_is_bloom(config)
> 
>         if is_bloom:
>             # Construct model with fake meta tensors, later will be replaced on devices during ds-inference ckpt load
>             with deepspeed.OnDevice(dtype=args.dtype, device="meta"):
>                 model = AutoModelForCausalLM.from_config(config, torch_dtype=args.dtype)
369,373c162,187
<             tokenizer_kwargs = {}
< 
<         encoded_prompt = tokenizer.encode(
<             preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt", **tokenizer_kwargs
<         )
---
>             with deepspeed.OnDevice(dtype=args.dtype, device=args.device):
>                 model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, torch_dtype=args.dtype)
>         model = model.eval()
> 
>         # Initialize the model
>         ds_inference_kwargs = {"dtype": args.dtype}
>         ds_inference_kwargs["tensor_parallel"] = {"tp_size": world_size}
>         ds_inference_kwargs["enable_cuda_graph"] = args.use_hpu_graphs
> 
>         # BLOOM is managed differently
>         if is_bloom:
>             checkpoints_json = "checkpoints.json"
>             write_checkpoints_json(args.model_name_or_path, args.local_rank, checkpoints_json)
> 
>             # Make sure all devices/nodes have access to the model checkpoints
>             torch.distributed.barrier()
> 
>             from transformers.models.bloom.modeling_bloom import BloomBlock
> 
>             ds_inference_kwargs["injection_policy"] = {BloomBlock: ("self_attention.dense", "mlp.dense_4h_to_h")}
>             ds_inference_kwargs["checkpoint"] = checkpoints_json
> 
>         model = deepspeed.init_inference(model, **ds_inference_kwargs)
>         if is_bloom:
>             model.module.split_lm_head()
>         model = model.module
375,377c189,191
<         prefix = args.prefix if args.prefix else args.padding_text
<         encoded_prompt = tokenizer.encode(prefix + prompt_text, add_special_tokens=False, return_tensors="pt")
<     encoded_prompt = encoded_prompt.to(args.device)
---
>         model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path)
>         model = model.eval().to(args.device)
>         is_bloom = model_is_bloom(model.config)
379,382c193,194
<     if encoded_prompt.size()[-1] == 0:
<         input_ids = None
<     else:
<         input_ids = encoded_prompt
---
>         if args.use_hpu_graphs:
>             from habana_frameworks.torch.hpu import wrap_in_hpu_graph
384,405c196
<     if args.jit:
<         jit_input_texts = ["jit"]
<         jit_inputs = prepare_jit_inputs(jit_input_texts, model, tokenizer)
<         torch._C._jit_set_texpr_fuser_enabled(False)
<         model.config.return_dict = False
<         traced_model = torch.jit.trace(model, jit_inputs, strict=False)
<         traced_model = torch.jit.freeze(traced_model.eval())
<         traced_model(*jit_inputs)
<         traced_model(*jit_inputs)
< 
<         model = _ModelFallbackWrapper(traced_model, model)
< 
<     output_sequences = model.generate(
<         input_ids=input_ids,
<         max_length=args.length + len(encoded_prompt[0]),
<         temperature=args.temperature,
<         top_k=args.k,
<         top_p=args.p,
<         repetition_penalty=args.repetition_penalty,
<         do_sample=True,
<         num_return_sequences=args.num_return_sequences,
<     )
---
>             model = wrap_in_hpu_graph(model)
407,409c198,201
<     # Remove the batch dimension when returning multiple sequences
<     if len(output_sequences.shape) > 2:
<         output_sequences.squeeze_()
---
>     # Some models like GPT2 do not have a PAD token so we have to set it if necessary
>     if tokenizer.pad_token is None:
>         tokenizer.pad_token = tokenizer.eos_token
>         model.generation_config.pad_token_id = model.generation_config.eos_token_id
411c203,204
<     generated_sequences = []
---
>     if rank in [-1, 0]:
>         logger.info(f"Args: {args}")
413,415c206,215
<     for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
<         print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===")
<         generated_sequence = generated_sequence.tolist()
---
>         use_bf16 = False
>         if use_deepspeed or (gaudi_config is not None and gaudi_config.use_habana_mixed_precision):
>             use_bf16 = True
>         logger.info(f"device: {args.device}, n_hpu: {world_size}, bf16: {use_bf16}")
> 
>     # Generation configuration
>     generation_config = GenerationConfig(
>         max_new_tokens=args.max_new_tokens,
>         use_cache=args.use_kv_cache,
>     )
417,418c217,320
<         # Decode text
<         text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
---
>     if args.dataset_name is None:
>         # Benchmark over the prompts below
>         input_sentences = [
>             "DeepSpeed is a machine learning framework",
>             "He is working on",
>             "He has a",
>             "He got all",
>             "Everyone is happy and I can",
>             "The new movie that got Oscar this year",
>             "In the far far distance from our galaxy,",
>             "Peace is the only way",
>         ]
> 
>         if args.batch_size > len(input_sentences):
>             # Dynamically extends to support larger batch sizes
>             num_sentences_to_add = args.batch_size - len(input_sentences)
>             for i in range(num_sentences_to_add):
>                 input_sentences.append(input_sentences[i % len(input_sentences)])
>         elif args.batch_size < len(input_sentences):
>             input_sentences = input_sentences[: args.batch_size]
> 
>         def generate():
>             """Generates sequences from the input sentences and returns them."""
> 
>             # Tokenization
>             input_tokens = tokenizer.batch_encode_plus(input_sentences, return_tensors="pt", padding=True)
> 
>             # Pad inputs to have static shapes during generation, this gives better performance than dynamic shapes on HPUs
>             input_token_len = input_tokens.input_ids.shape[-1]
>             input_tokens["input_ids"] = F.pad(
>                 input_tokens.input_ids, (0, args.max_new_tokens), value=model.config.pad_token_id
>             )
>             input_tokens["attention_mask"] = F.pad(input_tokens.attention_mask, (0, args.max_new_tokens), value=0)
>             if is_bloom:
>                 # token_idx is the current index in the generation process, it is incremented each time a new token is generated
>                 kwargs = {"token_idx": torch.tensor(input_token_len, device=args.device)}
>             else:
>                 kwargs = {}
> 
>             # Move inputs to target device(s)
>             for t in input_tokens:
>                 if torch.is_tensor(input_tokens[t]):
>                     input_tokens[t] = input_tokens[t].to(args.device)
> 
>             outputs = model.generate(
>                 **input_tokens,
>                 **kwargs,
>                 generation_config=generation_config,
>                 lazy_mode=True,
>                 hpu_graphs=args.use_hpu_graphs,
>             ).cpu()
>             return tokenizer.batch_decode(outputs, skip_special_tokens=True)
> 
>         # Compilation
>         if rank in [-1, 0]:
>             logger.info("Graph compilation...")
>         t0 = time.perf_counter()
>         # The first three iterations take longer because of graph compilation
>         for _ in range(3):
>             generate()
>         torch_hpu.synchronize()
>         compilation_duration = time.perf_counter() - t0
> 
>         total_new_tokens_generated = 0
>         if rank in [-1, 0]:
>             logger.info("Running generate...")
>         t0 = time.perf_counter()
>         # Benchmark over n_iterations iterations
>         for i in range(args.n_iterations):
>             generated = generate()
>         duration = time.perf_counter() - t0
>         total_new_tokens_generated = args.n_iterations * args.batch_size * args.max_new_tokens
>         throughput = total_new_tokens_generated / duration
> 
>         if rank in [-1, 0]:
>             stats = f"Throughput (including tokenization) = {throughput} tokens/second"
>             separator = "-" * len(stats)
>             print()
>             print("Stats:")
>             print(separator)
>             print(stats)
>             if args.use_hpu_graphs:
>                 print(f"Graph compilation duration = {compilation_duration} seconds")
>             print(separator)
>             print()
>             print("Input/outputs:")
>             print(separator)
>             for i, (input_sentence, output) in enumerate(zip(input_sentences, generated)):
>                 print(f"input {i+1}: {input_sentence}")
>                 print(f"output {i+1}: {output}")
>                 print(separator)
>     else:
>         # Downloading and loading a dataset from the hub.
>         from datasets import load_dataset
>         from torch.utils.data import DataLoader
> 
>         raw_dataset = load_dataset(args.dataset_name)
>         if "test" in raw_dataset:
>             split = "test"
>         elif "validation" in raw_dataset:
>             split = "validation"
>         else:
>             split = "train"
>         raw_dataset = raw_dataset[split]
420,421c322,332
<         # Remove all text after the stop token
<         text = text[: text.find(args.stop_token) if args.stop_token else None]
---
>         if args.column_name is None:
>             # If no column name is given, take the first column that has strings
>             column_name = [key for key in raw_dataset.features.keys() if raw_dataset.features[key].dtype == "string"][
>                 0
>             ]
>             if rank in [-1, 0]:
>                 logger.info(
>                     f"No column name was given so automatically choosing '{column_name}' for prompts. If you would like to use another column of the dataset, you can set the argument `--column_name`."
>                 )
>         else:
>             column_name = args.column_name
423,426c334,335
<         # 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)) :]
<         )
---
>         # Remove unused columns
>         raw_dataset = raw_dataset.remove_columns([name for name in raw_dataset.column_names if name != column_name])
428,429c337,338
<         generated_sequences.append(total_sequence)
<         print(total_sequence)
---
>         # Set the prompt length to 16
>         prompt_length = 16
431c340,387
<     return generated_sequences
---
>         def preprocess_function(examples):
>             # Tokenize the texts
>             return tokenizer(examples[column_name], padding="max_length", max_length=prompt_length, truncation=True)
> 
>         raw_dataset = raw_dataset.map(
>             preprocess_function,
>             batched=True,
>             desc="Running tokenizer on dataset",
>         )
>         # After tokenization, we can remove the column of interest
>         raw_dataset = raw_dataset.remove_columns([column_name])
>         raw_dataset.set_format(type="torch")
> 
>         separator = None
> 
>         logger.info("Running generation...")
> 
>         dataloader = DataLoader(raw_dataset, batch_size=args.batch_size)
>         for i, batch in enumerate(dataloader):
>             prompt = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True)
>             # Pad inputs to have static shapes during generation, this gives better performance than dynamic shapes on HPUs
>             batch["input_ids"] = F.pad(batch["input_ids"], (0, args.max_new_tokens), value=model.config.pad_token_id)
>             batch["attention_mask"] = F.pad(batch["attention_mask"], (0, args.max_new_tokens), value=0)
>             # prompt = batch.pop(column_name)
>             # Move inputs to target device(s)
>             for t in batch:
>                 if torch.is_tensor(batch[t]):
>                     batch[t] = batch[t].to(args.device)
>             if is_bloom:
>                 # token_idx is the current index in the generation process, it is incremented each time a new token is generated
>                 batch["token_idx"] = torch.tensor(prompt_length, device=args.device)
> 
>             # Generate new sequences
>             outputs = model.generate(
>                 **batch,
>                 generation_config=generation_config,
>                 lazy_mode=args.use_hpu_graphs,
>                 hpu_graphs=args.use_hpu_graphs,
>             ).cpu()
> 
>             # Print outputs
>             if separator is None:
>                 separator = "-" * len(prompt[0])
>             if rank in [-1, 0]:
>                 print(separator)
>                 print(f"Batch n°{i+1}")
>                 print(f"Input: {prompt[:args.batch_size]}")
>                 print(f"Output: {tokenizer.batch_decode(outputs, skip_special_tokens=True)[:args.batch_size]}")
