



"""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


lowerCamelCase				:     Union[str, Any]								=  {
    """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"""},
}

lowerCamelCase				:     Union[str, Any]								=  [
    """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 							A__ (     ):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			argparse.ArgumentParser(description='''Run and time generations on a big model using Accelerate.'''    )
  parser.add_argument('''model_name'''       ,							type=UpperCamelCase__       ,							default=UpperCamelCase__       ,							help='''The name of the model to try.'''    )
  parser.add_argument(
      '''--tokenizer_name'''       ,							type=UpperCamelCase__       ,							default=UpperCamelCase__       ,							help='''The name of the tokenizer (if different from the model.'''    )
  parser.add_argument('''--is_causal'''       ,							type=UpperCamelCase__       ,							default=UpperCamelCase__       ,							help='''Whether or not the model is causal.'''    )
  parser.add_argument(
      '''--model_revision'''       ,							type=UpperCamelCase__       ,							default=UpperCamelCase__       ,							help='''The revision to use for the model checkpoint.'''    )
  parser.add_argument('''--torch_dtype'''       ,							type=UpperCamelCase__       ,							default=UpperCamelCase__       ,							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 							A__ (     ):


  '''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(UpperCamelCase__       ,							'''torch_dtype'''       ,							torch.floataa    )
  else:
    _SCREAMING_SNAKE_CASE        =			getattr(UpperCamelCase__       ,							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'''       ,							**UpperCamelCase__    )
  _SCREAMING_SNAKE_CASE        =			end_measure(UpperCamelCase__    )
  log_measures(UpperCamelCase__       ,							'''Model loading'''    )

  _SCREAMING_SNAKE_CASE        =			compute_module_sizes(UpperCamelCase__    )
  _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(UpperCamelCase__       ,							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(UpperCamelCase__    ) if outputs[: len(UpperCamelCase__    )] != tokens else len(UpperCamelCase__    ) - len(UpperCamelCase__    )
    _SCREAMING_SNAKE_CASE        =			after_generate - before_generate

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

  _SCREAMING_SNAKE_CASE        =			end_measure(UpperCamelCase__    )
  log_measures(UpperCamelCase__       ,							'''Model generation'''    )

  _SCREAMING_SNAKE_CASE        =			[gen / tok for gen, tok in zip(UpperCamelCase__       ,							UpperCamelCase__    )]
  _SCREAMING_SNAKE_CASE        =			sum(UpperCamelCase__    ) / len(UpperCamelCase__    )
  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(UpperCamelCase__    ) - 1)
  print(F'''Average time of generation per token (excluding the first): {avg_gen:.2f}s'''    )


if __name__ == "__main__":
    main()





"""simple docstring"""






from transformers import ProcessorMixin




class    __snake_case(					__A       ):
    _A       = '''AutoFeatureExtractor'''
    _A       = '''AutoTokenizer'''
