
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


lowercase__							:							Dict													=     logging.get_logger(__name__)

lowercase__							:							Union[str, 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     SCREAMING_SNAKE_CASE__      (							_a      ):

							"""simple docstring"""






							_snake_case       						=					'mbart'
							_snake_case       						=					['past_key_values']
							_snake_case       						=					{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}



							def __init__(					self       ,		SCREAMING_SNAKE_CASE_=50265       ,		SCREAMING_SNAKE_CASE_=1024       ,		SCREAMING_SNAKE_CASE_=12       ,		SCREAMING_SNAKE_CASE_=4096       ,		SCREAMING_SNAKE_CASE_=16       ,		SCREAMING_SNAKE_CASE_=12       ,		SCREAMING_SNAKE_CASE_=4096       ,		SCREAMING_SNAKE_CASE_=16       ,		SCREAMING_SNAKE_CASE_=0.0       ,		SCREAMING_SNAKE_CASE_=0.0       ,		SCREAMING_SNAKE_CASE_=True       ,		SCREAMING_SNAKE_CASE_=True       ,		SCREAMING_SNAKE_CASE_="gelu"       ,		SCREAMING_SNAKE_CASE_=1024       ,		SCREAMING_SNAKE_CASE_=0.1       ,		SCREAMING_SNAKE_CASE_=0.0       ,		SCREAMING_SNAKE_CASE_=0.0       ,		SCREAMING_SNAKE_CASE_=0.0_2       ,		SCREAMING_SNAKE_CASE_=0.0       ,		SCREAMING_SNAKE_CASE_=False       ,		SCREAMING_SNAKE_CASE_=1       ,		SCREAMING_SNAKE_CASE_=0       ,		SCREAMING_SNAKE_CASE_=2       ,		SCREAMING_SNAKE_CASE_=2       ,		**SCREAMING_SNAKE_CASE_       ,		)->				Tuple:





													'''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=SCREAMING_SNAKE_CASE_       ,		bos_token_id=SCREAMING_SNAKE_CASE_       ,		eos_token_id=SCREAMING_SNAKE_CASE_       ,		is_encoder_decoder=SCREAMING_SNAKE_CASE_       ,		forced_eos_token_id=SCREAMING_SNAKE_CASE_       ,		**SCREAMING_SNAKE_CASE_       ,		)






class     SCREAMING_SNAKE_CASE__      (							_a      ):

							"""simple docstring"""






							@property
							def       A__						(					self     )->				Union[str, Any]:





													'''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_(SCREAMING_SNAKE_CASE_       ,		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(SCREAMING_SNAKE_CASE_     ):
																															__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       A__						(					self     )->				Union[str, Any]:





													'''simple docstring'''






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



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = False       ,		SCREAMING_SNAKE_CASE_ = None       ,		)->				Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase       					=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )

													# Generate decoder inputs
													__UpperCamelCase       					=     seq_length if not self.use_past else 1
													__UpperCamelCase       					=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     {F"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
													__UpperCamelCase       					=     dict(**SCREAMING_SNAKE_CASE_       ,		**SCREAMING_SNAKE_CASE_     )

													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(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )]       ,		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(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )
																			__UpperCamelCase       					=     max(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     ) - min_num_layers
																			__UpperCamelCase       					=     '''encoder''' if num_encoder_layers > num_decoder_layers else '''decoder'''

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



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = False       ,		SCREAMING_SNAKE_CASE_ = None       ,		)->				str:





													'''simple docstring'''






													__UpperCamelCase       					=     self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )

													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(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		dtype=SCREAMING_SNAKE_CASE_     )]       ,		dim=1     )
																			__UpperCamelCase       					=     [
																			    (torch.zeros(SCREAMING_SNAKE_CASE_     ), torch.zeros(SCREAMING_SNAKE_CASE_     )) for _ in range(SCREAMING_SNAKE_CASE_     )
																			]
													return common_inputs



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = False       ,		SCREAMING_SNAKE_CASE_ = None       ,		)->				Optional[int]:





													'''simple docstring'''






													__UpperCamelCase       					=     compute_effective_axis_dimension(
													    SCREAMING_SNAKE_CASE_       ,		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(SCREAMING_SNAKE_CASE_     )
													__UpperCamelCase       					=     compute_effective_axis_dimension(
													    SCREAMING_SNAKE_CASE_       ,		fixed_dimension=OnnxConfig.default_fixed_sequence       ,		num_token_to_add=SCREAMING_SNAKE_CASE_     )

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



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = -1       ,		SCREAMING_SNAKE_CASE_ = False       ,		SCREAMING_SNAKE_CASE_ = None       ,		)->				Optional[int]:





													'''simple docstring'''






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

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

													return common_inputs



							def       A__						(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )->				List[str]:





													'''simple docstring'''






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






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							:				Union[str, Any]	, snake_case							:				str	, snake_case							:				Dict	, snake_case							:				Optional[Any]=None	) ->      Optional[Any]:



						'''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							:				Tuple	, snake_case							:				Union[str, Any]	) ->      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							:				Dict	, snake_case							:				int	) ->      str:



						'''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							:				Union[str, os.PathLike]	, snake_case							:				Dict[str, torch.Tensor]	) ->      Dict:



						'''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     SCREAMING_SNAKE_CASE__      (							SCREAMING_SNAKE_CASE_      ):

							"""simple docstring"""






							def __init__(					self       ,		SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )->				str:





													'''simple docstring'''






													__UpperCamelCase       					=     dataset
													__UpperCamelCase       					=     prefix



							def __getitem__(					self       ,		SCREAMING_SNAKE_CASE_     )->				Union[str, Any]:





													'''simple docstring'''






													return self.dataset[F"{self.prefix}{key}"]



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





													'''simple docstring'''






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



							def __len__(					self     )->				str:





													'''simple docstring'''






													return len(self.dataset     )






class     SCREAMING_SNAKE_CASE__      (							SCREAMING_SNAKE_CASE_      ):

							"""simple docstring"""






							def __init__(					self       ,		SCREAMING_SNAKE_CASE_ = None       ,		SCREAMING_SNAKE_CASE_ = None       ,		SCREAMING_SNAKE_CASE_ = None       ,		SCREAMING_SNAKE_CASE_=None       ,		)->				Tuple:





													'''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(SCREAMING_SNAKE_CASE_       ,		'''index.json'''     )     ) as f:
																									__UpperCamelCase       					=     json.load(SCREAMING_SNAKE_CASE_     )
													__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       ,		SCREAMING_SNAKE_CASE_     )->				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=SCREAMING_SNAKE_CASE_     ) as f:
																									__UpperCamelCase       					=     f.get_tensor(weight_info.get('''weight_name'''       ,		SCREAMING_SNAKE_CASE_     )     )

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

													__UpperCamelCase       					=     os.path.join(self.save_folder       ,		F"{key}.dat"     )
													return load_offloaded_weight(SCREAMING_SNAKE_CASE_       ,		SCREAMING_SNAKE_CASE_     )



							def __iter__(					self     )->				Dict:





													'''simple docstring'''






													return iter(self.all_keys     )



							def __len__(					self     )->				Dict:





													'''simple docstring'''






													return len(self.all_keys     )





def 		A_	(							snake_case							:				Dict[str, torch.Tensor]	, snake_case							:				List[str]	) ->      Tuple:



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




