





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


__UpperCamelCase:	Optional[int] = logging.get_logger(__name__)


__UpperCamelCase:	List[str] = {
    """Salesforce/codegen-350M-nl""": """https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json""",
    """Salesforce/codegen-350M-multi""": """https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json""",
    """Salesforce/codegen-350M-mono""": """https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json""",
    """Salesforce/codegen-2B-nl""": """https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json""",
    """Salesforce/codegen-2B-multi""": """https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json""",
    """Salesforce/codegen-2B-mono""": """https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json""",
    """Salesforce/codegen-6B-nl""": """https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json""",
    """Salesforce/codegen-6B-multi""": """https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json""",
    """Salesforce/codegen-6B-mono""": """https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json""",
    """Salesforce/codegen-16B-nl""": """https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json""",
    """Salesforce/codegen-16B-multi""": """https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json""",
    """Salesforce/codegen-16B-mono""": """https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json""",
}


class       __lowerCAmelCase					(       _UpperCamelCase      ):


	'''simple docstring'''





	_A													=				"codegen"
	_A													=				{
	    "max_position_embeddings": "n_positions",
	    "hidden_size": "n_embd",
	    "num_attention_heads": "n_head",
	    "num_hidden_layers": "n_layer",
	}

	def __init__(						self:						List[str],					lowerCamelCase_:						Union[str, Any]=50400,					lowerCamelCase_:						Optional[int]=2048,					lowerCamelCase_:						Dict=2048,					lowerCamelCase_:						List[Any]=4096,					lowerCamelCase_:						Tuple=28,					lowerCamelCase_:						Dict=16,					lowerCamelCase_:						str=64,					lowerCamelCase_:						List[str]=None,					lowerCamelCase_:						Tuple="gelu_new",					lowerCamelCase_:						Dict=0.0,					lowerCamelCase_:						List[str]=0.0,					lowerCamelCase_:						Dict=0.0,					lowerCamelCase_:						Tuple=1E-5,					lowerCamelCase_:						List[str]=0.0_2,					lowerCamelCase_:						Optional[int]=True,					lowerCamelCase_:						str=50256,					lowerCamelCase_:						Tuple=50256,					lowerCamelCase_:						Tuple=False,					**lowerCamelCase_:						Any,					):
			lowercase__					:      Dict      =		vocab_size
			lowercase__					:      int      =		n_ctx
			lowercase__					:      str      =		n_positions
			lowercase__					:      Union[str, Any]      =		n_embd
			lowercase__					:      Union[str, Any]      =		n_layer
			lowercase__					:      Optional[int]      =		n_head
			lowercase__					:      Dict      =		n_inner
			lowercase__					:      Optional[Any]      =		rotary_dim
			lowercase__					:      str      =		activation_function
			lowercase__					:      List[Any]      =		resid_pdrop
			lowercase__					:      str      =		embd_pdrop
			lowercase__					:      List[str]      =		attn_pdrop
			lowercase__					:      str      =		layer_norm_epsilon
			lowercase__					:      Any      =		initializer_range
			lowercase__					:      List[str]      =		use_cache

			lowercase__					:      Tuple      =		bos_token_id
			lowercase__					:      List[Any]      =		eos_token_id

			super().__init__(
			    bos_token_id=__UpperCAmelCase,					eos_token_id=__UpperCAmelCase,					tie_word_embeddings=__UpperCAmelCase,					**__UpperCAmelCase			)


class       __lowerCAmelCase					(       _UpperCamelCase      ):


	'''simple docstring'''





	def __init__(						self:						Optional[int],					lowerCamelCase_:						Dict,					lowerCamelCase_:						str = "default",					lowerCamelCase_:						Any = None,					lowerCamelCase_:						List[Any] = False,					):
			super().__init__(__UpperCAmelCase,					task=__UpperCAmelCase,					patching_specs=__UpperCAmelCase,					use_past=__UpperCAmelCase			)
			if not getattr(self._config,					'pad_token_id',					__UpperCAmelCase			):
					# TODO: how to do that better?
					lowercase__					:      Union[str, Any]      =		0

	@property
	def    snake_case__(						self:						Any			):
			lowercase__					:      Optional[Any]      =		OrderedDict({'input_ids': {0: 'batch', 1: 'sequence'}}			)
			if self.use_past:
					self.fill_with_past_key_values_(__UpperCAmelCase,					direction='inputs'			)
					lowercase__					:      Any      =		{0: "batch", 1: "past_sequence + sequence"}
			else:
					lowercase__					:      List[Any]      =		{0: "batch", 1: "sequence"}

			return common_inputs

	@property
	def    snake_case__(						self:						Optional[Any]			):
			return self._config.n_layer

	@property
	def    snake_case__(						self:						Tuple			):
			return self._config.n_head

	def    snake_case__(						self:						str,					lowerCamelCase_:						Dict,					lowerCamelCase_:						Dict = -1,					lowerCamelCase_:						Dict = -1,					lowerCamelCase_:						List[Any] = False,					lowerCamelCase_:						Dict = None,					):
			lowercase__					:      Dict      =		super(__UpperCAmelCase,					self			).generate_dummy_inputs(
			    __UpperCAmelCase,					batch_size=__UpperCAmelCase,					seq_length=__UpperCAmelCase,					is_pair=__UpperCAmelCase,					framework=__UpperCAmelCase			)

			# We need to order the input in the way they appears in the forward()
			lowercase__					:      Union[str, Any]      =		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

							lowercase__					:      Union[str, Any]      =		common_inputs["input_ids"].shape
							# Not using the same length for past_key_values
							lowercase__					:      Any      =		seqlen + 2
							lowercase__					:      str      =		(
							    batch,
							    self.num_attention_heads,
							    past_key_values_length,
							    self._config.hidden_size // self.num_attention_heads,
							)
							lowercase__					:      str      =		[
							    (torch.zeros(__UpperCAmelCase			), torch.zeros(__UpperCAmelCase			)) for _ in range(self.num_layers			)
							]

			lowercase__					:      Union[str, Any]      =		common_inputs["attention_mask"]
			if self.use_past:
					lowercase__					:      str      =		ordered_inputs["attention_mask"].dtype
					lowercase__					:      Union[str, Any]      =		torch.cat(
					    [ordered_inputs['attention_mask'], torch.ones(__UpperCAmelCase,					__UpperCAmelCase,					dtype=__UpperCAmelCase			)],					dim=1			)

			return ordered_inputs





	@property
	def    snake_case__(						self:						List[str]			):
			return 13







import operator as op


lowerCamelCase		       =      """scaler.pt"""
lowerCamelCase		       =      """pytorch_model"""
lowerCamelCase		       =      """random_states"""
lowerCamelCase		       =      """optimizer"""
lowerCamelCase		       =      """scheduler"""
lowerCamelCase		       =      """pytorch_model.bin"""
lowerCamelCase		       =      """pytorch_model.bin.index.json"""
lowerCamelCase		       =      """model.safetensors"""
lowerCamelCase		       =      """model.safetensors.index.json"""
lowerCamelCase		       =      """1.10.2"""
lowerCamelCase		       =      """py38"""
lowerCamelCase		       =      """4.17.0"""
lowerCamelCase		       =      ["""ml.p3.16xlarge""", """ml.p3dn.24xlarge""", """ml.p4dn.24xlarge"""]
lowerCamelCase		       =      ["""FULL_SHARD""", """SHARD_GRAD_OP""", """NO_SHARD""", """HYBRID_SHARD""", """HYBRID_SHARD_ZERO2"""]
lowerCamelCase		       =      ["""TRANSFORMER_BASED_WRAP""", """SIZE_BASED_WRAP""", """NO_WRAP"""]
lowerCamelCase		       =      ["""BACKWARD_PRE""", """BACKWARD_POST""", """NO_PREFETCH"""]
lowerCamelCase		       =      ["""FULL_STATE_DICT""", """LOCAL_STATE_DICT""", """SHARDED_STATE_DICT"""]
lowerCamelCase		       =      """2.0.1"""
lowerCamelCase		       =      ["""pdsh""", """standard""", """openmpi""", """mvapich"""]
lowerCamelCase		       =      ["""default""", """reduce-overhead""", """max-autotune"""]

lowerCamelCase		       =      {""">""": op.gt, """>=""": op.ge, """==""": op.eq, """!=""": op.ne, """<=""": op.le, """<""": op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
lowerCamelCase		       =      [
    """nnodes""",
    """nproc_per_node""",
    """rdzv_backend""",
    """rdzv_endpoint""",
    """rdzv_id""",
    """rdzv_conf""",
    """standalone""",
    """max_restarts""",
    """monitor_interval""",
    """start_method""",
    """role""",
    """module""",
    """m""",
    """no_python""",
    """run_path""",
    """log_dir""",
    """r""",
    """redirects""",
    """t""",
    """tee""",
    """node_rank""",
    """master_addr""",
    """master_port""",
]

lowerCamelCase		       =      ["""DEEPSPEED""", """MULTI_GPU""", """FSDP""", """MEGATRON_LM"""]
lowerCamelCase		       =      ["""DEEPSPEED""", """MULTI_XPU""", """FSDP"""]


