
"""simple docstring"""




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

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


snake_case    =		logging.get_logger(__name__)

snake_case    =		{
    'gpt2': 'https://huggingface.co/gpt2/resolve/main/config.json',
    'gpt2-medium': 'https://huggingface.co/gpt2-medium/resolve/main/config.json',
    'gpt2-large': 'https://huggingface.co/gpt2-large/resolve/main/config.json',
    'gpt2-xl': 'https://huggingface.co/gpt2-xl/resolve/main/config.json',
    'distilgpt2': 'https://huggingface.co/distilgpt2/resolve/main/config.json',
}




class      UpperCamelCase	(       __magic_name__      ):

				"""simple docstring"""



				UpperCAmelCase_				:      List[str]						  =				"gpt2"
				UpperCAmelCase_				:      Any						  =				["past_key_values"]
				UpperCAmelCase_				:      Any						  =				{
				    "hidden_size": "n_embd",
				    "max_position_embeddings": "n_positions",
				    "num_attention_heads": "n_head",
				    "num_hidden_layers": "n_layer",
				}


				def __init__(				self		,	lowercase__=50257		,	lowercase__=1024		,	lowercase__=768		,	lowercase__=12		,	lowercase__=12		,	lowercase__=None		,	lowercase__="gelu_new"		,	lowercase__=0.1		,	lowercase__=0.1		,	lowercase__=0.1		,	lowercase__=1E-5		,	lowercase__=0.02		,	lowercase__="cls_index"		,	lowercase__=True		,	lowercase__=None		,	lowercase__=True		,	lowercase__=0.1		,	lowercase__=True		,	lowercase__=True		,	lowercase__=50256		,	lowercase__=50256		,	lowercase__=False		,	lowercase__=False		,	**lowercase__		,	)    ->    Dict:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	vocab_size
									SCREAMING_SNAKE_CASE				=	n_positions
									SCREAMING_SNAKE_CASE				=	n_embd
									SCREAMING_SNAKE_CASE				=	n_layer
									SCREAMING_SNAKE_CASE				=	n_head
									SCREAMING_SNAKE_CASE				=	n_inner
									SCREAMING_SNAKE_CASE				=	activation_function
									SCREAMING_SNAKE_CASE				=	resid_pdrop
									SCREAMING_SNAKE_CASE				=	embd_pdrop
									SCREAMING_SNAKE_CASE				=	attn_pdrop
									SCREAMING_SNAKE_CASE				=	layer_norm_epsilon
									SCREAMING_SNAKE_CASE				=	initializer_range
									SCREAMING_SNAKE_CASE				=	summary_type
									SCREAMING_SNAKE_CASE				=	summary_use_proj
									SCREAMING_SNAKE_CASE				=	summary_activation
									SCREAMING_SNAKE_CASE				=	summary_first_dropout
									SCREAMING_SNAKE_CASE				=	summary_proj_to_labels
									SCREAMING_SNAKE_CASE				=	scale_attn_weights
									SCREAMING_SNAKE_CASE				=	use_cache
									SCREAMING_SNAKE_CASE				=	scale_attn_by_inverse_layer_idx
									SCREAMING_SNAKE_CASE				=	reorder_and_upcast_attn

									SCREAMING_SNAKE_CASE				=	bos_token_id
									SCREAMING_SNAKE_CASE				=	eos_token_id

									super().__init__(bos_token_id=lowercase__		,	eos_token_id=lowercase__		,	**lowercase__			)






class      UpperCamelCase	(       __magic_name__      ):

				"""simple docstring"""



				def __init__(				self		,	lowercase__		,	lowercase__ = "default"		,	lowercase__ = None		,	lowercase__ = False		,	)    ->    Union[str, Any]:
									"""simple docstring"""







									super().__init__(lowercase__		,	task=lowercase__		,	patching_specs=lowercase__		,	use_past=lowercase__			)
									if not getattr(self._config		,	'pad_token_id'		,	lowercase__			):
														# TODO: how to do that better?
														SCREAMING_SNAKE_CASE				=	0


				@property
				def 							A					(				self			)    ->    Mapping[str, Mapping[int, str]]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	OrderedDict({'input_ids': {0: 'batch', 1: 'sequence'}}			)
									if self.use_past:
														self.fill_with_past_key_values_(lowercase__		,	direction='inputs'			)
														SCREAMING_SNAKE_CASE				=	{0: 'batch', 1: 'past_sequence + sequence'}
									else:
														SCREAMING_SNAKE_CASE				=	{0: 'batch', 1: 'sequence'}

									return common_inputs


				@property
				def 							A					(				self			)    ->    int:
									"""simple docstring"""







									return self._config.n_layer


				@property
				def 							A					(				self			)    ->    int:
									"""simple docstring"""







									return self._config.n_head


				def 							A					(				self		,	lowercase__		,	lowercase__ = -1		,	lowercase__ = -1		,	lowercase__ = False		,	lowercase__ = None		,	)    ->    Mapping[str, Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	super(lowercase__		,	self			).generate_dummy_inputs(
									    lowercase__		,	batch_size=lowercase__		,	seq_length=lowercase__		,	is_pair=lowercase__		,	framework=lowercase__			)

									# We need to order the input in the way they appears in the forward()
									SCREAMING_SNAKE_CASE				=	OrderedDict({'input_ids': common_inputs['input_ids']}			)

									# Need to add the past_keys
									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
																			# Not using the same length for past_key_values
																			SCREAMING_SNAKE_CASE				=	seqlen + 2
																			SCREAMING_SNAKE_CASE				=	(
																			    batch,
																			    self.num_attention_heads,
																			    past_key_values_length,
																			    self._config.hidden_size // self.num_attention_heads,
																			)
																			SCREAMING_SNAKE_CASE				=	[
																			    (torch.zeros(lowercase__			), torch.zeros(lowercase__			)) for _ in range(self.num_layers			)
																			]

									SCREAMING_SNAKE_CASE				=	common_inputs['attention_mask']
									if self.use_past:
														SCREAMING_SNAKE_CASE				=	ordered_inputs['attention_mask'].dtype
														SCREAMING_SNAKE_CASE				=	torch.cat(
														    [ordered_inputs['attention_mask'], torch.ones(lowercase__		,	lowercase__		,	dtype=lowercase__			)]		,	dim=1			)

									return ordered_inputs




				@property
				def 							A					(				self			)    ->    int:
									"""simple docstring"""







									return 13







"""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 SCREAMING_SNAKE_CASE_   (       snake_case   :					int						,  snake_case   :					bool = False		)->		Any:
       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 SCREAMING_SNAKE_CASE_   (       snake_case   :					Optional[RNGType] = None						,  snake_case   :					Optional[torch.Generator] = None		)->		str:
       # Get the proper rng state
       if rng_type == RNGType.TORCH:
              _lowerCamelCase    	=					torch.get_rng_state()
       elif rng_type == RNGType.CUDA:
              _lowerCamelCase    	=					torch.cuda.get_rng_state()
       elif rng_type == RNGType.XLA:
              assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
              _lowerCamelCase    	=					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."
              _lowerCamelCase    	=					torch.npu.get_rng_state()
       elif rng_type == RNGType.XPU:
              assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
              _lowerCamelCase    	=					torch.xpu.get_rng_state()
       elif rng_type == RNGType.GENERATOR:
              assert generator is not None, "Need a generator to synchronize its seed."
              _lowerCamelCase    	=					generator.get_state()

       # Broadcast the rng state from device 0 to other devices
       _lowerCamelCase    	=					AcceleratorState()
       if state.distributed_type == DistributedType.TPU:
              _lowerCamelCase    	=					rng_state.to(xm.xla_device()		)
              xm.collective_broadcast([rng_state]		)
              xm.mark_step()
              _lowerCamelCase    	=					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
       ):
              _lowerCamelCase    	=					rng_state.to(state.device		)
              torch.distributed.broadcast(snake_case						,  0		)
              _lowerCamelCase    	=					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 SCREAMING_SNAKE_CASE_   (       snake_case   :					List[Union[str, RNGType]]						,  snake_case   :					Optional[torch.Generator] = None		)->		int:
       for rng_type in rng_types:
              synchronize_rng_state(RNGType(snake_case		)						,  generator=snake_case		)






