




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


_A													=      logging.get_logger(__name__)


_A													=      {
    """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    (					a_						):
   _lowerCamelCase    :Tuple     				=     "codegen"
   _lowerCamelCase    :Union[str, Any]     				=     {
       "max_position_embeddings": "n_positions",
       "hidden_size": "n_embd",
       "num_attention_heads": "n_head",
       "num_hidden_layers": "n_layer",
   }






   def __init__(		self       :		List[str]   ,     UpperCamelCase       :		Optional[int]=5_04_00   ,     UpperCamelCase       :		str=20_48   ,     UpperCamelCase       :		int=20_48   ,     UpperCamelCase       :		Union[str, Any]=40_96   ,     UpperCamelCase       :		List[str]=28   ,     UpperCamelCase       :		Any=16   ,     UpperCamelCase       :		int=64   ,     UpperCamelCase       :		Tuple=None   ,     UpperCamelCase       :		List[Any]="gelu_new"   ,     UpperCamelCase       :		List[str]=0.0   ,     UpperCamelCase       :		Optional[int]=0.0   ,     UpperCamelCase       :		Optional[Any]=0.0   ,     UpperCamelCase       :		List[Any]=1E-5   ,     UpperCamelCase       :		List[Any]=0.02   ,     UpperCamelCase       :		str=True   ,     UpperCamelCase       :		List[str]=5_02_56   ,     UpperCamelCase       :		List[Any]=5_02_56   ,     UpperCamelCase       :		List[str]=False   ,     **UpperCamelCase       :		Union[str, Any]   ,     )			->						Optional[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			int   =	vocab_size
         lowerCAmelCase__      :			Dict   =	n_ctx
         lowerCAmelCase__      :			List[str]   =	n_positions
         lowerCAmelCase__      :			Tuple   =	n_embd
         lowerCAmelCase__      :			List[str]   =	n_layer
         lowerCAmelCase__      :			List[str]   =	n_head
         lowerCAmelCase__      :			Optional[int]   =	n_inner
         lowerCAmelCase__      :			List[Any]   =	rotary_dim
         lowerCAmelCase__      :			Union[str, Any]   =	activation_function
         lowerCAmelCase__      :			Tuple   =	resid_pdrop
         lowerCAmelCase__      :			int   =	embd_pdrop
         lowerCAmelCase__      :			Optional[int]   =	attn_pdrop
         lowerCAmelCase__      :			Tuple   =	layer_norm_epsilon
         lowerCAmelCase__      :			List[str]   =	initializer_range
         lowerCAmelCase__      :			Any   =	use_cache

         lowerCAmelCase__      :			List[str]   =	bos_token_id
         lowerCAmelCase__      :			Tuple   =	eos_token_id

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






class _lowerCamelCase    (					a_						):






   def __init__(		self       :		List[str]   ,     UpperCamelCase       :		PretrainedConfig   ,     UpperCamelCase       :		str = "default"   ,     UpperCamelCase       :		List[PatchingSpec] = None   ,     UpperCamelCase       :		bool = False   ,     )			->						List[Any]:

         """simple docstring"""
         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?
               lowerCAmelCase__      :			int   =	0






   @property
   def 	_lowerCAmelCase   (		self       :		List[str]   )			->						Mapping[str, Mapping[int, str]]:

         """simple docstring"""
         lowerCAmelCase__      :			Optional[Any]   =	OrderedDict({"""input_ids""": {0: """batch""", 1: """sequence"""}}   )
         if self.use_past:
               self.fill_with_past_key_values_(UpperCamelCase   ,     direction="""inputs"""   )
               lowerCAmelCase__      :			List[str]   =	{0: """batch""", 1: """past_sequence + sequence"""}
         else:
               lowerCAmelCase__      :			Optional[Any]   =	{0: """batch""", 1: """sequence"""}

         return common_inputs






   @property
   def 	_lowerCAmelCase   (		self       :		Any   )			->						int:

         """simple docstring"""
         return self._config.n_layer






   @property
   def 	_lowerCAmelCase   (		self       :		int   )			->						int:

         """simple docstring"""
         return self._config.n_head






   def 	_lowerCAmelCase   (		self       :		Optional[int]   ,     UpperCamelCase       :		PreTrainedTokenizer   ,     UpperCamelCase       :		int = -1   ,     UpperCamelCase       :		int = -1   ,     UpperCamelCase       :		bool = False   ,     UpperCamelCase       :		Optional[TensorType] = None   ,     )			->						Mapping[str, Any]:

         """simple docstring"""
         lowerCAmelCase__      :			Any   =	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()
         lowerCAmelCase__      :			int   =	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

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

         lowerCAmelCase__      :			Tuple   =	common_inputs["""attention_mask"""]
         if self.use_past:
               lowerCAmelCase__      :			Optional[Any]   =	ordered_inputs["""attention_mask"""].dtype
               lowerCAmelCase__      :			List[str]   =	torch.cat(
                   [ordered_inputs["""attention_mask"""], torch.ones(UpperCamelCase   ,     UpperCamelCase   ,     dtype=UpperCamelCase   )]   ,     dim=1   )

         return ordered_inputs






   @property
   def 	_lowerCAmelCase   (		self       :		Tuple   )			->						int:

         """simple docstring"""
         return 13











"""simple docstring"""


from argparse import ArgumentParser

from . import BaseTransformersCLICommand



def 	lowercase_				(			__UpperCAmelCase				)						->    List[str]:
      return DownloadCommand(args.model    ,		args.cache_dir    ,		args.force    ,		args.trust_remote_code				)


class _lowerCamelCase    (					a_						):






   @staticmethod
   def 	_lowerCAmelCase   (		UpperCamelCase       :		ArgumentParser   )			->						List[Any]:

         """simple docstring"""
         lowerCAmelCase__      :			int   =	parser.add_parser("""download"""   )
         download_parser.add_argument(
             """--cache-dir"""   ,     type=UpperCamelCase   ,     default=UpperCamelCase   ,     help="""Path to location to store the models"""   )
         download_parser.add_argument(
             """--force"""   ,     action="""store_true"""   ,     help="""Force the model to be download even if already in cache-dir"""   )
         download_parser.add_argument(
             """--trust-remote-code"""   ,     action="""store_true"""   ,     help="""Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you've reviewed the code as it will execute on your local machine"""   ,     )
         download_parser.add_argument("""model"""   ,     type=UpperCamelCase   ,     help="""Name of the model to download"""   )
         download_parser.set_defaults(func=UpperCamelCase   )






   def __init__(		self       :		Optional[int]   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		str   ,     UpperCamelCase       :		bool   ,     UpperCamelCase       :		bool   )			->						Any:

         """simple docstring"""
         lowerCAmelCase__      :			int   =	model
         lowerCAmelCase__      :			Union[str, Any]   =	cache
         lowerCAmelCase__      :			Optional[int]   =	force
         lowerCAmelCase__      :			Dict   =	trust_remote_code






   def 	_lowerCAmelCase   (		self       :		int   )			->						Any:

         """simple docstring"""
         from ..models.auto import AutoModel, AutoTokenizer

         AutoModel.from_pretrained(
             self._model   ,     cache_dir=self._cache   ,     force_download=self._force   ,     trust_remote_code=self._trust_remote_code   )
         AutoTokenizer.from_pretrained(
             self._model   ,     cache_dir=self._cache   ,     force_download=self._force   ,     trust_remote_code=self._trust_remote_code   )





