




import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


lowercase_		=							[nn.LayerNorm]

lowercase_		=							logging.get_logger(__name__)

lowercase_		=							version.parse(version.parse(torch.__version__).base_version)

lowercase_		=							parsed_torch_version_base >= version.parse('''2.0''')
lowercase_		=							parsed_torch_version_base >= version.parse('''1.12''')
lowercase_		=							parsed_torch_version_base >= version.parse('''1.11''')
lowercase_		=							parsed_torch_version_base < version.parse('''1.11''')







def 			lowerCAmelCase  (     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase,     self	)  ->Any:






					"""simple docstring"""






					from torch import _softmax_backward_data

					if is_torch_less_than_1_11:
										return _softmax_backward_data(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     parent.dim,     self	)
					else:
										return _softmax_backward_data(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     parent.dim,     self.dtype	)







def 			lowerCAmelCase  (     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase = 0	)  ->nn.Linear:






					"""simple docstring"""






					__magic_name__ :							List[str]		       =							index.to(layer.weight.device	)
					__magic_name__ :							List[str]		       =							layer.weight.index_select(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_	).clone().detach()
					if layer.bias is not None:
										if dim == 1:
															__magic_name__ :							Any		       =							layer.bias.clone().detach()
										else:
															__magic_name__ :							Dict		       =							layer.bias[index].clone().detach()
					__magic_name__ :							List[str]		       =							list(layer.weight.size()	)
					__magic_name__ :							Tuple		       =							len(SCREAMING_SNAKE_CASE_	)
					__magic_name__ :							Union[str, Any]		       =							nn.Linear(new_size[1],     new_size[0],     bias=layer.bias is not None	).to(layer.weight.device	)
					__magic_name__ :							List[str]		       =							False
					new_layer.weight.copy_(W.contiguous()	)
					__magic_name__ :							str		       =							True
					if layer.bias is not None:
										__magic_name__ :							Tuple		       =							False
										new_layer.bias.copy_(b.contiguous()	)
										__magic_name__ :							int		       =							True
					return new_layer





class 		A__    (       nn.Module			):






	def __init__(     self    ,					lowerCamelCase    ,					lowerCamelCase						)				->       Optional[int]:




						"""simple docstring"""






						super().__init__()
						__magic_name__ :							Dict		       =							nf
						__magic_name__ :							int		       =							nn.Parameter(torch.empty(lowerCamelCase    ,					lowerCamelCase						)						)
						__magic_name__ :							Tuple		       =							nn.Parameter(torch.zeros(lowerCamelCase						)						)
						nn.init.normal_(self.weight    ,					std=0.0_2						)






	def 			lowercase				(     self    ,					lowerCamelCase						)				->       Tuple:




						"""simple docstring"""






						__magic_name__ :							Optional[int]		       =							x.size()[:-1] + (self.nf,)
						__magic_name__ :							Dict		       =							torch.addmm(self.bias    ,					x.view(-1    ,					x.size(-1						)						)    ,					self.weight						)
						__magic_name__ :							Dict		       =							x.view(lowerCamelCase						)
						return x







def 			lowerCAmelCase  (     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase = 1	)  ->ConvaD:






					"""simple docstring"""






					__magic_name__ :							Union[str, Any]		       =							index.to(layer.weight.device	)
					__magic_name__ :							List[Any]		       =							layer.weight.index_select(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_	).clone().detach()
					if dim == 0:
										__magic_name__ :							List[str]		       =							layer.bias.clone().detach()
					else:
										__magic_name__ :							int		       =							layer.bias[index].clone().detach()
					__magic_name__ :							int		       =							list(layer.weight.size()	)
					__magic_name__ :							Optional[Any]		       =							len(SCREAMING_SNAKE_CASE_	)
					__magic_name__ :							Any		       =							ConvaD(new_size[1],     new_size[0]	).to(layer.weight.device	)
					__magic_name__ :							int		       =							False
					new_layer.weight.copy_(W.contiguous()	)
					__magic_name__ :							Union[str, Any]		       =							True
					__magic_name__ :							Union[str, Any]		       =							False
					new_layer.bias.copy_(b.contiguous()	)
					__magic_name__ :							Dict		       =							True
					return new_layer







def 			lowerCAmelCase  (     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase = None	)  ->Union[nn.Linear, ConvaD]:






					"""simple docstring"""






					if isinstance(SCREAMING_SNAKE_CASE_,     nn.Linear	):
										return prune_linear_layer(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     dim=0 if dim is None else dim	)
					elif isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_	):
										return prune_convad_layer(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     dim=1 if dim is None else dim	)
					else:
										raise ValueError(F'''Can\'t prune layer of class {layer.__class__}'''	)







def 			lowerCAmelCase  (     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase,     *UpperCAmelCase	)  ->torch.Tensor:






					"""simple docstring"""






					assert len(SCREAMING_SNAKE_CASE_	) > 0, F'''{input_tensors} has to be a tuple/list of tensors'''

					# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
					__magic_name__ :							Union[str, Any]		       =							len(inspect.signature(SCREAMING_SNAKE_CASE_	).parameters	)
					if num_args_in_forward_chunk_fn != len(SCREAMING_SNAKE_CASE_	):
										raise ValueError(
										    F'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(SCREAMING_SNAKE_CASE_	)} input '''
										    '''tensors are given'''	)

					if chunk_size > 0:
										__magic_name__ :							int		       =							input_tensors[0].shape[chunk_dim]
										for input_tensor in input_tensors:
															if input_tensor.shape[chunk_dim] != tensor_shape:
																				raise ValueError(
																				    F'''All input tenors have to be of the same shape: {tensor_shape}, '''
																				    F'''found shape {input_tensor.shape[chunk_dim]}'''	)

										if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
															raise ValueError(
															    F'''The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk '''
															    F'''size {chunk_size}'''	)

										__magic_name__ :							Any		       =							input_tensors[0].shape[chunk_dim] // chunk_size

										# chunk input tensor into tuples
										__magic_name__ :							int		       =							tuple(input_tensor.chunk(SCREAMING_SNAKE_CASE_,     dim=SCREAMING_SNAKE_CASE_	) for input_tensor in input_tensors	)
										# apply forward fn to every tuple
										__magic_name__ :							Any		       =							tuple(forward_fn(*SCREAMING_SNAKE_CASE_	) for input_tensors_chunk in zip(*SCREAMING_SNAKE_CASE_	)	)
										# concatenate output at same dimension
										return torch.cat(SCREAMING_SNAKE_CASE_,     dim=SCREAMING_SNAKE_CASE_	)

					return forward_fn(*SCREAMING_SNAKE_CASE_	)







def 			lowerCAmelCase  (     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase,     UpperCAmelCase	)  ->Tuple[Set[int], torch.LongTensor]:






					"""simple docstring"""






					__magic_name__ :							Optional[Any]		       =							torch.ones(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_	)
					__magic_name__ :							List[Any]		       =							set(SCREAMING_SNAKE_CASE_	) - already_pruned_heads  # Convert to set and remove already pruned heads
					for head in heads:
										# Compute how many pruned heads are before the head and move the index accordingly
										__magic_name__ :							Tuple		       =							head - sum(1 if h < head else 0 for h in already_pruned_heads	)
										__magic_name__ :							Tuple		       =							0
					__magic_name__ :							Dict		       =							mask.view(-1	).contiguous().eq(1	)
					__magic_name__ :							str		       =							torch.arange(len(SCREAMING_SNAKE_CASE_	)	)[mask].long()
					return heads, index







def 			lowerCAmelCase  (     *UpperCAmelCase,     UpperCAmelCase = None	)  ->Tuple[torch.Tensor, ...]:






					"""simple docstring"""






					return torch.meshgrid(*SCREAMING_SNAKE_CASE_,     indexing=SCREAMING_SNAKE_CASE_	)







def 			lowerCAmelCase  (     UpperCAmelCase	)  ->Tuple[torch.device, int, int]:






					"""simple docstring"""






					return tensor.device, storage_ptr(SCREAMING_SNAKE_CASE_	), storage_size(SCREAMING_SNAKE_CASE_	)









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

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


lowercase_		=							logging.get_logger(__name__)

lowercase_		=							{
    '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/config.json''',
    # See all M2M100 models at https://huggingface.co/models?filter=m2m_100
}





class 		A__    (       __SCREAMING_SNAKE_CASE			):
	lowerCamelCase__		:			Optional[int]       ="m2m_100"
	lowerCamelCase__		:			List[str]       =["past_key_values"]
	lowerCamelCase__		:			Union[str, Any]       ={"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}






	def __init__(     self    ,					lowerCamelCase=128112    ,					lowerCamelCase=1024    ,					lowerCamelCase=12    ,					lowerCamelCase=4096    ,					lowerCamelCase=16    ,					lowerCamelCase=12    ,					lowerCamelCase=4096    ,					lowerCamelCase=16    ,					lowerCamelCase=0.0_5    ,					lowerCamelCase=0.0_5    ,					lowerCamelCase=True    ,					lowerCamelCase=True    ,					lowerCamelCase="relu"    ,					lowerCamelCase=1024    ,					lowerCamelCase=0.1    ,					lowerCamelCase=0.1    ,					lowerCamelCase=0.0    ,					lowerCamelCase=0.0_2    ,					lowerCamelCase=2    ,					lowerCamelCase=True    ,					lowerCamelCase=1    ,					lowerCamelCase=0    ,					lowerCamelCase=2    ,					**lowerCamelCase    ,					)				->       List[str]:




						"""simple docstring"""






						__magic_name__ :							List[Any]		       =							vocab_size
						__magic_name__ :							str		       =							max_position_embeddings
						__magic_name__ :							Union[str, Any]		       =							d_model
						__magic_name__ :							Dict		       =							encoder_ffn_dim
						__magic_name__ :							List[Any]		       =							encoder_layers
						__magic_name__ :							Optional[int]		       =							encoder_attention_heads
						__magic_name__ :							Dict		       =							decoder_ffn_dim
						__magic_name__ :							List[str]		       =							decoder_layers
						__magic_name__ :							Optional[Any]		       =							decoder_attention_heads
						__magic_name__ :							List[Any]		       =							dropout
						__magic_name__ :							List[str]		       =							attention_dropout
						__magic_name__ :							Optional[int]		       =							activation_dropout
						__magic_name__ :							Dict		       =							activation_function
						__magic_name__ :							Tuple		       =							init_std
						__magic_name__ :							Optional[int]		       =							encoder_layerdrop
						__magic_name__ :							int		       =							decoder_layerdrop
						__magic_name__ :							Union[str, Any]		       =							use_cache
						__magic_name__ :							Optional[int]		       =							encoder_layers
						__magic_name__ :							Dict		       =							scale_embedding  # scale factor will be sqrt(d_model) if True

						super().__init__(
						    pad_token_id=lowerCamelCase    ,					bos_token_id=lowerCamelCase    ,					eos_token_id=lowerCamelCase    ,					is_encoder_decoder=lowerCamelCase    ,					decoder_start_token_id=lowerCamelCase    ,					**lowerCamelCase    ,					)





class 		A__    (       __SCREAMING_SNAKE_CASE			):






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




						"""simple docstring"""






						__magic_name__ :							Dict		       =							OrderedDict(
						    [
						        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
						        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
						    ]						)

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

						if self.use_past:
											self.fill_with_past_key_values_(lowerCamelCase    ,					direction='''inputs'''						)
						return common_inputs






	def 			lowercase				(     self    ,					lowerCamelCase    ,					lowerCamelCase = -1    ,					lowerCamelCase = -1    ,					lowerCamelCase = False    ,					lowerCamelCase = None    ,					)				->       Mapping[str, Any]:




						"""simple docstring"""






						__magic_name__ :							int		       =							compute_effective_axis_dimension(
						    lowerCamelCase    ,					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
						__magic_name__ :							str		       =							tokenizer.num_special_tokens_to_add(lowerCamelCase						)
						__magic_name__ :							Optional[Any]		       =							compute_effective_axis_dimension(
						    lowerCamelCase    ,					fixed_dimension=OnnxConfig.default_fixed_sequence    ,					num_token_to_add=lowerCamelCase						)

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






	def 			lowercase				(     self    ,					lowerCamelCase    ,					lowerCamelCase = -1    ,					lowerCamelCase = -1    ,					lowerCamelCase = False    ,					lowerCamelCase = None    ,					)				->       Mapping[str, Any]:




						"""simple docstring"""






						__magic_name__ :							Any		       =							self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    lowerCamelCase    ,					lowerCamelCase    ,					lowerCamelCase    ,					lowerCamelCase    ,					lowerCamelCase						)

						# Generate decoder inputs
						__magic_name__ :							int		       =							seq_length if not self.use_past else 1
						__magic_name__ :							Optional[Any]		       =							self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
						    lowerCamelCase    ,					lowerCamelCase    ,					lowerCamelCase    ,					lowerCamelCase    ,					lowerCamelCase						)
						__magic_name__ :							List[str]		       =							{F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
						__magic_name__ :							Optional[Any]		       =							dict(**lowerCamelCase    ,					**lowerCamelCase						)

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

											__magic_name__ :							int		       =							torch.cat(
											    [common_inputs['''decoder_attention_mask'''], torch.ones(lowerCamelCase    ,					lowerCamelCase						)]    ,					dim=1						)

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

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





	lowerCamelCase__		:			List[str]       =_generate_dummy_inputs_for_default_and_seqaseq_lm



