
'''simple docstring'''



from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


UpperCamelCase__ :						str    							=					logging.get_logger(__name__)

UpperCamelCase__ :						Dict    							=					{
    "facebook/m2m100_418M": "https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json",
    # See all M2M100 models at https://huggingface.co/models?filter=m2m_100
}

class 				_a			(_lowerCamelCase):





       """simple docstring"""




       SCREAMING_SNAKE_CASE =					'm2m_100'
       SCREAMING_SNAKE_CASE =					['past_key_values']
       SCREAMING_SNAKE_CASE =					{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}






       def __init__(							self    ,							A__=12_81_12    ,							A__=10_24    ,							A__=12    ,							A__=40_96    ,							A__=16    ,							A__=12    ,							A__=40_96    ,							A__=16    ,							A__=0.05    ,							A__=0.05    ,							A__=True    ,							A__=True    ,							A__="relu"    ,							A__=10_24    ,							A__=0.1    ,							A__=0.1    ,							A__=0.0    ,							A__=0.02    ,							A__=2    ,							A__=True    ,							A__=1    ,							A__=0    ,							A__=2    ,							**A__    ,							)      ->     Tuple:
          _SCREAMING_SNAKE_CASE       	=			vocab_size
          _SCREAMING_SNAKE_CASE       	=			max_position_embeddings
          _SCREAMING_SNAKE_CASE       	=			d_model
          _SCREAMING_SNAKE_CASE       	=			encoder_ffn_dim
          _SCREAMING_SNAKE_CASE       	=			encoder_layers
          _SCREAMING_SNAKE_CASE       	=			encoder_attention_heads
          _SCREAMING_SNAKE_CASE       	=			decoder_ffn_dim
          _SCREAMING_SNAKE_CASE       	=			decoder_layers
          _SCREAMING_SNAKE_CASE       	=			decoder_attention_heads
          _SCREAMING_SNAKE_CASE       	=			dropout
          _SCREAMING_SNAKE_CASE       	=			attention_dropout
          _SCREAMING_SNAKE_CASE       	=			activation_dropout
          _SCREAMING_SNAKE_CASE       	=			activation_function
          _SCREAMING_SNAKE_CASE       	=			init_std
          _SCREAMING_SNAKE_CASE       	=			encoder_layerdrop
          _SCREAMING_SNAKE_CASE       	=			decoder_layerdrop
          _SCREAMING_SNAKE_CASE       	=			use_cache
          _SCREAMING_SNAKE_CASE       	=			encoder_layers
          _SCREAMING_SNAKE_CASE       	=			scale_embedding  # scale factor will be sqrt(d_model) if True

          super().__init__(
              pad_token_id=A__    ,							bos_token_id=A__    ,							eos_token_id=A__    ,							is_encoder_decoder=A__    ,							decoder_start_token_id=A__    ,							**A__    ,							)






class 				_a			(_lowerCamelCase):





       """simple docstring"""




       @property
       def     UpperCamelCase			(							self )      ->     Mapping[str, Mapping[int, str]]:
          _SCREAMING_SNAKE_CASE       	=			OrderedDict(
              [
                  ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                  ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
              ] )

          if self.use_past:
             _SCREAMING_SNAKE_CASE       	=			{0: """batch"""}
             _SCREAMING_SNAKE_CASE       	=			{0: """batch""", 1: """past_decoder_sequence + sequence"""}
          else:
             _SCREAMING_SNAKE_CASE       	=			{0: """batch""", 1: """decoder_sequence"""}
             _SCREAMING_SNAKE_CASE       	=			{0: """batch""", 1: """decoder_sequence"""}

          if self.use_past:
             self.fill_with_past_key_values_(A__    ,							direction="""inputs""" )
          return common_inputs






       def     UpperCamelCase			(							self    ,							A__    ,							A__ = -1    ,							A__ = -1    ,							A__ = False    ,							A__ = None    ,							)      ->     Mapping[str, Any]:
          # Copied from OnnxConfig.generate_dummy_inputs
          # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
          # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
          _SCREAMING_SNAKE_CASE       	=			compute_effective_axis_dimension(
              A__    ,							fixed_dimension=OnnxConfig.default_fixed_batch    ,							num_token_to_add=0 )

          # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
          _SCREAMING_SNAKE_CASE       	=			tokenizer.num_special_tokens_to_add(A__ )
          _SCREAMING_SNAKE_CASE       	=			compute_effective_axis_dimension(
              A__    ,							fixed_dimension=OnnxConfig.default_fixed_sequence    ,							num_token_to_add=A__ )

          # Generate dummy inputs according to compute batch and sequence
          _SCREAMING_SNAKE_CASE       	=			[""" """.join([tokenizer.unk_token] ) * seq_length] * batch_size
          _SCREAMING_SNAKE_CASE       	=			dict(tokenizer(A__    ,							return_tensors=A__ ) )
          return common_inputs






       def     UpperCamelCase			(							self    ,							A__    ,							A__ = -1    ,							A__ = -1    ,							A__ = False    ,							A__ = None    ,							)      ->     Mapping[str, Any]:
          _SCREAMING_SNAKE_CASE       	=			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              A__    ,							A__    ,							A__    ,							A__    ,							A__ )

          # Generate decoder inputs
          _SCREAMING_SNAKE_CASE       	=			seq_length if not self.use_past else 1
          _SCREAMING_SNAKE_CASE       	=			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
              A__    ,							A__    ,							A__    ,							A__    ,							A__ )
          _SCREAMING_SNAKE_CASE       	=			{F"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
          _SCREAMING_SNAKE_CASE       	=			dict(**A__    ,							**A__ )

          if self.use_past:
             if not is_torch_available():
                raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed.""" )
             else:
                import torch
             _SCREAMING_SNAKE_CASE	,							_SCREAMING_SNAKE_CASE       	=			common_inputs["""input_ids"""].shape
             _SCREAMING_SNAKE_CASE       	=			common_inputs["""decoder_input_ids"""].shape[1]
             _SCREAMING_SNAKE_CASE	,							_SCREAMING_SNAKE_CASE       	=			self.num_attention_heads
             _SCREAMING_SNAKE_CASE       	=			(
                 batch,
                 num_encoder_attention_heads,
                 encoder_seq_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             _SCREAMING_SNAKE_CASE       	=			decoder_seq_length + 3
             _SCREAMING_SNAKE_CASE       	=			(
                 batch,
                 num_decoder_attention_heads,
                 decoder_past_length,
                 self._config.hidden_size // num_decoder_attention_heads,
             )

             _SCREAMING_SNAKE_CASE       	=			torch.cat(
                 [common_inputs["""decoder_attention_mask"""], torch.ones(A__    ,							A__ )]    ,							dim=1 )

             _SCREAMING_SNAKE_CASE       	=			[]
             # If the number of encoder and decoder layers are present in the model configuration, both are considered
             _SCREAMING_SNAKE_CASE	,							_SCREAMING_SNAKE_CASE       	=			self.num_layers
             _SCREAMING_SNAKE_CASE       	=			min(A__    ,							A__ )
             _SCREAMING_SNAKE_CASE       	=			max(A__    ,							A__ ) - min_num_layers
             _SCREAMING_SNAKE_CASE       	=			"""encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

             for _ in range(A__ ):
                common_inputs["past_key_values"].append(
                    (
                        torch.zeros(A__ ),
                        torch.zeros(A__ ),
                        torch.zeros(A__ ),
                        torch.zeros(A__ ),
                    ) )
             # TODO: test this.
             _SCREAMING_SNAKE_CASE       	=			encoder_shape if remaining_side_name == """encoder""" else decoder_shape
             for _ in range(A__    ,							A__ ):
                common_inputs["past_key_values"].append((torch.zeros(A__ ), torch.zeros(A__ )) )
          return common_inputs


       SCREAMING_SNAKE_CASE =					_generate_dummy_inputs_for_default_and_seqaseq_lm








'''simple docstring'''



import argparse
import time

import torch
import transformers
from measures_util import end_measure, log_measures, start_measure
from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeqaSeqLM, AutoTokenizer

from accelerate.utils import compute_module_sizes


UpperCamelCase__ :						Tuple    							=					{
    "gpt-j-6b": {"is_causal": True, "model": "sgugger/sharded-gpt-j-6B", "tokenizer": "EleutherAI/gpt-j-6B"},
    "gpt-neox": {"is_causal": True, "model": "EleutherAI/gpt-neox-20b"},
    "opt": {"is_causal": True, "model": "facebook/opt-30b"},
    "T0pp": {"is_causal": False, "model": "bigscience/T0pp", "model_revision": "sharded"},
}

UpperCamelCase__ :						int    							=					[
    "Hello, my name is",
    "Are unicorns real? Unicorns are",
    "For the first time in several years,",
    "My name is Julien and I am",
    "The goal of life is",
    "Whenever I'm sad, I like to",
]
def 				lowerCAmelCase_		(			)       ->       str:




   """simple docstring"""
   _SCREAMING_SNAKE_CASE       	=			argparse.ArgumentParser(description="""Run and time generations on a big model using Accelerate."""  )
   parser.add_argument("""model_name"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The name of the model to try."""  )
   parser.add_argument(
       """--tokenizer_name"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The name of the tokenizer (if different from the model."""  )
   parser.add_argument("""--is_causal"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""Whether or not the model is causal."""  )
   parser.add_argument(
       """--model_revision"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The revision to use for the model checkpoint."""  )
   parser.add_argument("""--torch_dtype"""      , type=SCREAMING_SNAKE_CASE_      , default=SCREAMING_SNAKE_CASE_      , help="""The dtype for the model."""  )
   parser.add_argument("""--disk_offload"""      , action="""store_true"""  )

   _SCREAMING_SNAKE_CASE       	=			parser.parse_args()

   # Sanitize args
   if args.model_name in DEFAULT_MODELS:
      _SCREAMING_SNAKE_CASE       	=			DEFAULT_MODELS[args.model_name]
      _SCREAMING_SNAKE_CASE       	=			defaults["""model"""]
      if args.tokenizer_name is None:
         _SCREAMING_SNAKE_CASE       	=			defaults.get("""tokenizer"""      , args.model_name  )
      if args.is_causal is None:
         _SCREAMING_SNAKE_CASE       	=			defaults["""is_causal"""]
      if args.model_revision is None:
         _SCREAMING_SNAKE_CASE       	=			defaults.get("""model_revision"""      , """main"""  )

   if args.is_causal is None:
      raise ValueError("""Could not infer the default for `--is_causal`, pass either True or False for it."""  )
   if args.tokenizer_name is None:
      _SCREAMING_SNAKE_CASE       	=			args.model_name
   if args.model_revision is None:
      _SCREAMING_SNAKE_CASE       	=			"""main"""

   return args
def 				lowerCAmelCase_		(			)       ->       Optional[int]:




   """simple docstring"""
   transformers.utils.logging.set_verbosity_error()
   _SCREAMING_SNAKE_CASE       	=			parse_args()

   if args.torch_dtype is None:
      _SCREAMING_SNAKE_CASE       	=			AutoConfig.from_pretrained(args.model_name  )
      _SCREAMING_SNAKE_CASE       	=			getattr(SCREAMING_SNAKE_CASE_      , """torch_dtype"""      , torch.floataa  )
   else:
      _SCREAMING_SNAKE_CASE       	=			getattr(SCREAMING_SNAKE_CASE_      , args.torch_dtype  )
   _SCREAMING_SNAKE_CASE       	=			AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
   _SCREAMING_SNAKE_CASE       	=			{
       """torch_dtype""": torch_dtype,
       """revision""": args.model_revision,
   }
   if args.disk_offload:
      _SCREAMING_SNAKE_CASE       	=			"""tmp_offload"""
      _SCREAMING_SNAKE_CASE       	=			True

   _SCREAMING_SNAKE_CASE       	=			start_measure()
   _SCREAMING_SNAKE_CASE       	=			model_cls.from_pretrained(args.model_name      , device_map="""auto"""      , **SCREAMING_SNAKE_CASE_  )
   _SCREAMING_SNAKE_CASE       	=			end_measure(SCREAMING_SNAKE_CASE_  )
   log_measures(SCREAMING_SNAKE_CASE_      , """Model loading"""  )

   _SCREAMING_SNAKE_CASE       	=			compute_module_sizes(SCREAMING_SNAKE_CASE_  )
   _SCREAMING_SNAKE_CASE       	=			{v: 0 for v in model.hf_device_map.values()}
   for module, device in model.hf_device_map.items():
      device_size[device] += module_sizes[module]
   _SCREAMING_SNAKE_CASE       	=			"""\n""".join([F"- {device}: {size // 2**20}MiB" for device, size in device_size.items()]  )
   print(F"\nTheoretical use:\n{message}"  )

   _SCREAMING_SNAKE_CASE       	=			AutoTokenizer.from_pretrained(args.tokenizer_name  )

   _SCREAMING_SNAKE_CASE       	=			start_measure()
   _SCREAMING_SNAKE_CASE       	=			[]
   _SCREAMING_SNAKE_CASE       	=			[]
   _SCREAMING_SNAKE_CASE       	=			[]
   for prompt in PROMPTS:
      _SCREAMING_SNAKE_CASE       	=			tokenizer(SCREAMING_SNAKE_CASE_      , return_tensors="""pt"""  ).to(0  )
      _SCREAMING_SNAKE_CASE       	=			inputs["""input_ids"""][0].tolist()
      _SCREAMING_SNAKE_CASE       	=			time.time()
      _SCREAMING_SNAKE_CASE       	=			model.generate(inputs["""input_ids"""]  )
      _SCREAMING_SNAKE_CASE       	=			time.time()
      _SCREAMING_SNAKE_CASE       	=			outputs[0].tolist()
      _SCREAMING_SNAKE_CASE       	=			len(SCREAMING_SNAKE_CASE_  ) if outputs[: len(SCREAMING_SNAKE_CASE_  )] != tokens else len(SCREAMING_SNAKE_CASE_  ) - len(SCREAMING_SNAKE_CASE_  )
      _SCREAMING_SNAKE_CASE       	=			after_generate - before_generate

      _SCREAMING_SNAKE_CASE       	=			tokenizer.decode(SCREAMING_SNAKE_CASE_      , skip_special_tokens=SCREAMING_SNAKE_CASE_  )
      texts_outs.append(SCREAMING_SNAKE_CASE_  )
      generation_times.append(SCREAMING_SNAKE_CASE_  )
      gen_tokens.append(SCREAMING_SNAKE_CASE_  )
      print(F"Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n"  )

   _SCREAMING_SNAKE_CASE       	=			end_measure(SCREAMING_SNAKE_CASE_  )
   log_measures(SCREAMING_SNAKE_CASE_      , """Model generation"""  )

   _SCREAMING_SNAKE_CASE       	=			[gen / tok for gen, tok in zip(SCREAMING_SNAKE_CASE_      , SCREAMING_SNAKE_CASE_  )]
   _SCREAMING_SNAKE_CASE       	=			sum(SCREAMING_SNAKE_CASE_  ) / len(SCREAMING_SNAKE_CASE_  )
   print(F"Average time of generation per token: {avg_gen:.2f}s"  )
   print(F"First generation (avg time per token): {generation_times_per_token[0]:.2f}s"  )
   _SCREAMING_SNAKE_CASE       	=			sum(generation_times_per_token[1:]  ) / (len(SCREAMING_SNAKE_CASE_  ) - 1)
   print(F"Average time of generation per token (excluding the first): {avg_gen:.2f}s"  )


if __name__ == "__main__":
     main()






