






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

__UpperCamelCase:  str   =   {
    """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 			__lowerCAmelCase				(					_UpperCamelCase			):


	'''simple docstring'''





	_A       	=							"mbart"
	_A       	=							["past_key_values"]
	_A       	=							{"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}

	def __init__(	self:	Any,				lowerCamelCase_:	Any=50265,				lowerCamelCase_:	Optional[int]=1024,				lowerCamelCase_:	List[Any]=12,				lowerCamelCase_:	Union[str, Any]=4096,				lowerCamelCase_:	Any=16,				lowerCamelCase_:	Optional[Any]=12,				lowerCamelCase_:	Any=4096,				lowerCamelCase_:	Optional[int]=16,				lowerCamelCase_:	Tuple=0.0,				lowerCamelCase_:	Optional[Any]=0.0,				lowerCamelCase_:	str=True,				lowerCamelCase_:	int=True,				lowerCamelCase_:	Optional[Any]="gelu",				lowerCamelCase_:	str=1024,				lowerCamelCase_:	List[Any]=0.1,				lowerCamelCase_:	Union[str, Any]=0.0,				lowerCamelCase_:	Union[str, Any]=0.0,				lowerCamelCase_:	Union[str, Any]=0.0_2,				lowerCamelCase_:	List[Any]=0.0,				lowerCamelCase_:	Tuple=False,				lowerCamelCase_:	List[str]=1,				lowerCamelCase_:	str=0,				lowerCamelCase_:	Any=2,				lowerCamelCase_:	Optional[Any]=2,				**lowerCamelCase_:	List[str],				):
			lowercase__				:			List[str]            =		vocab_size
			lowercase__				:			Optional[Any]            =		max_position_embeddings
			lowercase__				:			List[Any]            =		d_model
			lowercase__				:			Dict            =		encoder_ffn_dim
			lowercase__				:			Any            =		encoder_layers
			lowercase__				:			List[str]            =		encoder_attention_heads
			lowercase__				:			Any            =		decoder_ffn_dim
			lowercase__				:			List[str]            =		decoder_layers
			lowercase__				:			Any            =		decoder_attention_heads
			lowercase__				:			Union[str, Any]            =		dropout
			lowercase__				:			Dict            =		attention_dropout
			lowercase__				:			Union[str, Any]            =		activation_dropout
			lowercase__				:			Union[str, Any]            =		activation_function
			lowercase__				:			str            =		init_std
			lowercase__				:			List[Any]            =		encoder_layerdrop
			lowercase__				:			Union[str, Any]            =		decoder_layerdrop
			lowercase__				:			Any            =		classifier_dropout
			lowercase__				:			Union[str, Any]            =		use_cache
			lowercase__				:			Dict            =		encoder_layers
			lowercase__				:			int            =		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_,				forced_eos_token_id=lowerCamelCase_,				**lowerCamelCase_,				)


class 			__lowerCAmelCase				(					_UpperCamelCase			):


	'''simple docstring'''





	@property
	def       snake_case__(	self:	str ):
			if self.task in ["default", "seq2seq-lm"]:
					lowercase__				:			int            =		OrderedDict(
					    [
					        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
					        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
					    ] )

					if self.use_past:
							lowercase__				:			str            =		{0: 'batch'}
							lowercase__				:			int            =		{0: 'batch', 1: 'past_decoder_sequence + sequence'}
					else:
							lowercase__				:			int            =		{0: 'batch', 1: 'decoder_sequence'}
							lowercase__				:			Any            =		{0: 'batch', 1: 'decoder_sequence'}

					if self.use_past:
							self.fill_with_past_key_values_(lowerCamelCase_,				direction='inputs' )
			elif self.task == "causal-lm":
					# TODO: figure this case out.
					lowercase__				:			Any            =		OrderedDict(
					    [
					        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
					        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
					    ] )
					if self.use_past:
							lowercase__				:			Optional[int]            =		self.num_layers
							for i in range(lowerCamelCase_ ):
									lowercase__				:			Any            =		{0: 'batch', 2: 'past_sequence + sequence'}
									lowercase__				:			Any            =		{0: 'batch', 2: 'past_sequence + sequence'}
			else:
					lowercase__				:			Union[str, Any]            =		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       snake_case__(	self:	Tuple ):
			if self.task in ["default", "seq2seq-lm"]:
					lowercase__				:			List[Any]            =		super().outputs
			else:
					lowercase__				:			Union[str, Any]            =		super(lowerCamelCase_,				self ).outputs
					if self.use_past:
							lowercase__				:			Any            =		self.num_layers
							for i in range(lowerCamelCase_ ):
									lowercase__				:			Optional[int]            =		{0: 'batch', 2: 'past_sequence + sequence'}
									lowercase__				:			int            =		{0: 'batch', 2: 'past_sequence + sequence'}
			return common_outputs

	def       snake_case__(	self:	Optional[Any],				lowerCamelCase_:	PreTrainedTokenizer,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	bool = False,				lowerCamelCase_:	Optional[TensorType] = None,				):
			lowercase__				:			int            =		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
			    lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_ )

			# Generate decoder inputs
			lowercase__				:			Optional[Any]            =		seq_length if not self.use_past else 1
			lowercase__				:			List[str]            =		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
			    lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_ )
			lowercase__				:			List[Any]            =		{F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
			lowercase__				:			int            =		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
					lowercase__				:			int            =		common_inputs['input_ids'].shape
					lowercase__				:			int            =		common_inputs['decoder_input_ids'].shape[1]
					lowercase__				:			str            =		self.num_attention_heads
					lowercase__				:			List[str]            =		(
					    batch,
					    num_encoder_attention_heads,
					    encoder_seq_length,
					    self._config.hidden_size // num_encoder_attention_heads,
					)
					lowercase__				:			Union[str, Any]            =		decoder_seq_length + 3
					lowercase__				:			Optional[Any]            =		(
					    batch,
					    num_decoder_attention_heads,
					    decoder_past_length,
					    self._config.hidden_size // num_decoder_attention_heads,
					)

					lowercase__				:			List[str]            =		torch.cat(
					    [common_inputs['decoder_attention_mask'], torch.ones(lowerCamelCase_,				lowerCamelCase_ )],				dim=1 )

					lowercase__				:			Union[str, Any]            =		[]
					# If the number of encoder and decoder layers are present in the model configuration, both are considered
					lowercase__				:			Optional[Any]            =		self.num_layers
					lowercase__				:			int            =		min(lowerCamelCase_,				lowerCamelCase_ )
					lowercase__				:			Any            =		max(lowerCamelCase_,				lowerCamelCase_ ) - min_num_layers
					lowercase__				:			Tuple            =		'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.
					lowercase__				:			Union[str, Any]            =		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

	def       snake_case__(	self:	Optional[int],				lowerCamelCase_:	PreTrainedTokenizer,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	bool = False,				lowerCamelCase_:	Optional[TensorType] = None,				):
			lowercase__				:			Dict            =		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
			    lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_,				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
					lowercase__				:			Dict            =		common_inputs['input_ids'].shape
					# Not using the same length for past_key_values
					lowercase__				:			str            =		seqlen + 2
					lowercase__				:			Tuple            =		self.num_layers
					lowercase__				:			str            =		self.num_attention_heads
					lowercase__				:			List[Any]            =		(
					    batch,
					    num_encoder_attention_heads,
					    past_key_values_length,
					    self._config.hidden_size // num_encoder_attention_heads,
					)

					lowercase__				:			Optional[Any]            =		common_inputs['attention_mask'].dtype
					lowercase__				:			Dict            =		torch.cat(
					    [common_inputs['attention_mask'], torch.ones(lowerCamelCase_,				lowerCamelCase_,				dtype=lowerCamelCase_ )],				dim=1 )
					lowercase__				:			Optional[int]            =		[
					    (torch.zeros(lowerCamelCase_ ), torch.zeros(lowerCamelCase_ )) for _ in range(lowerCamelCase_ )
					]
			return common_inputs

	def       snake_case__(	self:	Tuple,				lowerCamelCase_:	PreTrainedTokenizer,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	bool = False,				lowerCamelCase_:	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
			lowercase__				:			List[str]            =		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
			lowercase__				:			List[str]            =		tokenizer.num_special_tokens_to_add(lowerCamelCase_ )
			lowercase__				:			int            =		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
			lowercase__				:			str            =		[' '.join([tokenizer.unk_token] ) * seq_length] * batch_size
			lowercase__				:			List[str]            =		dict(tokenizer(lowerCamelCase_,				return_tensors=lowerCamelCase_ ) )
			return common_inputs

	def       snake_case__(	self:	Union[str, Any],				lowerCamelCase_:	PreTrainedTokenizer,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	int = -1,				lowerCamelCase_:	bool = False,				lowerCamelCase_:	Optional[TensorType] = None,				):
			if self.task in ["default", "seq2seq-lm"]:
					lowercase__				:			str            =		self._generate_dummy_inputs_for_default_and_seqaseq_lm(
					    lowerCamelCase_,				batch_size=lowerCamelCase_,				seq_length=lowerCamelCase_,				is_pair=lowerCamelCase_,				framework=lowerCamelCase_ )

			elif self.task == "causal-lm":
					lowercase__				:			Tuple            =		self._generate_dummy_inputs_for_causal_lm(
					    lowerCamelCase_,				batch_size=lowerCamelCase_,				seq_length=lowerCamelCase_,				is_pair=lowerCamelCase_,				framework=lowerCamelCase_ )
			else:
					lowercase__				:			Optional[int]            =		self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
					    lowerCamelCase_,				batch_size=lowerCamelCase_,				seq_length=lowerCamelCase_,				is_pair=lowerCamelCase_,				framework=lowerCamelCase_ )

			return common_inputs





	def       snake_case__(	self:	Union[str, Any],				lowerCamelCase_:	List[Any],				lowerCamelCase_:	List[str],				lowerCamelCase_:	List[str],				lowerCamelCase_:	List[str] ):
			if self.task in ["default", "seq2seq-lm"]:
					lowercase__				:			Any            =		super()._flatten_past_key_values_(lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_ )
			else:
					lowercase__				:			str            =		super(lowerCamelCase_,				self )._flatten_past_key_values_(
					    lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_,				lowerCamelCase_ )









import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union


class 			__lowerCAmelCase				(					enum.Enum			):


	'''simple docstring'''





	_A       	=							"timm"
	_A       	=							"transformers"


def 			SCREAMING_SNAKE_CASE__     (					_lowercase    :   Optional[Iterable[str]]      ,						_lowercase    :   Optional[Iterable[int]]      ,						_lowercase    :   Optional[Iterable[str]] )      ->       Optional[int]:




		'''simple docstring'''
		if stage_names is None:
				raise ValueError('Stage_names must be set for transformers backbones' )

		if out_features is not None:
				if not isinstance(_lowercase      ,						(list,) ):
						raise ValueError(f"""out_features must be a list {type(_lowercase )}""" )
				if any(feat not in stage_names for feat in out_features ):
						raise ValueError(f"""out_features must be a subset of stage_names: {stage_names} got {out_features}""" )

		if out_indices is not None:
				if not isinstance(_lowercase      ,						(list, tuple) ):
						raise ValueError(f"""out_indices must be a list or tuple, got {type(_lowercase )}""" )
				if any(idx >= len(_lowercase ) for idx in out_indices ):
						raise ValueError('out_indices must be valid indices for stage_names {stage_names}, got {out_indices}' )

		if out_features is not None and out_indices is not None:
				if len(_lowercase ) != len(_lowercase ):
						raise ValueError('out_features and out_indices should have the same length if both are set' )
				if out_features != [stage_names[idx] for idx in out_indices]:
						raise ValueError('out_features and out_indices should correspond to the same stages if both are set' )


def 			SCREAMING_SNAKE_CASE__     (					_lowercase    :   Optional[List[str]]      ,						_lowercase    :   Optional[Union[List[int], Tuple[int]]]      ,						_lowercase    :   List[str]      ,						)      ->       Dict:




		'''simple docstring'''
		if out_indices is None and out_features is None:
				lowercase__				:			str            =		[len(_lowercase ) - 1]
				lowercase__				:			Union[str, Any]            =		[stage_names[-1]]
		elif out_indices is None and out_features is not None:
				lowercase__				:			Optional[int]            =		[stage_names.index(_lowercase ) for layer in out_features]
		elif out_features is None and out_indices is not None:
				lowercase__				:			str            =		[stage_names[idx] for idx in out_indices]
		return out_features, out_indices


def 			SCREAMING_SNAKE_CASE__     (					_lowercase    :   Optional[List[str]]      ,						_lowercase    :   Optional[Union[List[int], Tuple[int]]]      ,						_lowercase    :   List[str]      ,						)      ->       Tuple[List[str], List[int]]:




		'''simple docstring'''
		verify_out_features_out_indices(out_features=_lowercase      ,						out_indices=_lowercase      ,						stage_names=_lowercase )
		lowercase__      ,							lowercase__				:			List[str]            =		_align_output_features_output_indices(
		    out_features=_lowercase      ,						out_indices=_lowercase      ,						stage_names=_lowercase )
		# Verify that the aligned out_features and out_indices are valid
		verify_out_features_out_indices(out_features=_lowercase      ,						out_indices=_lowercase      ,						stage_names=_lowercase )
		return output_features, output_indices


class 			__lowerCAmelCase				:


	'''simple docstring'''





	_A       	=							None

	def       snake_case__(	self:	Dict,				lowerCamelCase_:	Optional[Any] ):
			if getattr(self,				'_backbone',				lowerCamelCase_ ) is None:
					raise ValueError('self._backbone must be set before calling _init_timm_backbone' )

			# These will diagree with the defaults for the transformers models e.g. for resnet50
			# the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
			# the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
			lowercase__				:			Tuple            =		[stage['module'] for stage in self._backbone.feature_info.info]
			lowercase__				:			Union[str, Any]            =		[stage['num_chs'] for stage in self._backbone.feature_info.info]
			lowercase__				:			Dict            =		self._backbone.feature_info.out_indices
			lowercase__				:			Tuple            =		self._backbone.feature_info.module_name()

			# We verify the out indices and out features are valid
			verify_out_features_out_indices(
			    out_features=lowerCamelCase_,				out_indices=lowerCamelCase_,				stage_names=self.stage_names )
			lowercase__      ,							lowercase__				:			List[Any]            =		out_features, out_indices

	def       snake_case__(	self:	Any,				lowerCamelCase_:	List[str] ):
			lowercase__				:			Dict            =		getattr(lowerCamelCase_,				'stage_names' )
			lowercase__				:			List[str]            =		getattr(lowerCamelCase_,				'out_features',				lowerCamelCase_ )
			lowercase__				:			Optional[int]            =		getattr(lowerCamelCase_,				'out_indices',				lowerCamelCase_ )

			lowercase__				:			Optional[int]            =		stage_names
			lowercase__      ,							lowercase__				:			Union[str, Any]            =		get_aligned_output_features_output_indices(
			    out_features=lowerCamelCase_,				out_indices=lowerCamelCase_,				stage_names=lowerCamelCase_ )
			# Number of channels for each stage. This is set in the transformer backbone model init
			lowercase__				:			Dict            =		None

	def       snake_case__(	self:	List[Any],				lowerCamelCase_:	List[str] ):
			lowercase__				:			List[Any]            =		config

			lowercase__				:			Optional[int]            =		getattr(lowerCamelCase_,				'use_timm_backbone',				lowerCamelCase_ )
			lowercase__				:			List[Any]            =		BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

			if self.backbone_type == BackboneType.TIMM:
					self._init_timm_backbone(lowerCamelCase_ )
			elif self.backbone_type == BackboneType.TRANSFORMERS:
					self._init_transformers_backbone(lowerCamelCase_ )
			else:
					raise ValueError(F"""backbone_type {self.backbone_type} not supported.""" )

	@property
	def       snake_case__(	self:	List[str] ):
			return self._out_features

	@out_features.setter
	def       snake_case__(	self:	Union[str, Any],				lowerCamelCase_:	List[str] ):
			lowercase__      ,							lowercase__				:			Optional[int]            =		get_aligned_output_features_output_indices(
			    out_features=lowerCamelCase_,				out_indices=lowerCamelCase_,				stage_names=self.stage_names )

	@property
	def       snake_case__(	self:	List[Any] ):
			return self._out_indices

	@out_indices.setter
	def       snake_case__(	self:	Tuple,				lowerCamelCase_:	Union[Tuple[int], List[int]] ):
			lowercase__      ,							lowercase__				:			List[str]            =		get_aligned_output_features_output_indices(
			    out_features=lowerCamelCase_,				out_indices=lowerCamelCase_,				stage_names=self.stage_names )

	@property
	def       snake_case__(	self:	str ):
			# the current backbones will output the number of channels for each stage
			# even if that stage is not in the out_features list.
			return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names )}

	@property
	def       snake_case__(	self:	Optional[int] ):
			return [self.out_feature_channels[name] for name in self.out_features]

	def       snake_case__(	self:	List[Any],				*lowerCamelCase_:	List[Any],				**lowerCamelCase_:	Optional[Any] ):
			lowercase__				:			Any            =		dict(inspect.signature(self.forward ).parameters )
			lowercase__				:			List[Any]            =		{k: v for k, v in kwargs.items() if k in signature}
			return self(*lowerCamelCase_,				**lowerCamelCase_ )

	def       snake_case__(	self:	int,				lowerCamelCase_:	str,				lowerCamelCase_:	Optional[bool] = None,				lowerCamelCase_:	Optional[bool] = None,				lowerCamelCase_:	Optional[bool] = None,				):
			raise NotImplementedError('This method should be implemented by the derived class.' )





	def       snake_case__(	self:	Dict ):
			lowercase__				:			List[Any]            =		super().to_dict()
			lowercase__				:			Union[str, Any]            =		output.pop('_out_features' )
			lowercase__				:			Any            =		output.pop('_out_indices' )
			return output


class 			__lowerCAmelCase				:


	'''simple docstring'''





	@property
	def       snake_case__(	self:	Optional[int] ):
			return self._out_features

	@out_features.setter
	def       snake_case__(	self:	List[Any],				lowerCamelCase_:	List[str] ):
			lowercase__      ,							lowercase__				:			List[Any]            =		get_aligned_output_features_output_indices(
			    out_features=lowerCamelCase_,				out_indices=lowerCamelCase_,				stage_names=self.stage_names )

	@property
	def       snake_case__(	self:	Dict ):
			return self._out_indices

	@out_indices.setter
	def       snake_case__(	self:	Optional[int],				lowerCamelCase_:	Union[Tuple[int], List[int]] ):
			lowercase__      ,							lowercase__				:			Any            =		get_aligned_output_features_output_indices(
			    out_features=lowerCamelCase_,				out_indices=lowerCamelCase_,				stage_names=self.stage_names )





	def       snake_case__(	self:	Union[str, Any] ):
			lowercase__				:			List[str]            =		super().to_dict()
			lowercase__				:			str            =		output.pop('_out_features' )
			lowercase__				:			str            =		output.pop('_out_indices' )
			return output


