




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


_SCREAMING_SNAKE_CASE							:     Optional[int]   	=				{
    "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"},
}

_SCREAMING_SNAKE_CASE							:     List[str]   	=				[
    "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   UpperCamelCase_(			):

							'''simple docstring'''






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

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

							snake_case_			      =   compute_module_sizes(snake_case  )
							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**2_0}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(snake_case							,       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(snake_case  ) if outputs[: len(snake_case  )] != tokens else len(snake_case  ) - len(snake_case  )
														snake_case_			      =   after_generate - before_generate

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

							snake_case_			      =   end_measure(snake_case  )
							log_measures(snake_case							,       "Model generation"  )

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


if __name__ == "__main__":
		main()










'''simple docstring'''





import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
		import torch_xla.core.xla_model as xm


def   UpperCamelCase_(			snake_case							: int							,       snake_case							: bool = False  ):

							'''simple docstring'''






							if device_specific:
														seed += AcceleratorState().process_index
							random.seed(snake_case  )
							np.random.seed(snake_case  )
							torch.manual_seed(snake_case  )
							if is_xpu_available():
														torch.xpu.manual_seed_all(snake_case  )
							elif is_npu_available():
														torch.npu.manual_seed_all(snake_case  )
							else:
														torch.cuda.manual_seed_all(snake_case  )
							# ^^ safe to call this function even if cuda is not available
							if is_tpu_available():
														xm.set_rng_state(snake_case  )


def   UpperCamelCase_(			snake_case							: Optional[RNGType] = None							,       snake_case							: Optional[torch.Generator] = None  ):

							'''simple docstring'''






							if rng_type == RNGType.TORCH:
														snake_case_			      =   torch.get_rng_state()
							elif rng_type == RNGType.CUDA:
														snake_case_			      =   torch.cuda.get_rng_state()
							elif rng_type == RNGType.XLA:
														assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
														snake_case_			      =   torch.tensor(xm.get_rng_state()  )
							elif rng_type == RNGType.NPU:
														assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
														snake_case_			      =   torch.npu.get_rng_state()
							elif rng_type == RNGType.XPU:
														assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
														snake_case_			      =   torch.xpu.get_rng_state()
							elif rng_type == RNGType.GENERATOR:
														assert generator is not None, "Need a generator to synchronize its seed."
														snake_case_			      =   generator.get_state()

							# Broadcast the rng state from device 0 to other devices
							snake_case_			      =   AcceleratorState()
							if state.distributed_type == DistributedType.TPU:
														snake_case_			      =   rng_state.to(xm.xla_device()  )
														xm.collective_broadcast([rng_state]  )
														xm.mark_step()
														snake_case_			      =   rng_state.cpu()
							elif (
							    state.distributed_type in CUDA_DISTRIBUTED_TYPES
							    or state.distributed_type == DistributedType.MULTI_NPU
							    or state.distributed_type == DistributedType.MULTI_XPU
							):
														snake_case_			      =   rng_state.to(state.device  )
														torch.distributed.broadcast(snake_case							,       0  )
														snake_case_			      =   rng_state.cpu()
							elif state.distributed_type == DistributedType.MULTI_CPU:
														torch.distributed.broadcast(snake_case							,       0  )

							# Set the broadcast rng state
							if rng_type == RNGType.TORCH:
														torch.set_rng_state(snake_case  )
							elif rng_type == RNGType.CUDA:
														torch.cuda.set_rng_state(snake_case  )
							elif rng_type == RNGType.NPU:
														torch.npu.set_rng_state(snake_case  )
							elif rng_type == RNGType.XPU:
														torch.xpu.set_rng_state(snake_case  )
							elif rng_type == RNGType.XLA:
														xm.set_rng_state(rng_state.item()  )
							elif rng_type == RNGType.GENERATOR:
														generator.set_state(snake_case  )




def   UpperCamelCase_(			snake_case							: List[Union[str, RNGType]]							,       snake_case							: Optional[torch.Generator] = None  ):

							'''simple docstring'''






							for rng_type in rng_types:
														synchronize_rng_state(RNGType(snake_case  )							,       generator=snake_case  )




