




'''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 	_lowercase	(			):



		'''simple docstring'''







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

		_snake_case			 = parser.parse_args()

		# Sanitize args
		if args.model_name in DEFAULT_MODELS:
				_snake_case			 = DEFAULT_MODELS[args.model_name]
				_snake_case			 = defaults["model"]
				if args.tokenizer_name is None:
						_snake_case			 = defaults.get("tokenizer"		,							args.model_name						)
				if args.is_causal is None:
						_snake_case			 = defaults["is_causal"]
				if args.model_revision is None:
						_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:
				_snake_case			 = args.model_name
		if args.model_revision is None:
				_snake_case			 = "main"

		return args






def 	_lowercase	(			):



		'''simple docstring'''







		transformers.utils.logging.set_verbosity_error()
		_snake_case			 = parse_args()

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

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

		_snake_case			 = compute_module_sizes(lowercase_						)
		_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]
		_snake_case			 = "\n".join([f'''- {device}: {size // 2**20}MiB''' for device, size in device_size.items()]						)
		print(f'''\nTheoretical use:\n{message}'''						)

		_snake_case			 = AutoTokenizer.from_pretrained(args.tokenizer_name						)

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

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

		_snake_case			 = end_measure(lowercase_						)
		log_measures(lowercase_		,							"Model generation"						)

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


if __name__ == "__main__":
						main()











'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCamelCase							:		Any										=						{'configuration_pegasus': ['PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'PegasusConfig']}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		Optional[Any]										=						['PegasusTokenizer']

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		Optional[Any]										=						['PegasusTokenizerFast']

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		List[str]										=						[
						    'PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'PegasusForCausalLM',
						    'PegasusForConditionalGeneration',
						    'PegasusModel',
						    'PegasusPreTrainedModel',
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		Union[str, Any]										=						[
						    'TFPegasusForConditionalGeneration',
						    'TFPegasusModel',
						    'TFPegasusPreTrainedModel',
						]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_lowerCamelCase							:		List[Any]										=						[
						    'FlaxPegasusForConditionalGeneration',
						    'FlaxPegasusModel',
						    'FlaxPegasusPreTrainedModel',
						]


if TYPE_CHECKING:
						from .configuration_pegasus import PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP, PegasusConfig

						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_pegasus import PegasusTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_pegasus_fast import PegasusTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_pegasus import (
												    PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
												    PegasusForCausalLM,
												    PegasusForConditionalGeneration,
												    PegasusModel,
												    PegasusPreTrainedModel,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_pegasus import TFPegasusForConditionalGeneration, TFPegasusModel, TFPegasusPreTrainedModel

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_flax_pegasus import (
												    FlaxPegasusForConditionalGeneration,
												    FlaxPegasusModel,
												    FlaxPegasusPreTrainedModel,
												)

else:
						import sys

						_lowerCamelCase							:		List[Any]										=						_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





