






import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


_UpperCAmelCase			:		Any   				=    logging.get_logger(__name__)





def 	lowerCAmelCase_		(lowercase__						:       Tuple     ,					lowercase__						:       Optional[int]     ,					lowercase__						:       List[Any]     ,					lowercase__						:       Optional[int]=False       )		->    Dict:


				'''simple docstring'''




				try:
								import torch  # noqa: F401
				except ImportError:
								logger.error(
								    '''Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see'''
								    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
								    ''' instructions.'''       )
								raise

				if not is_sharded:
								lowerCAmelCase__      			=						os.path.abspath(lowercase__       )
								logger.info(f'Loading PyTorch weights from {pt_path}'       )

								lowerCAmelCase__      			=						torch.load(lowercase__     ,					map_location='''cpu'''       )
								logger.info(f'PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()       ):,} parameters.'       )

								lowerCAmelCase__      			=						convert_pytorch_state_dict_to_flax(lowercase__     ,					lowercase__       )
				else:
								# model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
								lowerCAmelCase__      			=						convert_pytorch_sharded_state_dict_to_flax(lowercase__     ,					lowercase__       )
				return flax_state_dict





def 	lowerCAmelCase_		(lowercase__						:       Tuple[str]     ,					lowercase__						:       np.ndarray     ,					lowercase__						:       Dict[str, jnp.ndarray]     ,					lowercase__						:       str     ,					)		->    (Tuple[str], np.ndarray):


				'''simple docstring'''




				def is_key_or_prefix_key_in_dict(lowercase__						:       Tuple[str]       ) -> bool:
								return len(set(lowercase__       ) & {key, (model_prefix,) + key}       ) > 0

				# layer norm
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''scale''',)
				if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(lowercase__       ):
								return renamed_pt_tuple_key, pt_tensor

				# batch norm layer mean
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''mean''',)
				if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(lowercase__       ):
								return renamed_pt_tuple_key, pt_tensor

				# batch norm layer var
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''var''',)
				if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(lowercase__       ):
								return renamed_pt_tuple_key, pt_tensor

				# embedding
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''embedding''',)
				if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(lowercase__       ):
								return renamed_pt_tuple_key, pt_tensor

				# conv layer
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''kernel''',)
				if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(lowercase__       ):
								lowerCAmelCase__      			=						pt_tensor.transpose(2     ,					3     ,					1     ,					0       )
								return renamed_pt_tuple_key, pt_tensor

				# linear layer
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''kernel''',)
				if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(lowercase__       ):
								lowerCAmelCase__      			=						pt_tensor.T
								return renamed_pt_tuple_key, pt_tensor

				# old PyTorch layer norm weight
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''weight''',)
				if pt_tuple_key[-1] == "gamma":
								return renamed_pt_tuple_key, pt_tensor

				# old PyTorch layer norm bias
				lowerCAmelCase__      			=						pt_tuple_key[:-1] + ('''bias''',)
				if pt_tuple_key[-1] == "beta":
								return renamed_pt_tuple_key, pt_tensor

				# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
				lowerCAmelCase__      			=						None
				if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
								lowerCAmelCase__      			=						pt_tuple_key[-2] + '''_g'''
				elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
								lowerCAmelCase__      			=						pt_tuple_key[-2] + '''_v'''
				if name is not None:
								lowerCAmelCase__      			=						pt_tuple_key[:-3] + (name,)
								return renamed_pt_tuple_key, pt_tensor

				return pt_tuple_key, pt_tensor





def 	lowerCAmelCase_		(lowercase__						:       Optional[int]     ,					lowercase__						:       List[str]       )		->    Tuple:


				'''simple docstring'''




				lowerCAmelCase__      			=						{k: v.numpy() for k, v in pt_state_dict.items()}

				lowerCAmelCase__      			=						flax_model.base_model_prefix

				# use params dict if the model contains batch norm layers
				if "params" in flax_model.params:
								lowerCAmelCase__      			=						flax_model.params['''params''']
				else:
								lowerCAmelCase__      			=						flax_model.params
				lowerCAmelCase__      			=						flatten_dict(lowercase__       )

				# add batch_stats keys,values to dict
				if "batch_stats" in flax_model.params:
								lowerCAmelCase__      			=						flatten_dict(flax_model.params['''batch_stats''']       )
								random_flax_state_dict.update(lowercase__       )

				lowerCAmelCase__      			=						{}

				lowerCAmelCase__      			=						(model_prefix not in flax_model_params) and (
				    model_prefix in {k.split('''.'''       )[0] for k in pt_state_dict.keys()}
				)
				lowerCAmelCase__      			=						(model_prefix in flax_model_params) and (
				    model_prefix not in {k.split('''.'''       )[0] for k in pt_state_dict.keys()}
				)

				# Need to change some parameters name to match Flax names
				for pt_key, pt_tensor in pt_state_dict.items():
								lowerCAmelCase__      			=						tuple(pt_key.split('''.'''       )       )

								# remove base model prefix if necessary
								lowerCAmelCase__      			=						pt_tuple_key[0] == model_prefix
								if load_model_with_head_into_base_model and has_base_model_prefix:
												lowerCAmelCase__      			=						pt_tuple_key[1:]

								# Correctly rename weight parameters
								lowerCAmelCase__			,		lowerCAmelCase__      			=						rename_key_and_reshape_tensor(
								    lowercase__     ,					lowercase__     ,					lowercase__     ,					lowercase__       )

								# add model prefix if necessary
								lowerCAmelCase__      			=						(model_prefix,) + flax_key in random_flax_state_dict
								if load_base_model_into_model_with_head and require_base_model_prefix:
												lowerCAmelCase__      			=						(model_prefix,) + flax_key

								if flax_key in random_flax_state_dict:
												if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																raise ValueError(
																    f'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
																    f'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'       )

        # add batch stats if the model contains batchnorm layers
								if "batch_stats" in flax_model.params:
												if "mean" in flax_key[-1] or "var" in flax_key[-1]:
																lowerCAmelCase__      			=						jnp.asarray(lowercase__       )
																continue
												# remove num_batches_tracked key
												if "num_batches_tracked" in flax_key[-1]:
																flax_state_dict.pop(lowercase__     ,					lowercase__       )
																continue

												# also add unexpected weight so that warning is thrown
												lowerCAmelCase__      			=						jnp.asarray(lowercase__       )

								else:
												# also add unexpected weight so that warning is thrown
												lowerCAmelCase__      			=						jnp.asarray(lowercase__       )

				return unflatten_dict(lowercase__       )





def 	lowerCAmelCase_		(lowercase__						:       Any     ,					lowercase__						:       Union[str, Any]       )		->    Optional[Any]:


				'''simple docstring'''




				import torch

				# Load the index
				lowerCAmelCase__      			=						{}
				for shard_file in shard_filenames:
								# load using msgpack utils
								lowerCAmelCase__      			=						torch.load(lowercase__       )
								lowerCAmelCase__      			=						{k: v.numpy() for k, v in pt_state_dict.items()}

								lowerCAmelCase__      			=						flax_model.base_model_prefix

								# use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
								if "batch_stats" in flax_model.params:
												lowerCAmelCase__      			=						flax_model.params['''params''']

												lowerCAmelCase__      			=						flatten_dict(lowercase__       )
												random_flax_state_dict.update(flatten_dict(flax_model.params['''batch_stats''']       )       )
								else:
												lowerCAmelCase__      			=						flax_model.params
												lowerCAmelCase__      			=						flatten_dict(lowercase__       )

								lowerCAmelCase__      			=						(model_prefix not in flax_model_params) and (
								    model_prefix in {k.split('''.'''       )[0] for k in pt_state_dict.keys()}
								)
								lowerCAmelCase__      			=						(model_prefix in flax_model_params) and (
								    model_prefix not in {k.split('''.'''       )[0] for k in pt_state_dict.keys()}
								)
								# Need to change some parameters name to match Flax names
								for pt_key, pt_tensor in pt_state_dict.items():
												lowerCAmelCase__      			=						tuple(pt_key.split('''.'''       )       )

												# remove base model prefix if necessary
												lowerCAmelCase__      			=						pt_tuple_key[0] == model_prefix
												if load_model_with_head_into_base_model and has_base_model_prefix:
																lowerCAmelCase__      			=						pt_tuple_key[1:]

												# Correctly rename weight parameters
												lowerCAmelCase__			,		lowerCAmelCase__      			=						rename_key_and_reshape_tensor(
												    lowercase__     ,					lowercase__     ,					lowercase__     ,					lowercase__       )
												# add model prefix if necessary
												lowerCAmelCase__      			=						(model_prefix,) + flax_key in random_flax_state_dict
												if load_base_model_into_model_with_head and require_base_model_prefix:
																lowerCAmelCase__      			=						(model_prefix,) + flax_key

												if flax_key in random_flax_state_dict:
																if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																				raise ValueError(
																				    f'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
																				    f'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'       )

            # add batch stats if the model contains batchnorm layers
												if "batch_stats" in flax_model.params:
																if "mean" in flax_key[-1]:
																				lowerCAmelCase__      			=						jnp.asarray(lowercase__       )
																				continue
																if "var" in flax_key[-1]:
																				lowerCAmelCase__      			=						jnp.asarray(lowercase__       )
																				continue
																# remove num_batches_tracked key
																if "num_batches_tracked" in flax_key[-1]:
																				flax_state_dict.pop(lowercase__     ,					lowercase__       )
																				continue

																# also add unexpected weight so that warning is thrown
																lowerCAmelCase__      			=						jnp.asarray(lowercase__       )

												else:
																# also add unexpected weight so that warning is thrown
																lowerCAmelCase__      			=						jnp.asarray(lowercase__       )
				return unflatten_dict(lowercase__       )





def 	lowerCAmelCase_		(lowercase__						:       Dict     ,					lowercase__						:       Tuple       )		->    Union[str, Any]:


				'''simple docstring'''




				lowerCAmelCase__      			=						os.path.abspath(lowercase__       )
				logger.info(f'Loading Flax weights from {flax_checkpoint_path}'       )

				# import correct flax class
				lowerCAmelCase__      			=						getattr(lowercase__     ,					'''Flax''' + model.__class__.__name__       )

				# load flax weight dict
				with open(lowercase__     ,					'''rb'''       ) as state_f:
								try:
												lowerCAmelCase__      			=						from_bytes(lowercase__     ,					state_f.read()       )
								except UnpicklingError:
												raise EnvironmentError(f'Unable to convert {flax_checkpoint_path} to Flax deserializable object. '       )

				return load_flax_weights_in_pytorch_model(lowercase__     ,					lowercase__       )





def 	lowerCAmelCase_		(lowercase__						:       Union[str, Any]     ,					lowercase__						:       Optional[Any]       )		->    Optional[Any]:


				'''simple docstring'''




				try:
								import torch  # noqa: F401
				except ImportError:
								logger.error(
								    '''Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see'''
								    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
								    ''' instructions.'''       )
								raise

				# check if we have bf16 weights
				lowerCAmelCase__      			=						flatten_dict(jax.tree_util.tree_map(lambda lowercase__       : x.dtype == jnp.bfloataa     ,					lowercase__       )       ).values()
				if any(lowercase__       ):
								# convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
								# and bf16 is not fully supported in PT yet.
								logger.warning(
								    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
								    '''before loading those in PyTorch model.'''       )
								lowerCAmelCase__      			=						jax.tree_util.tree_map(
								    lambda lowercase__       : params.astype(np.floataa       ) if params.dtype == jnp.bfloataa else params     ,					lowercase__       )

				lowerCAmelCase__      			=						flatten_dict(lowercase__       )
				lowerCAmelCase__      			=						pt_model.state_dict()

				lowerCAmelCase__      			=						(pt_model.base_model_prefix in flax_state) and (
				    pt_model.base_model_prefix not in {k.split('''.'''       )[0] for k in pt_model_dict.keys()}
				)
				lowerCAmelCase__      			=						(pt_model.base_model_prefix not in flax_state) and (
				    pt_model.base_model_prefix in {k.split('''.'''       )[0] for k in pt_model_dict.keys()}
				)

				# keep track of unexpected & missing keys
				lowerCAmelCase__      			=						[]
				lowerCAmelCase__      			=						set(pt_model_dict.keys()       )

				for flax_key_tuple, flax_tensor in flax_state_dict.items():
								lowerCAmelCase__      			=						flax_key_tuple[0] == pt_model.base_model_prefix
								lowerCAmelCase__      			=						'''.'''.join((pt_model.base_model_prefix,) + flax_key_tuple       ) in pt_model_dict

								# adapt flax_key to prepare for loading from/to base model only
								if load_model_with_head_into_base_model and has_base_model_prefix:
												lowerCAmelCase__      			=						flax_key_tuple[1:]
								elif load_base_model_into_model_with_head and require_base_model_prefix:
												lowerCAmelCase__      			=						(pt_model.base_model_prefix,) + flax_key_tuple

								# rename flax weights to PyTorch format
								if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(lowercase__       ) not in pt_model_dict:
												# conv layer
												lowerCAmelCase__      			=						flax_key_tuple[:-1] + ('''weight''',)
												lowerCAmelCase__      			=						jnp.transpose(lowercase__     ,					(3, 2, 0, 1)       )
								elif flax_key_tuple[-1] == "kernel" and ".".join(lowercase__       ) not in pt_model_dict:
												# linear layer
												lowerCAmelCase__      			=						flax_key_tuple[:-1] + ('''weight''',)
												lowerCAmelCase__      			=						flax_tensor.T
								elif flax_key_tuple[-1] in ["scale", "embedding"]:
												lowerCAmelCase__      			=						flax_key_tuple[:-1] + ('''weight''',)

								# adding batch stats from flax batch norm to pt
								elif "mean" in flax_key_tuple[-1]:
												lowerCAmelCase__      			=						flax_key_tuple[:-1] + ('''running_mean''',)
								elif "var" in flax_key_tuple[-1]:
												lowerCAmelCase__      			=						flax_key_tuple[:-1] + ('''running_var''',)

								if "batch_stats" in flax_state:
												lowerCAmelCase__      			=						'''.'''.join(flax_key_tuple[1:]       )  # Remove the params/batch_stats header
								else:
												lowerCAmelCase__      			=						'''.'''.join(lowercase__       )

								# We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
								lowerCAmelCase__      			=						{}
								# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
								for key in pt_model_dict:
												lowerCAmelCase__      			=						key.split('''.'''       )
												lowerCAmelCase__      			=						None
												if key_components[-3::2] == ["parametrizations", "original0"]:
																lowerCAmelCase__      			=						key_components[-2] + '''_g'''
												elif key_components[-3::2] == ["parametrizations", "original1"]:
																lowerCAmelCase__      			=						key_components[-2] + '''_v'''
												if name is not None:
																lowerCAmelCase__      			=						key_components[:-3] + [name]
																lowerCAmelCase__      			=						'''.'''.join(lowercase__       )
																lowerCAmelCase__      			=						key

								if flax_key in special_pt_names:
												lowerCAmelCase__      			=						special_pt_names[flax_key]

								if flax_key in pt_model_dict:
												if flax_tensor.shape != pt_model_dict[flax_key].shape:
																raise ValueError(
																    f'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
																    f'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'       )
												else:
																# add weight to pytorch dict
																lowerCAmelCase__      			=						np.asarray(lowercase__       ) if not isinstance(lowercase__     ,					np.ndarray       ) else flax_tensor
																lowerCAmelCase__      			=						torch.from_numpy(lowercase__       )
																# remove from missing keys
																missing_keys.remove(lowercase__       )
								else:
												# weight is not expected by PyTorch model
												unexpected_keys.append(lowercase__       )

				pt_model.load_state_dict(lowercase__       )

				# re-transform missing_keys to list
				lowerCAmelCase__      			=						list(lowercase__       )

				if len(lowercase__       ) > 0:
								logger.warning(
								    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
								    f' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
								    f' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
								    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
								    f' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
								    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
								    ''' FlaxBertForSequenceClassification model).'''       )
				else:
								logger.warning(f'All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'       )
				if len(lowercase__       ) > 0:
								logger.warning(
								    f'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
								    f' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
								    ''' use it for predictions and inference.'''       )
				else:
								logger.warning(
								    f'All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'
								    '''If your task is similar to the task the model of the checkpoint was trained on, '''
								    f'you can already use {pt_model.__class__.__name__} for predictions without further training.'       )

				return pt_model













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			:		Tuple   				=    logging.get_logger(__name__)


_UpperCAmelCase			:		int   				=    {
    "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_       (				snake_case__					):
							UpperCamelCase_		:int		      =    'codegen'
							UpperCamelCase_		:int		      =    {
							    'max_position_embeddings': 'n_positions',
							    'hidden_size': 'n_embd',
							    'num_attention_heads': 'n_head',
							    'num_hidden_layers': 'n_layer',
							}




							def __init__(   self     :   Optional[int] ,							SCREAMING_SNAKE_CASE_     :   str=50_400 ,							SCREAMING_SNAKE_CASE_     :   str=2_048 ,							SCREAMING_SNAKE_CASE_     :   int=2_048 ,							SCREAMING_SNAKE_CASE_     :   Any=4_096 ,							SCREAMING_SNAKE_CASE_     :   List[Any]=28 ,							SCREAMING_SNAKE_CASE_     :   str=16 ,							SCREAMING_SNAKE_CASE_     :   str=64 ,							SCREAMING_SNAKE_CASE_     :   Tuple=None ,							SCREAMING_SNAKE_CASE_     :   Dict="gelu_new" ,							SCREAMING_SNAKE_CASE_     :   str=0.0 ,							SCREAMING_SNAKE_CASE_     :   Any=0.0 ,							SCREAMING_SNAKE_CASE_     :   Any=0.0 ,							SCREAMING_SNAKE_CASE_     :   Any=1e-5 ,							SCREAMING_SNAKE_CASE_     :   Any=0.02 ,							SCREAMING_SNAKE_CASE_     :   Union[str, Any]=True ,							SCREAMING_SNAKE_CASE_     :   Optional[int]=50_256 ,							SCREAMING_SNAKE_CASE_     :   Any=50_256 ,							SCREAMING_SNAKE_CASE_     :   List[str]=False ,							**SCREAMING_SNAKE_CASE_     :   Union[str, Any] ,							):
											lowerCAmelCase__      			=						vocab_size
											lowerCAmelCase__      			=						n_ctx
											lowerCAmelCase__      			=						n_positions
											lowerCAmelCase__      			=						n_embd
											lowerCAmelCase__      			=						n_layer
											lowerCAmelCase__      			=						n_head
											lowerCAmelCase__      			=						n_inner
											lowerCAmelCase__      			=						rotary_dim
											lowerCAmelCase__      			=						activation_function
											lowerCAmelCase__      			=						resid_pdrop
											lowerCAmelCase__      			=						embd_pdrop
											lowerCAmelCase__      			=						attn_pdrop
											lowerCAmelCase__      			=						layer_norm_epsilon
											lowerCAmelCase__      			=						initializer_range
											lowerCAmelCase__      			=						use_cache

											lowerCAmelCase__      			=						bos_token_id
											lowerCAmelCase__      			=						eos_token_id

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



class  lowerCAmelCase_       (				snake_case__					):




							def __init__(   self     :   str ,							SCREAMING_SNAKE_CASE_     :   PretrainedConfig ,							SCREAMING_SNAKE_CASE_     :   str = "default" ,							SCREAMING_SNAKE_CASE_     :   List[PatchingSpec] = None ,							SCREAMING_SNAKE_CASE_     :   bool = False ,							):
											super().__init__(SCREAMING_SNAKE_CASE_ ,							task=SCREAMING_SNAKE_CASE_ ,							patching_specs=SCREAMING_SNAKE_CASE_ ,							use_past=SCREAMING_SNAKE_CASE_	)
											if not getattr(self._config ,							'''pad_token_id''' ,							SCREAMING_SNAKE_CASE_	):
															# TODO: how to do that better?
															lowerCAmelCase__      			=						0




							@property
							def 						__snake_case					(   self     :   str	):
											lowerCAmelCase__      			=						OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}}	)
											if self.use_past:
															self.fill_with_past_key_values_(SCREAMING_SNAKE_CASE_ ,							direction='''inputs'''	)
															lowerCAmelCase__      			=						{0: '''batch''', 1: '''past_sequence + sequence'''}
											else:
															lowerCAmelCase__      			=						{0: '''batch''', 1: '''sequence'''}

											return common_inputs




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




							@property
							def 						__snake_case					(   self     :   Union[str, Any]	):
											return self._config.n_head




							def 						__snake_case					(   self     :   Optional[Any] ,							SCREAMING_SNAKE_CASE_     :   PreTrainedTokenizer ,							SCREAMING_SNAKE_CASE_     :   int = -1 ,							SCREAMING_SNAKE_CASE_     :   int = -1 ,							SCREAMING_SNAKE_CASE_     :   bool = False ,							SCREAMING_SNAKE_CASE_     :   Optional[TensorType] = None ,							):
											lowerCAmelCase__      			=						super(SCREAMING_SNAKE_CASE_ ,							self	).generate_dummy_inputs(
											    SCREAMING_SNAKE_CASE_ ,							batch_size=SCREAMING_SNAKE_CASE_ ,							seq_length=SCREAMING_SNAKE_CASE_ ,							is_pair=SCREAMING_SNAKE_CASE_ ,							framework=SCREAMING_SNAKE_CASE_	)

											# We need to order the input in the way they appears in the forward()
											lowerCAmelCase__      			=						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__      			=						common_inputs['''input_ids'''].shape
																			# Not using the same length for past_key_values
																			lowerCAmelCase__      			=						seqlen + 2
																			lowerCAmelCase__      			=						(
																			    batch,
																			    self.num_attention_heads,
																			    past_key_values_length,
																			    self._config.hidden_size // self.num_attention_heads,
																			)
																			lowerCAmelCase__      			=						[
																			    (torch.zeros(SCREAMING_SNAKE_CASE_	), torch.zeros(SCREAMING_SNAKE_CASE_	)) for _ in range(self.num_layers	)
																			]

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

											return ordered_inputs




							@property
							def 						__snake_case					(   self     :   Optional[int]	):
											return 13





