





"""simple docstring"""





import argparse
import os
from pathlib import Path

import torch
from bark.generation import _load_model as _bark_load_model
from huggingface_hub import hf_hub_download

from transformers import EncodecConfig, EncodecModel, set_seed
from transformers.models.bark.configuration_bark import (
    BarkCoarseConfig,
    BarkConfig,
    BarkFineConfig,
    BarkSemanticConfig,
)
from transformers.models.bark.generation_configuration_bark import (
    BarkCoarseGenerationConfig,
    BarkFineGenerationConfig,
    BarkGenerationConfig,
    BarkSemanticGenerationConfig,
)
from transformers.models.bark.modeling_bark import BarkCoarseModel, BarkFineModel, BarkModel, BarkSemanticModel
from transformers.utils import logging


logging.set_verbosity_info()
_UpperCamelCase					      =   logging.get_logger(__name__)

set_seed(770)


_UpperCamelCase					      =   {
    'c_attn': 'att_proj',
    'c_proj': 'out_proj',
    'c_fc': 'in_proj',
    'transformer.': '',
    'h.': 'layers.',
    'ln_1': 'layernorm_1',
    'ln_2': 'layernorm_2',
    'ln_f': 'layernorm_final',
    'wpe': 'position_embeds_layer',
    'wte': 'input_embeds_layer',
}


_UpperCamelCase					      =   {
    'text_small': {
        'repo_id': 'suno/bark',
        'file_name': 'text.pt',
    },
    'coarse_small': {
        'repo_id': 'suno/bark',
        'file_name': 'coarse.pt',
    },
    'fine_small': {
        'repo_id': 'suno/bark',
        'file_name': 'fine.pt',
    },
    'text': {
        'repo_id': 'suno/bark',
        'file_name': 'text_2.pt',
    },
    'coarse': {
        'repo_id': 'suno/bark',
        'file_name': 'coarse_2.pt',
    },
    'fine': {
        'repo_id': 'suno/bark',
        'file_name': 'fine_2.pt',
    },
}

_UpperCamelCase					      =   os.path.dirname(os.path.abspath(__file__))
_UpperCamelCase					      =   os.path.join(os.path.expanduser('~'), '.cache')
_UpperCamelCase					      =   os.path.join(os.getenv('XDG_CACHE_HOME', default_cache_dir), 'suno', 'bark_v0')



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Union[str, Any]     ,    SCREAMING_SNAKE_CASE							:				Dict=False ):


							'''simple docstring'''







							__lowerCamelCase   :   Union[str, Any]       =model_type
							if use_small:
														key += "_small"
							return os.path.join(SCREAMING_SNAKE_CASE     ,    REMOTE_MODEL_PATHS[key]['''file_name'''] )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				int     ,    SCREAMING_SNAKE_CASE							:				Dict ):


							'''simple docstring'''







							os.makedirs(SCREAMING_SNAKE_CASE     ,    exist_ok=SCREAMING_SNAKE_CASE )
							hf_hub_download(repo_id=SCREAMING_SNAKE_CASE     ,    filename=SCREAMING_SNAKE_CASE     ,    local_dir=SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				List[str]     ,    SCREAMING_SNAKE_CASE							:				Tuple     ,    SCREAMING_SNAKE_CASE							:				Dict=False     ,    SCREAMING_SNAKE_CASE							:				List[str]="text" ):


							'''simple docstring'''







							if model_type == "text":
														__lowerCamelCase   :   Union[str, Any]       =BarkSemanticModel
														__lowerCamelCase   :   Union[str, Any]       =BarkSemanticConfig
														__lowerCamelCase   :   int       =BarkSemanticGenerationConfig
							elif model_type == "coarse":
														__lowerCamelCase   :   List[Any]       =BarkCoarseModel
														__lowerCamelCase   :   int       =BarkCoarseConfig
														__lowerCamelCase   :   Tuple       =BarkCoarseGenerationConfig
							elif model_type == "fine":
														__lowerCamelCase   :   Optional[Any]       =BarkFineModel
														__lowerCamelCase   :   str       =BarkFineConfig
														__lowerCamelCase   :   Dict       =BarkFineGenerationConfig
							else:
														raise NotImplementedError()
							__lowerCamelCase   :   Tuple       =F'{model_type}_small' if use_small else model_type
							__lowerCamelCase   :   List[Any]       =REMOTE_MODEL_PATHS[model_key]
							if not os.path.exists(SCREAMING_SNAKE_CASE ):
														logger.info(F'{model_type} model not found, downloading into `{CACHE_DIR}`.' )
														_download(model_info['''repo_id''']     ,    model_info['''file_name'''] )
							__lowerCamelCase   :   Optional[int]       =torch.load(SCREAMING_SNAKE_CASE     ,    map_location=SCREAMING_SNAKE_CASE )
							# this is a hack
							__lowerCamelCase   :   Tuple       =checkpoint['''model_args''']
							if "input_vocab_size" not in model_args:
														__lowerCamelCase   :   Dict       =model_args['''vocab_size''']
														__lowerCamelCase   :   int       =model_args['''vocab_size''']
														del model_args["vocab_size"]

							# convert Bark model arguments to HF Bark model arguments
							__lowerCamelCase   :   Dict       =model_args.pop('''n_head''' )
							__lowerCamelCase   :   Dict       =model_args.pop('''n_embd''' )
							__lowerCamelCase   :   Optional[int]       =model_args.pop('''n_layer''' )

							__lowerCamelCase   :   Optional[int]       =ConfigClass(**checkpoint['''model_args'''] )
							__lowerCamelCase   :   Any       =ModelClass(config=SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   List[str]       =GenerationConfigClass()

							__lowerCamelCase   :   Optional[Any]       =model_generation_config
							__lowerCamelCase   :   Any       =checkpoint['''model''']
							# fixup checkpoint
							__lowerCamelCase   :   Union[str, Any]       ='''_orig_mod.'''
							for k, v in list(state_dict.items() ):
														if k.startswith(SCREAMING_SNAKE_CASE ):
																					# replace part of the key with corresponding layer name in HF implementation
																					__lowerCamelCase   :   List[Any]       =k[len(SCREAMING_SNAKE_CASE ) :]
																					for old_layer_name in new_layer_name_dict:
																												__lowerCamelCase   :   Tuple       =new_k.replace(SCREAMING_SNAKE_CASE     ,    new_layer_name_dict[old_layer_name] )

																					__lowerCamelCase   :   Dict       =state_dict.pop(SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   Any       =set(state_dict.keys() ) - set(model.state_dict().keys() )
							__lowerCamelCase   :   Optional[int]       ={k for k in extra_keys if not k.endswith('''.attn.bias''' )}
							__lowerCamelCase   :   Optional[int]       =set(model.state_dict().keys() ) - set(state_dict.keys() )
							__lowerCamelCase   :   List[Any]       ={k for k in missing_keys if not k.endswith('''.attn.bias''' )}
							if len(SCREAMING_SNAKE_CASE ) != 0:
														raise ValueError(F'extra keys found: {extra_keys}' )
							if len(SCREAMING_SNAKE_CASE ) != 0:
														raise ValueError(F'missing keys: {missing_keys}' )
							model.load_state_dict(SCREAMING_SNAKE_CASE     ,    strict=SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   int       =model.num_parameters(exclude_embeddings=SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   Union[str, Any]       =checkpoint['''best_val_loss'''].item()
							logger.info(F'model loaded: {round(n_params/1E6     ,    1 )}M params, {round(SCREAMING_SNAKE_CASE     ,    3 )} loss' )
							model.eval()
							model.to(SCREAMING_SNAKE_CASE )
							del checkpoint, state_dict

							return model



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Optional[int]     ,    SCREAMING_SNAKE_CASE							:				Optional[int]=False     ,    SCREAMING_SNAKE_CASE							:				Optional[Any]="text" ):


							'''simple docstring'''







							if model_type not in ("text", "coarse", "fine"):
														raise NotImplementedError()

							__lowerCamelCase   :   Tuple       ='''cpu'''  # do conversion on cpu

							__lowerCamelCase   :   Dict       =_get_ckpt_path(SCREAMING_SNAKE_CASE     ,    use_small=SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   str       =_load_model(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    model_type=SCREAMING_SNAKE_CASE     ,    use_small=SCREAMING_SNAKE_CASE )

							# load bark initial model
							__lowerCamelCase   :   List[str]       =_bark_load_model(SCREAMING_SNAKE_CASE     ,    '''cpu'''     ,    model_type=SCREAMING_SNAKE_CASE     ,    use_small=SCREAMING_SNAKE_CASE )

							if model_type == "text":
														__lowerCamelCase   :   int       =bark_model['''model''']

							if model.num_parameters(exclude_embeddings=SCREAMING_SNAKE_CASE ) != bark_model.get_num_params():
														raise ValueError('''initial and new models don\'t have the same number of parameters''' )

							# check if same output as the bark model
							__lowerCamelCase   :   Dict       =5
							__lowerCamelCase   :   Optional[int]       =10

							if model_type in ["text", "coarse"]:
														__lowerCamelCase   :   List[Any]       =torch.randint(256     ,    (batch_size, sequence_length)     ,    dtype=torch.int )
														__lowerCamelCase   :   int       =bark_model(SCREAMING_SNAKE_CASE )[0]

														__lowerCamelCase   :   List[Any]       =model(SCREAMING_SNAKE_CASE )

														# take last logits
														__lowerCamelCase   :   str       =output_new_model_total.logits[:, [-1], :]

							else:
														__lowerCamelCase   :   str       =3
														__lowerCamelCase   :   Union[str, Any]       =8
														__lowerCamelCase   :   Union[str, Any]       =torch.randint(256     ,    (batch_size, sequence_length, n_codes_total)     ,    dtype=torch.int )

														__lowerCamelCase   :   List[Any]       =model(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )
														__lowerCamelCase   :   int       =bark_model(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

														__lowerCamelCase   :   Dict       =output_new_model_total.logits

							# output difference should come from the difference of self-attention implementation design
							if output_new_model.shape != output_old_model.shape:
														raise ValueError('''initial and new outputs don\'t have the same shape''' )
							if (output_new_model - output_old_model).abs().max().item() > 1E-3:
														raise ValueError('''initial and new outputs are not equal''' )

							Path(SCREAMING_SNAKE_CASE ).mkdir(exist_ok=SCREAMING_SNAKE_CASE )
							model.save_pretrained(SCREAMING_SNAKE_CASE )



def 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Union[str, Any]     ,    SCREAMING_SNAKE_CASE							:				str     ,    SCREAMING_SNAKE_CASE							:				Any     ,    SCREAMING_SNAKE_CASE							:				Union[str, Any]     ,    SCREAMING_SNAKE_CASE							:				List[str]     ,    SCREAMING_SNAKE_CASE							:				Any     ,    ):


							'''simple docstring'''







							__lowerCamelCase   :   str       =os.path.join(SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   List[Any]       =BarkSemanticConfig.from_pretrained(os.path.join(SCREAMING_SNAKE_CASE     ,    '''config.json''' ) )
							__lowerCamelCase   :   Any       =BarkCoarseConfig.from_pretrained(os.path.join(SCREAMING_SNAKE_CASE     ,    '''config.json''' ) )
							__lowerCamelCase   :   List[str]       =BarkFineConfig.from_pretrained(os.path.join(SCREAMING_SNAKE_CASE     ,    '''config.json''' ) )
							__lowerCamelCase   :   Optional[Any]       =EncodecConfig.from_pretrained('''facebook/encodec_24khz''' )

							__lowerCamelCase   :   Any       =BarkSemanticModel.from_pretrained(SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   Optional[int]       =BarkCoarseModel.from_pretrained(SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   Optional[int]       =BarkFineModel.from_pretrained(SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   List[Any]       =EncodecModel.from_pretrained('''facebook/encodec_24khz''' )

							__lowerCamelCase   :   Any       =BarkConfig.from_sub_model_configs(
							    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE     ,    SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   Tuple       =BarkGenerationConfig.from_sub_model_configs(
							    semantic.generation_config     ,    coarseAcoustic.generation_config     ,    fineAcoustic.generation_config )

							__lowerCamelCase   :   Optional[int]       =BarkModel(SCREAMING_SNAKE_CASE )

							__lowerCamelCase   :   List[Any]       =semantic
							__lowerCamelCase   :   Union[str, Any]       =coarseAcoustic
							__lowerCamelCase   :   List[Any]       =fineAcoustic
							__lowerCamelCase   :   str       =codec

							__lowerCamelCase   :   int       =bark_generation_config

							Path(SCREAMING_SNAKE_CASE ).mkdir(exist_ok=SCREAMING_SNAKE_CASE )
							bark.save_pretrained(SCREAMING_SNAKE_CASE     ,    repo_id=SCREAMING_SNAKE_CASE     ,    push_to_hub=SCREAMING_SNAKE_CASE )


if __name__ == "__main__":
							_UpperCamelCase					      =   argparse.ArgumentParser()
							# Required parameters

							parser.add_argument('model_type', type=str, help='text, coarse or fine.')
							parser.add_argument('pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
							parser.add_argument('--is_small', action='store_true', help='convert the small version instead of the large.')

							_UpperCamelCase					      =   parser.parse_args()

							load_model(args.pytorch_dump_folder_path, model_type=args.model_type, use_small=args.is_small)











"""simple docstring"""





from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


_UpperCamelCase					      =   logging.get_logger(__name__)

_UpperCamelCase					      =   {
    'Helsinki-NLP/opus-mt-en-de': 'https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json',
    # See all Marian models at https://huggingface.co/models?filter=marian
}
class 		SCREAMING_SNAKE_CASE_				(     snake_case__ ):
				"""simple docstring"""
				__snake_case :							Tuple									=		"""marian"""
				__snake_case :							Any									=		["""past_key_values"""]
				__snake_case :							Optional[Any]									=		{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}




				def __init__(  self	:List[Any]    ,     __lowercase	:Any=5_8101    ,     __lowercase	:Tuple=None    ,     __lowercase	:Union[str, Any]=1024    ,     __lowercase	:Dict=12    ,     __lowercase	:int=4096    ,     __lowercase	:int=16    ,     __lowercase	:List[Any]=12    ,     __lowercase	:Dict=4096    ,     __lowercase	:Dict=16    ,     __lowercase	:Tuple=0.0    ,     __lowercase	:Tuple=0.0    ,     __lowercase	:List[Any]=True    ,     __lowercase	:int=True    ,     __lowercase	:Tuple="gelu"    ,     __lowercase	:str=1024    ,     __lowercase	:Optional[int]=0.1    ,     __lowercase	:List[str]=0.0    ,     __lowercase	:Union[str, Any]=0.0    ,     __lowercase	:Dict=0.02    ,     __lowercase	:Tuple=5_8100    ,     __lowercase	:Optional[Any]=False    ,     __lowercase	:int=5_8100    ,     __lowercase	:Any=0    ,     __lowercase	:str=0    ,     __lowercase	:str=True    ,     **__lowercase	:Any    ,     ):
											__lowerCamelCase   :   List[Any]       =vocab_size
											__lowerCamelCase   :   Optional[int]       =decoder_vocab_size or vocab_size
											__lowerCamelCase   :   Tuple       =max_position_embeddings
											__lowerCamelCase   :   List[Any]       =d_model
											__lowerCamelCase   :   Any       =encoder_ffn_dim
											__lowerCamelCase   :   str       =encoder_layers
											__lowerCamelCase   :   List[str]       =encoder_attention_heads
											__lowerCamelCase   :   str       =decoder_ffn_dim
											__lowerCamelCase   :   Tuple       =decoder_layers
											__lowerCamelCase   :   Any       =decoder_attention_heads
											__lowerCamelCase   :   List[Any]       =dropout
											__lowerCamelCase   :   Any       =attention_dropout
											__lowerCamelCase   :   Union[str, Any]       =activation_dropout
											__lowerCamelCase   :   Optional[int]       =activation_function
											__lowerCamelCase   :   Dict       =init_std
											__lowerCamelCase   :   List[Any]       =encoder_layerdrop
											__lowerCamelCase   :   Optional[Any]       =decoder_layerdrop
											__lowerCamelCase   :   Any       =use_cache
											__lowerCamelCase   :   Any       =encoder_layers
											__lowerCamelCase   :   Optional[int]       =scale_embedding  # scale factor will be sqrt(d_model) if True
											__lowerCamelCase   :   int       =share_encoder_decoder_embeddings
											super().__init__(
											    pad_token_id=__lowercase    ,     eos_token_id=__lowercase    ,     is_encoder_decoder=__lowercase    ,     decoder_start_token_id=__lowercase    ,     forced_eos_token_id=__lowercase    ,     **__lowercase    ,     )

class 		SCREAMING_SNAKE_CASE_				(     snake_case__ ):
				"""simple docstring"""




				@property
				# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
				def        __lowercase	(  self	:int      ):
											if self.task in ["default", "seq2seq-lm"]:
																		__lowerCamelCase   :   Tuple       =OrderedDict(
																		    [
																		        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
																		        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
																		    ]      )

																		if self.use_past:
																									__lowerCamelCase   :   Optional[Any]       ={0: '''batch'''}
																									__lowerCamelCase   :   Any       ={0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
																		else:
																									__lowerCamelCase   :   Optional[Any]       ={0: '''batch''', 1: '''decoder_sequence'''}
																									__lowerCamelCase   :   List[Any]       ={0: '''batch''', 1: '''decoder_sequence'''}

																		if self.use_past:
																									self.fill_with_past_key_values_(__lowercase    ,     direction='''inputs'''      )
											elif self.task == "causal-lm":
																		# TODO: figure this case out.
																		__lowerCamelCase   :   Any       =OrderedDict(
																		    [
																		        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
																		        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
																		    ]      )
																		if self.use_past:
																									__lowerCamelCase     ,    __lowerCamelCase   :   str       =self.num_layers
																									for i in range(__lowercase      ):
																																__lowerCamelCase   :   Optional[int]       ={0: '''batch''', 2: '''past_sequence + sequence'''}
																																__lowerCamelCase   :   List[str]       ={0: '''batch''', 2: '''past_sequence + sequence'''}
											else:
																		__lowerCamelCase   :   List[str]       =OrderedDict(
																		    [
																		        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
																		        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
																		        ('''decoder_input_ids''', {0: '''batch''', 1: '''decoder_sequence'''}),
																		        ('''decoder_attention_mask''', {0: '''batch''', 1: '''decoder_sequence'''}),
																		    ]      )

											return common_inputs




				@property
				# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
				def        __lowercase	(  self	:Optional[int]      ):
											if self.task in ["default", "seq2seq-lm"]:
																		__lowerCamelCase   :   Tuple       =super().outputs
											else:
																		__lowerCamelCase   :   List[str]       =super(__lowercase    ,     self      ).outputs
																		if self.use_past:
																									__lowerCamelCase     ,    __lowerCamelCase   :   int       =self.num_layers
																									for i in range(__lowercase      ):
																																__lowerCamelCase   :   Any       ={0: '''batch''', 2: '''past_sequence + sequence'''}
																																__lowerCamelCase   :   Optional[int]       ={0: '''batch''', 2: '''past_sequence + sequence'''}
											return common_outputs




				def        __lowercase	(  self	:str    ,     __lowercase	:PreTrainedTokenizer    ,     __lowercase	:int = -1    ,     __lowercase	:int = -1    ,     __lowercase	:bool = False    ,     __lowercase	:Optional[TensorType] = None    ,     ):
											__lowerCamelCase   :   List[str]       =self._generate_dummy_inputs_for_encoder_and_decoder(
											    __lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase      )

											# Generate decoder inputs
											__lowerCamelCase   :   Optional[Any]       =seq_length if not self.use_past else 1
											__lowerCamelCase   :   List[Any]       =self._generate_dummy_inputs_for_encoder_and_decoder(
											    __lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase      )
											__lowerCamelCase   :   Dict       ={f'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
											__lowerCamelCase   :   str       =dict(**__lowercase    ,     **__lowercase      )

											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
																		__lowerCamelCase   :   Optional[Any]       =common_inputs['''decoder_input_ids'''].shape[1]
																		__lowerCamelCase     ,    __lowerCamelCase   :   Optional[int]       =self.num_attention_heads
																		__lowerCamelCase   :   Any       =(
																		    batch,
																		    num_encoder_attention_heads,
																		    encoder_seq_length,
																		    self._config.hidden_size // num_encoder_attention_heads,
																		)
																		__lowerCamelCase   :   Dict       =decoder_seq_length + 3
																		__lowerCamelCase   :   Tuple       =(
																		    batch,
																		    num_decoder_attention_heads,
																		    decoder_past_length,
																		    self._config.hidden_size // num_decoder_attention_heads,
																		)

																		__lowerCamelCase   :   Optional[int]       =torch.cat(
																		    [common_inputs['''decoder_attention_mask'''], torch.ones(__lowercase    ,     __lowercase      )]    ,     dim=1      )

																		__lowerCamelCase   :   Any       =[]
																		# If the number of encoder and decoder layers are present in the model configuration, both are considered
																		__lowerCamelCase     ,    __lowerCamelCase   :   str       =self.num_layers
																		__lowerCamelCase   :   List[Any]       =min(__lowercase    ,     __lowercase      )
																		__lowerCamelCase   :   int       =max(__lowercase    ,     __lowercase      ) - min_num_layers
																		__lowerCamelCase   :   Any       ='''encoder''' if num_encoder_layers > num_decoder_layers else '''decoder'''

																		for _ in range(__lowercase      ):
																									common_inputs["past_key_values"].append(
																									    (
																									        torch.zeros(__lowercase      ),
																									        torch.zeros(__lowercase      ),
																									        torch.zeros(__lowercase      ),
																									        torch.zeros(__lowercase      ),
																									    )      )
																		# TODO: test this.
																		__lowerCamelCase   :   Dict       =encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
																		for _ in range(__lowercase    ,     __lowercase      ):
																									common_inputs["past_key_values"].append((torch.zeros(__lowercase      ), torch.zeros(__lowercase      ))      )
											return common_inputs




				def        __lowercase	(  self	:Dict    ,     __lowercase	:PreTrainedTokenizer    ,     __lowercase	:int = -1    ,     __lowercase	:int = -1    ,     __lowercase	:bool = False    ,     __lowercase	:Optional[TensorType] = None    ,     ):
											__lowerCamelCase   :   List[str]       =self._generate_dummy_inputs_for_encoder_and_decoder(
											    __lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase      )

											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   :   str       =seqlen + 2
																		__lowerCamelCase     ,    __lowerCamelCase   :   List[str]       =self.num_layers
																		__lowerCamelCase     ,    __lowerCamelCase   :   Optional[Any]       =self.num_attention_heads
																		__lowerCamelCase   :   List[str]       =(
																		    batch,
																		    num_encoder_attention_heads,
																		    past_key_values_length,
																		    self._config.hidden_size // num_encoder_attention_heads,
																		)

																		__lowerCamelCase   :   Any       =common_inputs['''attention_mask'''].dtype
																		__lowerCamelCase   :   Optional[int]       =torch.cat(
																		    [common_inputs['''attention_mask'''], torch.ones(__lowercase    ,     __lowercase    ,     dtype=__lowercase      )]    ,     dim=1      )
																		__lowerCamelCase   :   List[str]       =[
																		    (torch.zeros(__lowercase      ), torch.zeros(__lowercase      )) for _ in range(__lowercase      )
																		]
											return common_inputs




				def        __lowercase	(  self	:int    ,     __lowercase	:PreTrainedTokenizer    ,     __lowercase	:int = -1    ,     __lowercase	:int = -1    ,     __lowercase	:bool = False    ,     __lowercase	:Optional[TensorType] = None    ,     ):
											# Copied from OnnxConfig.generate_dummy_inputs
											# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
											# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
											__lowerCamelCase   :   Union[str, Any]       =compute_effective_axis_dimension(
											    __lowercase    ,     fixed_dimension=OnnxConfig.default_fixed_batch    ,     num_token_to_add=0      )

											# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
											__lowerCamelCase   :   Optional[int]       =tokenizer.num_special_tokens_to_add(__lowercase      )
											__lowerCamelCase   :   Optional[int]       =compute_effective_axis_dimension(
											    __lowercase    ,     fixed_dimension=OnnxConfig.default_fixed_sequence    ,     num_token_to_add=__lowercase      )

											# Generate dummy inputs according to compute batch and sequence
											__lowerCamelCase   :   Any       =[''' '''.join([tokenizer.unk_token]      ) * seq_length] * batch_size
											__lowerCamelCase   :   Union[str, Any]       =dict(tokenizer(__lowercase    ,     return_tensors=__lowercase      )      )
											return common_inputs




				def        __lowercase	(  self	:Tuple    ,     __lowercase	:PreTrainedTokenizer    ,     __lowercase	:int = -1    ,     __lowercase	:int = -1    ,     __lowercase	:bool = False    ,     __lowercase	:Optional[TensorType] = None    ,     ):
											if self.task in ["default", "seq2seq-lm"]:
																		__lowerCamelCase   :   int       =self._generate_dummy_inputs_for_default_and_seqaseq_lm(
																		    __lowercase    ,     batch_size=__lowercase    ,     seq_length=__lowercase    ,     is_pair=__lowercase    ,     framework=__lowercase      )

											else:
																		__lowerCamelCase   :   List[Any]       =self._generate_dummy_inputs_for_causal_lm(
																		    __lowercase    ,     batch_size=__lowercase    ,     seq_length=__lowercase    ,     is_pair=__lowercase    ,     framework=__lowercase      )

											return common_inputs




				def        __lowercase	(  self	:Optional[int]    ,     __lowercase	:Tuple    ,     __lowercase	:Tuple    ,     __lowercase	:Dict    ,     __lowercase	:List[str]      ):
											if self.task in ["default", "seq2seq-lm"]:
																		__lowerCamelCase   :   str       =super()._flatten_past_key_values_(__lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase      )
											else:
																		__lowerCamelCase   :   Optional[Any]       =super(__lowercase    ,     self      )._flatten_past_key_values_(
																		    __lowercase    ,     __lowercase    ,     __lowercase    ,     __lowercase      )




				@property
				def        __lowercase	(  self	:List[str]      ):
											return 1e-4




