




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


__a			:	Tuple							     =		logging.get_logger(__name__)

__a			:	List[Any]							     =		{
    """facebook/mbart-large-cc25""": """https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json""",
    # See all MBART models at https://huggingface.co/models?filter=mbart
}




class     A					(				lowerCamelCase_       ):
							_SCREAMING_SNAKE_CASE  :		List[Any]			  =       '''mbart'''
							_SCREAMING_SNAKE_CASE  :		int			  =       ['''past_key_values''']
							_SCREAMING_SNAKE_CASE  :		List[str]			  =       {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}







							def __init__(    self   :	Tuple   ,     __UpperCAmelCase   :	Tuple=50265   ,     __UpperCAmelCase   :	List[str]=1024   ,     __UpperCAmelCase   :	List[Any]=12   ,     __UpperCAmelCase   :	List[str]=4096   ,     __UpperCAmelCase   :	List[str]=16   ,     __UpperCAmelCase   :	int=12   ,     __UpperCAmelCase   :	Union[str, Any]=4096   ,     __UpperCAmelCase   :	Any=16   ,     __UpperCAmelCase   :	Tuple=0.0   ,     __UpperCAmelCase   :	Union[str, Any]=0.0   ,     __UpperCAmelCase   :	Dict=True   ,     __UpperCAmelCase   :	Dict=True   ,     __UpperCAmelCase   :	Any="gelu"   ,     __UpperCAmelCase   :	Any=1024   ,     __UpperCAmelCase   :	int=0.1   ,     __UpperCAmelCase   :	Any=0.0   ,     __UpperCAmelCase   :	Optional[int]=0.0   ,     __UpperCAmelCase   :	Any=0.02   ,     __UpperCAmelCase   :	Optional[Any]=0.0   ,     __UpperCAmelCase   :	List[Any]=False   ,     __UpperCAmelCase   :	int=1   ,     __UpperCAmelCase   :	int=0   ,     __UpperCAmelCase   :	Union[str, Any]=2   ,     __UpperCAmelCase   :	Optional[int]=2   ,     **__UpperCAmelCase   :	Optional[Any]   ,     )		->					List[str]:

										"""simple docstring"""



										UpperCamelCase_     				= vocab_size
										UpperCamelCase_     				= max_position_embeddings
										UpperCamelCase_     				= d_model
										UpperCamelCase_     				= encoder_ffn_dim
										UpperCamelCase_     				= encoder_layers
										UpperCamelCase_     				= encoder_attention_heads
										UpperCamelCase_     				= decoder_ffn_dim
										UpperCamelCase_     				= decoder_layers
										UpperCamelCase_     				= decoder_attention_heads
										UpperCamelCase_     				= dropout
										UpperCamelCase_     				= attention_dropout
										UpperCamelCase_     				= activation_dropout
										UpperCamelCase_     				= activation_function
										UpperCamelCase_     				= init_std
										UpperCamelCase_     				= encoder_layerdrop
										UpperCamelCase_     				= decoder_layerdrop
										UpperCamelCase_     				= classifier_dropout
										UpperCamelCase_     				= use_cache
										UpperCamelCase_     				= encoder_layers
										UpperCamelCase_     				= scale_embedding  # scale factor will be sqrt(d_model) if True
										super().__init__(
										    pad_token_id=__UpperCAmelCase   ,     bos_token_id=__UpperCAmelCase   ,     eos_token_id=__UpperCAmelCase   ,     is_encoder_decoder=__UpperCAmelCase   ,     forced_eos_token_id=__UpperCAmelCase   ,     **__UpperCAmelCase   ,     )







class     A					(				lowerCamelCase_       ):







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

										"""simple docstring"""



										if self.task in ["default", "seq2seq-lm"]:
													UpperCamelCase_     				= OrderedDict(
													    [
													        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
													        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
													    ]				)

													if self.use_past:
																UpperCamelCase_     				= {0: 'batch'}
																UpperCamelCase_     				= {0: 'batch', 1: 'past_decoder_sequence + sequence'}
													else:
																UpperCamelCase_     				= {0: 'batch', 1: 'decoder_sequence'}
																UpperCamelCase_     				= {0: 'batch', 1: 'decoder_sequence'}

													if self.use_past:
																self.fill_with_past_key_values_(__UpperCAmelCase   ,     direction='inputs'				)
										elif self.task == "causal-lm":
													# TODO: figure this case out.
													UpperCamelCase_     				= OrderedDict(
													    [
													        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
													        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
													    ]				)
													if self.use_past:
																UpperCamelCase_				,   UpperCamelCase_     				= self.num_layers
																for i in range(__UpperCAmelCase				):
																			UpperCamelCase_     				= {0: 'batch', 2: 'past_sequence + sequence'}
																			UpperCamelCase_     				= {0: 'batch', 2: 'past_sequence + sequence'}
										else:
													UpperCamelCase_     				= 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
							def        lowercase__							(    self   :	Optional[Any]				)		->					Mapping[str, Mapping[int, str]]:

										"""simple docstring"""



										if self.task in ["default", "seq2seq-lm"]:
													UpperCamelCase_     				= super().outputs
										else:
													UpperCamelCase_     				= super(__UpperCAmelCase   ,     self				).outputs
													if self.use_past:
																UpperCamelCase_				,   UpperCamelCase_     				= self.num_layers
																for i in range(__UpperCAmelCase				):
																			UpperCamelCase_     				= {0: 'batch', 2: 'past_sequence + sequence'}
																			UpperCamelCase_     				= {0: 'batch', 2: 'past_sequence + sequence'}
										return common_outputs







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

										"""simple docstring"""



										UpperCamelCase_     				= self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase				)

										# Generate decoder inputs
										UpperCamelCase_     				= seq_length if not self.use_past else 1
										UpperCamelCase_     				= self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase				)
										UpperCamelCase_     				= {f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
										UpperCamelCase_     				= dict(**__UpperCAmelCase   ,     **__UpperCAmelCase				)

										if self.use_past:
													if not is_torch_available():
																raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'				)
													else:
																import torch
													UpperCamelCase_				,   UpperCamelCase_     				= common_inputs['input_ids'].shape
													UpperCamelCase_     				= common_inputs['decoder_input_ids'].shape[1]
													UpperCamelCase_				,   UpperCamelCase_     				= self.num_attention_heads
													UpperCamelCase_     				= (
													    batch,
													    num_encoder_attention_heads,
													    encoder_seq_length,
													    self._config.hidden_size // num_encoder_attention_heads,
													)
													UpperCamelCase_     				= decoder_seq_length + 3
													UpperCamelCase_     				= (
													    batch,
													    num_decoder_attention_heads,
													    decoder_past_length,
													    self._config.hidden_size // num_decoder_attention_heads,
													)

													UpperCamelCase_     				= torch.cat(
													    [common_inputs['decoder_attention_mask'], torch.ones(__UpperCAmelCase   ,     __UpperCAmelCase				)]   ,     dim=1				)

													UpperCamelCase_     				= []
													# If the number of encoder and decoder layers are present in the model configuration, both are considered
													UpperCamelCase_				,   UpperCamelCase_     				= self.num_layers
													UpperCamelCase_     				= min(__UpperCAmelCase   ,     __UpperCAmelCase				)
													UpperCamelCase_     				= max(__UpperCAmelCase   ,     __UpperCAmelCase				) - min_num_layers
													UpperCamelCase_     				= 'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

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







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

										"""simple docstring"""



										UpperCamelCase_     				= self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase				)

										if self.use_past:
													if not is_torch_available():
																raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'				)
													else:
																import torch
													UpperCamelCase_				,   UpperCamelCase_     				= common_inputs['input_ids'].shape
													# Not using the same length for past_key_values
													UpperCamelCase_     				= seqlen + 2
													UpperCamelCase_				,   UpperCamelCase_     				= self.num_layers
													UpperCamelCase_				,   UpperCamelCase_     				= self.num_attention_heads
													UpperCamelCase_     				= (
													    batch,
													    num_encoder_attention_heads,
													    past_key_values_length,
													    self._config.hidden_size // num_encoder_attention_heads,
													)

													UpperCamelCase_     				= common_inputs['attention_mask'].dtype
													UpperCamelCase_     				= torch.cat(
													    [common_inputs['attention_mask'], torch.ones(__UpperCAmelCase   ,     __UpperCAmelCase   ,     dtype=__UpperCAmelCase				)]   ,     dim=1				)
													UpperCamelCase_     				= [
													    (torch.zeros(__UpperCAmelCase				), torch.zeros(__UpperCAmelCase				)) for _ in range(__UpperCAmelCase				)
													]
										return common_inputs







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

										"""simple docstring"""



										UpperCamelCase_     				= compute_effective_axis_dimension(
										    __UpperCAmelCase   ,     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
										UpperCamelCase_     				= tokenizer.num_special_tokens_to_add(__UpperCAmelCase				)
										UpperCamelCase_     				= compute_effective_axis_dimension(
										    __UpperCAmelCase   ,     fixed_dimension=OnnxConfig.default_fixed_sequence   ,     num_token_to_add=__UpperCAmelCase				)

										# Generate dummy inputs according to compute batch and sequence
										UpperCamelCase_     				= [' '.join([tokenizer.unk_token]				) * seq_length] * batch_size
										UpperCamelCase_     				= dict(tokenizer(__UpperCAmelCase   ,     return_tensors=__UpperCAmelCase				)				)
										return common_inputs







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

										"""simple docstring"""



										if self.task in ["default", "seq2seq-lm"]:
													UpperCamelCase_     				= self._generate_dummy_inputs_for_default_and_seqaseq_lm(
													    __UpperCAmelCase   ,     batch_size=__UpperCAmelCase   ,     seq_length=__UpperCAmelCase   ,     is_pair=__UpperCAmelCase   ,     framework=__UpperCAmelCase				)

										elif self.task == "causal-lm":
													UpperCamelCase_     				= self._generate_dummy_inputs_for_causal_lm(
													    __UpperCAmelCase   ,     batch_size=__UpperCAmelCase   ,     seq_length=__UpperCAmelCase   ,     is_pair=__UpperCAmelCase   ,     framework=__UpperCAmelCase				)
										else:
													UpperCamelCase_     				= self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    __UpperCAmelCase   ,     batch_size=__UpperCAmelCase   ,     seq_length=__UpperCAmelCase   ,     is_pair=__UpperCAmelCase   ,     framework=__UpperCAmelCase				)

										return common_inputs







							def        lowercase__							(    self   :	Any   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	Any   ,     __UpperCAmelCase   :	Union[str, Any]   ,     __UpperCAmelCase   :	Any				)		->					Optional[int]:

										"""simple docstring"""



										if self.task in ["default", "seq2seq-lm"]:
													UpperCamelCase_     				= super()._flatten_past_key_values_(__UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase				)
										else:
													UpperCamelCase_     				= super(__UpperCAmelCase   ,     self				)._flatten_past_key_values_(
													    __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase   ,     __UpperCAmelCase				)






import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case=None					)	->			str:

			'''simple docstring'''


			UpperCamelCase_     				= None
			# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
			if str(weight.dtype					) == "torch.bfloat16":
						# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
						UpperCamelCase_     				= weight.view(torch.intaa					)
						UpperCamelCase_     				= 'bfloat16'
			UpperCamelCase_     				= weight.cpu().numpy()
			UpperCamelCase_     				= os.path.join(__snake_case     ,		F'''{weight_name}.dat'''					)
			if index is not None:
						if dtype is None:
									UpperCamelCase_     				= str(array.dtype					)
						UpperCamelCase_     				= {'dtype': dtype, 'shape': list(array.shape					)}
			if array.ndim == 0:
						UpperCamelCase_     				= array[None]
			UpperCamelCase_     				= np.memmap(__snake_case     ,		dtype=array.dtype     ,		mode='w+'     ,		shape=array.shape					)
			UpperCamelCase_     				= array[:]
			file_array.flush()
			return index




def  a_		(							__snake_case     ,		__snake_case					)	->			Dict:

			'''simple docstring'''


			UpperCamelCase_     				= tuple(weight_info['shape']					)
			if shape == ():
						# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
						UpperCamelCase_     				= (1,)

			UpperCamelCase_     				= weight_info['dtype']
			if dtype == "bfloat16":
						# NumPy does not support bfloat16 so this was saved as a int16
						UpperCamelCase_     				= 'int16'

			UpperCamelCase_     				= np.memmap(__snake_case     ,		dtype=__snake_case     ,		shape=__snake_case     ,		mode='r'					)

			if len(weight_info['shape']					) == 0:
						UpperCamelCase_     				= weight[0]
			UpperCamelCase_     				= torch.tensor(__snake_case					)
			if weight_info["dtype"] == "bfloat16":
						UpperCamelCase_     				= weight.view(torch.bfloataa					)

			return weight




def  a_		(							__snake_case     ,		__snake_case					)	->			Optional[int]:

			'''simple docstring'''


			if index is None or len(__snake_case					) == 0:
						# Nothing to save
						return

			UpperCamelCase_     				= os.path.join(__snake_case     ,		'index.json'					)
			if os.path.isfile(__snake_case					):
						with open(__snake_case     ,		'r'     ,		encoding='utf-8'					) as f:
									UpperCamelCase_     				= json.load(__snake_case					)
			else:
						UpperCamelCase_     				= {}
			current_index.update(__snake_case					)

			with open(__snake_case     ,		'w'     ,		encoding='utf-8'					) as f:
						json.dump(__snake_case     ,		__snake_case     ,		indent=2					)




def  a_		(							__snake_case     ,		__snake_case					)	->			Tuple:

			'''simple docstring'''


			os.makedirs(__snake_case     ,		exist_ok=__snake_case					)
			UpperCamelCase_     				= {}
			for name, parameter in state_dict.items():
						UpperCamelCase_     				= offload_weight(__snake_case     ,		__snake_case     ,		__snake_case     ,		index=__snake_case					)

			# Update index
			save_offload_index(__snake_case     ,		__snake_case					)




class     A					(				lowerCamelCase_       ):







							def __init__(    self   :	int   ,     __UpperCAmelCase   :	Mapping   ,     __UpperCAmelCase   :	str				)		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= dataset
										UpperCamelCase_     				= prefix







							def __getitem__(    self   :	List[Any]   ,     __UpperCAmelCase   :	Dict				)		->					Tuple:

										"""simple docstring"""



										return self.dataset[f'''{self.prefix}{key}''']







							def __iter__(    self   :	int				)		->					List[str]:

										"""simple docstring"""



										return iter([key for key in self.dataset if key.startswith(self.prefix				)]				)







							def __len__(    self   :	Any				)		->					Optional[Any]:

										"""simple docstring"""



										return len(self.dataset				)







class     A					(				lowerCamelCase_       ):







							def __init__(    self   :	int   ,     __UpperCAmelCase   :	Dict[str, torch.Tensor] = None   ,     __UpperCAmelCase   :	Optional[Union[str, os.PathLike]] = None   ,     __UpperCAmelCase   :	Mapping = None   ,     __UpperCAmelCase   :	Optional[int]=None   ,     )		->					List[str]:

										"""simple docstring"""



										if state_dict is None and save_folder is None:
													raise ValueError('Need either a `state_dict` or a `save_folder` containing offloaded weights.'				)

										UpperCamelCase_     				= {} if state_dict is None else state_dict
										UpperCamelCase_     				= save_folder
										if index is None and save_folder is not None:
													with open(os.path.join(__UpperCAmelCase   ,     'index.json'				)				) as f:
																UpperCamelCase_     				= json.load(__UpperCAmelCase				)
										UpperCamelCase_     				= {} if index is None else index
										UpperCamelCase_     				= list(self.state_dict.keys()				)
										self.all_keys.extend([key for key in self.index if key not in self.all_keys]				)
										UpperCamelCase_     				= device







							def __getitem__(    self   :	Optional[Any]   ,     __UpperCAmelCase   :	str				)		->					Tuple:

										"""simple docstring"""



										if key in self.state_dict:
													return self.state_dict[key]
										UpperCamelCase_     				= self.index[key]
										if weight_info.get('safetensors_file'				) is not None:
													if not is_safetensors_available():
																raise ImportError('These offloaded weights require the use of safetensors: `pip install safetensors`.'				)

													from safetensors import safe_open

													UpperCamelCase_     				= 'cpu' if self.device is None else self.device
													with safe_open(weight_info['safetensors_file']   ,     framework='pt'   ,     device=__UpperCAmelCase				) as f:
																UpperCamelCase_     				= f.get_tensor(weight_info.get('weight_name'   ,     __UpperCAmelCase				)				)

													if "dtype" in weight_info:
																return tensor.to(getattr(__UpperCAmelCase   ,     weight_info['dtype']				)				)
													else:
																return tensor

										UpperCamelCase_     				= os.path.join(self.save_folder   ,     f'''{key}.dat'''				)
										return load_offloaded_weight(__UpperCAmelCase   ,     __UpperCAmelCase				)







							def __iter__(    self   :	Optional[int]				)		->					Union[str, Any]:

										"""simple docstring"""



										return iter(self.all_keys				)







							def __len__(    self   :	Any				)		->					Any:

										"""simple docstring"""



										return len(self.all_keys				)







def  a_		(							__snake_case     ,		__snake_case					)	->			Any:

			'''simple docstring'''


			UpperCamelCase_     				= {}
			for module_name in submodule_names:
						# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
						# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
						result.update(
						    {
						        key: param
						        for key, param in state_dict.items()
						        if key == module_name or key.startswith(module_name + '.'					)
						    }					)
			return result
