






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


__SCREAMING_SNAKE_CASE   =   logging.get_logger(__name__)


__SCREAMING_SNAKE_CASE   =   {
    '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 						__UpperCamelCase     (     UpperCamelCase						):
					lowercase_    :				Tuple      =  """codegen"""
					lowercase_    :				Any      =  {
					    """max_position_embeddings""": """n_positions""",
					    """hidden_size""": """n_embd""",
					    """num_attention_heads""": """n_head""",
					    """num_hidden_layers""": """n_layer""",
					}
					def __init__(							self							:	int     , UpperCAmelCase							:	Any=5_0400     , UpperCAmelCase							:	Optional[Any]=2048     , UpperCAmelCase							:	int=2048     , UpperCAmelCase							:	Any=4096     , UpperCAmelCase							:	int=28     , UpperCAmelCase							:	Optional[int]=16     , UpperCAmelCase							:	Optional[int]=64     , UpperCAmelCase							:	Optional[int]=None     , UpperCAmelCase							:	str="gelu_new"     , UpperCAmelCase							:	List[Any]=0.0     , UpperCAmelCase							:	str=0.0     , UpperCAmelCase							:	List[str]=0.0     , UpperCAmelCase							:	List[str]=1e-5     , UpperCAmelCase							:	Tuple=0.0_2     , UpperCAmelCase							:	Union[str, Any]=True     , UpperCAmelCase							:	List[Any]=5_0256     , UpperCAmelCase							:	List[Any]=5_0256     , UpperCAmelCase							:	List[Any]=False     , **UpperCAmelCase							:	Dict     , )	->       Optional[Any]:
											lowerCAmelCase   :List[Any]						=			vocab_size
											lowerCAmelCase   :List[Any]						=			n_ctx
											lowerCAmelCase   :str						=			n_positions
											lowerCAmelCase   :Optional[Any]						=			n_embd
											lowerCAmelCase   :Tuple						=			n_layer
											lowerCAmelCase   :Union[str, Any]						=			n_head
											lowerCAmelCase   :Tuple						=			n_inner
											lowerCAmelCase   :int						=			rotary_dim
											lowerCAmelCase   :Optional[Any]						=			activation_function
											lowerCAmelCase   :Optional[int]						=			resid_pdrop
											lowerCAmelCase   :Tuple						=			embd_pdrop
											lowerCAmelCase   :int						=			attn_pdrop
											lowerCAmelCase   :Union[str, Any]						=			layer_norm_epsilon
											lowerCAmelCase   :List[Any]						=			initializer_range
											lowerCAmelCase   :Optional[int]						=			use_cache

											lowerCAmelCase   :Any						=			bos_token_id
											lowerCAmelCase   :List[str]						=			eos_token_id

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




class 						__UpperCamelCase     (     UpperCamelCase						):
					def __init__(							self							:	int     , UpperCAmelCase							:	PretrainedConfig     , UpperCAmelCase							:	str = "default"     , UpperCAmelCase							:	List[PatchingSpec] = None     , UpperCAmelCase							:	bool = False     , )	->       Dict:
											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   :List[Any]						=			0
					@property
					def        UpperCAmelCase__ (							self							:	str   )	->       Mapping[str, Mapping[int, str]]:
											lowerCAmelCase   :str						=			OrderedDict({'input_ids': {0: 'batch', 1: 'sequence'}}   )
											if self.use_past:
																	self.fill_with_past_key_values_(UpperCAmelCase     , direction='inputs'   )
																	lowerCAmelCase   :Tuple						=			{0: 'batch', 1: 'past_sequence + sequence'}
											else:
																	lowerCAmelCase   :List[str]						=			{0: 'batch', 1: 'sequence'}

											return common_inputs
					@property
					def        UpperCAmelCase__ (							self							:	int   )	->       int:
											return self._config.n_layer
					@property
					def        UpperCAmelCase__ (							self							:	str   )	->       int:
											return self._config.n_head
					def        UpperCAmelCase__ (							self							:	Optional[Any]     , UpperCAmelCase							:	PreTrainedTokenizer     , UpperCAmelCase							:	int = -1     , UpperCAmelCase							:	int = -1     , UpperCAmelCase							:	bool = False     , UpperCAmelCase							:	Optional[TensorType] = None     , )	->       Mapping[str, Any]:
											lowerCAmelCase   :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()
											lowerCAmelCase   :str						=			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   :int						=			common_inputs['input_ids'].shape
																							# Not using the same length for past_key_values
																							lowerCAmelCase   :Optional[Any]						=			seqlen + 2
																							lowerCAmelCase   :Optional[int]						=			(
																							    batch,
																							    self.num_attention_heads,
																							    past_key_values_length,
																							    self._config.hidden_size // self.num_attention_heads,
																							)
																							lowerCAmelCase   :Union[str, Any]						=			[
																							    (torch.zeros(UpperCAmelCase   ), torch.zeros(UpperCAmelCase   )) for _ in range(self.num_layers   )
																							]

											lowerCAmelCase   :Any						=			common_inputs['attention_mask']
											if self.use_past:
																	lowerCAmelCase   :Tuple						=			ordered_inputs['attention_mask'].dtype
																	lowerCAmelCase   :Optional[Any]						=			torch.cat(
																	    [ordered_inputs['attention_mask'], torch.ones(UpperCAmelCase     , UpperCAmelCase     , dtype=UpperCAmelCase   )]     , dim=1   )

											return ordered_inputs
					@property
					def        UpperCAmelCase__ (							self							:	int   )	->       int:
											return 13







"""simple docstring"""


from __future__ import annotations

from collections import namedtuple
from dataclasses import dataclass




@dataclass
class 						__UpperCamelCase     :
					lowercase_    :				int
					lowercase_    :				TreeNode | None      =  None
					lowercase_    :				TreeNode | None      =  None


__SCREAMING_SNAKE_CASE   =   namedtuple('CoinsDistribResult', 'moves excess')


def        UpperCAmelCase      (					a__ ):





						'''simple docstring'''






						if root is None:
												return 0

						# Validation
						def count_nodes(a__ ) -> int:
												if node is None:
																		return 0

												return count_nodes(node.left ) + count_nodes(node.right ) + 1

						def count_coins(a__ ) -> int:
												if node is None:
																		return 0

												return count_coins(node.left ) + count_coins(node.right ) + node.data

						if count_nodes(a__ ) != count_coins(a__ ):
												raise ValueError('The nodes number should be same as the number of coins' )

						# Main calculation
						def get_distrib(a__ ) -> CoinsDistribResult:

												if node is None:
																		return CoinsDistribResult(0     ,      1 )

												lowerCAmelCase	, lowerCAmelCase   :Optional[int]						=			get_distrib(node.left )
												lowerCAmelCase	, lowerCAmelCase   :str						=			get_distrib(node.right )

												lowerCAmelCase   :str						=			1 - left_distrib_excess
												lowerCAmelCase   :Optional[Any]						=			1 - right_distrib_excess

												lowerCAmelCase   :Any						=			(
												    left_distrib_moves
												    + right_distrib_moves
												    + abs(a__ )
												    + abs(a__ )
												)
												lowerCAmelCase   :Tuple						=			node.data - coins_to_left - coins_to_right

												return CoinsDistribResult(a__     ,      a__ )

						return get_distrib(a__ )[0]


if __name__ == "__main__":
		import doctest

		doctest.testmod()