



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


__lowerCAmelCase									=   logging.get_logger(__name__)


__lowerCAmelCase									=   {
    '''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 	__magic_name__						(    _UpperCamelCase   ):
							lowerCAmelCase       :    str   				=					'codegen'
							lowerCAmelCase       :    Optional[Any]   				=					{
							    'max_position_embeddings': 'n_positions',
							    'hidden_size': 'n_embd',
							    'num_attention_heads': 'n_head',
							    'num_hidden_layers': 'n_layer',
							}







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

									_a     :      Any        =							bos_token_id
									_a     :      Dict        =							eos_token_id

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

class 	__magic_name__						(    _UpperCamelCase   ):







							def __init__(							self					:       List[str]	,_UpperCAmelCase					:       PretrainedConfig	,_UpperCAmelCase					:       str = "default"	,_UpperCAmelCase					:       List[PatchingSpec] = None	,_UpperCAmelCase					:       bool = 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?
											_a     :      Dict        =							0







							@property
							def 							__lowercase			(							self					:       Tuple		):
									_a     :      List[Any]        =							OrderedDict({'input_ids': {0: 'batch', 1: 'sequence'}}		)
									if self.use_past:
											self.fill_with_past_key_values_(_UpperCAmelCase	,direction='inputs'		)
											_a     :      Union[str, Any]        =							{0: 'batch', 1: 'past_sequence + sequence'}
									else:
											_a     :      Any        =							{0: 'batch', 1: 'sequence'}

									return common_inputs







							@property
							def 							__lowercase			(							self					:       Dict		):
									return self._config.n_layer







							@property
							def 							__lowercase			(							self					:       Optional[int]		):
									return self._config.n_head







							def 							__lowercase			(							self					:       int	,_UpperCAmelCase					:       PreTrainedTokenizer	,_UpperCAmelCase					:       int = -1	,_UpperCAmelCase					:       int = -1	,_UpperCAmelCase					:       bool = False	,_UpperCAmelCase					:       Optional[TensorType] = None	,):
									_a     :      List[str]        =							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()
									_a     :      List[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

													_a       ,					_a     :      int        =							common_inputs['input_ids'].shape
													# Not using the same length for past_key_values
													_a     :      Optional[int]        =							seqlen + 2
													_a     :      int        =							(
													    batch,
													    self.num_attention_heads,
													    past_key_values_length,
													    self._config.hidden_size // self.num_attention_heads,
													)
													_a     :      Any        =							[
													    (torch.zeros(_UpperCAmelCase		), torch.zeros(_UpperCAmelCase		)) for _ in range(self.num_layers		)
													]

									_a     :      Optional[Any]        =							common_inputs['attention_mask']
									if self.use_past:
											_a     :      str        =							ordered_inputs['attention_mask'].dtype
											_a     :      Dict        =							torch.cat(
											    [ordered_inputs['attention_mask'], torch.ones(_UpperCAmelCase	,_UpperCAmelCase	,dtype=_UpperCAmelCase		)]	,dim=1		)

									return ordered_inputs







							@property
							def 							__lowercase			(							self					:       Dict		):
									return 13





"""simple docstring"""




import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
				from deepspeed import DeepSpeedEngine

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






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if is_torch_version('<'				,			'2.0.0' ) or not hasattr(UpperCamelCase__				,			'_dynamo' ):
														return False
							return isinstance(UpperCamelCase__				,			torch._dynamo.eval_frame.OptimizedModule )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ = True ):






							"""simple docstring"""





							A__		   =							(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

							A__		   =							is_compiled_module(UpperCamelCase__ )
							if is_compiled:
														A__		   =							model
														A__		   =							model._orig_mod

							if is_deepspeed_available():
														options += (DeepSpeedEngine,)

							while isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
														A__		   =							model.module

							if not keep_fpaa_wrapper:
														A__		   =							getattr(UpperCamelCase__				,			'forward' )
														A__		   =							model.__dict__.pop('_original_forward'				,			UpperCamelCase__ )
														if original_forward is not None:
																					while hasattr(UpperCamelCase__				,			'__wrapped__' ):
																												A__		   =							forward.__wrapped__
																												if forward == original_forward:
																																			break
																					A__		   =							forward
														if getattr(UpperCamelCase__				,			'_converted_to_transformer_engine'				,			UpperCamelCase__ ):
																					convert_model(UpperCamelCase__				,			to_transformer_engine=UpperCamelCase__ )

							if is_compiled:
														A__		   =							model
														A__		   =							compiled_model

							return model






def      UpperCAmelCase    (		):






							"""simple docstring"""





							PartialState().wait_for_everyone()






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							if PartialState().distributed_type == DistributedType.TPU:
														xm.save(UpperCamelCase__				,			UpperCamelCase__ )
							elif PartialState().local_process_index == 0:
														torch.save(UpperCamelCase__				,			UpperCamelCase__ )






@contextmanager
def      UpperCAmelCase    (		**UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in kwargs.items():
														A__		   =							str(UpperCamelCase__ )

							yield

							for key in kwargs:
														if key.upper() in os.environ:
																					del os.environ[key.upper()]






def      UpperCAmelCase    (		UpperCamelCase__ ):






							"""simple docstring"""





							if not hasattr(UpperCamelCase__				,			'__qualname__' ) and not hasattr(UpperCamelCase__				,			'__name__' ):
														A__		   =							getattr(UpperCamelCase__				,			'__class__'				,			UpperCamelCase__ )
							if hasattr(UpperCamelCase__				,			'__qualname__' ):
														return obj.__qualname__
							if hasattr(UpperCamelCase__				,			'__name__' ):
														return obj.__name__
							return str(UpperCamelCase__ )






def      UpperCAmelCase    (		UpperCamelCase__				,			UpperCamelCase__ ):






							"""simple docstring"""





							for key, value in source.items():
														if isinstance(UpperCamelCase__				,			UpperCamelCase__ ):
																					A__		   =							destination.setdefault(UpperCamelCase__				,			{} )
																					merge_dicts(UpperCamelCase__				,			UpperCamelCase__ )
														else:
																					A__		   =							value

							return destination






def      UpperCAmelCase    (		UpperCamelCase__ = None ):






							"""simple docstring"""





							if port is None:
														A__		   =							29_500
							with socket.socket(socket.AF_INET				,			socket.SOCK_STREAM ) as s:
														return s.connect_ex(('localhost', port) ) == 0





