


'''simple docstring'''



import warnings
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    :    List[str]			     =  logging.get_logger(__name__)

_lowercase    :    Union[str, Any]			     =  {
    """facebook/bart-large""": """https://huggingface.co/facebook/bart-large/resolve/main/config.json""",
    # See all BART models at https://huggingface.co/models?filter=bart
}






class  UpperCamelCase__(     lowerCAmelCase ):
				__magic_name__       :   Any							= "bart"
				__magic_name__       :   Tuple							= ["past_key_values"]
				__magic_name__       :   Dict							= {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




				def __init__(	self    :				List[str]   ,					lowerCAmelCase    :				Tuple=50265   ,					lowerCAmelCase    :				str=1024   ,					lowerCAmelCase    :				int=12   ,					lowerCAmelCase    :				Optional[Any]=4096   ,					lowerCAmelCase    :				List[str]=16   ,					lowerCAmelCase    :				Dict=12   ,					lowerCAmelCase    :				Optional[Any]=4096   ,					lowerCAmelCase    :				Union[str, Any]=16   ,					lowerCAmelCase    :				Dict=0.0   ,					lowerCAmelCase    :				Optional[int]=0.0   ,					lowerCAmelCase    :				List[str]="gelu"   ,					lowerCAmelCase    :				Dict=1024   ,					lowerCAmelCase    :				List[str]=0.1   ,					lowerCAmelCase    :				List[Any]=0.0   ,					lowerCAmelCase    :				Optional[Any]=0.0   ,					lowerCAmelCase    :				Dict=0.02   ,					lowerCAmelCase    :				Optional[int]=0.0   ,					lowerCAmelCase    :				int=False   ,					lowerCAmelCase    :				Any=True   ,					lowerCAmelCase    :				Union[str, Any]=3   ,					lowerCAmelCase    :				int=1   ,					lowerCAmelCase    :				List[Any]=0   ,					lowerCAmelCase    :				Tuple=2   ,					lowerCAmelCase    :				List[Any]=True   ,					lowerCAmelCase    :				List[str]=2   ,					lowerCAmelCase    :				Any=2   ,					**lowerCAmelCase    :				Dict   ,					)->  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__(
										    num_labels=lowerCAmelCase   ,					pad_token_id=lowerCAmelCase   ,					bos_token_id=lowerCAmelCase   ,					eos_token_id=lowerCAmelCase   ,					is_encoder_decoder=lowerCAmelCase   ,					decoder_start_token_id=lowerCAmelCase   ,					forced_eos_token_id=lowerCAmelCase   ,					**lowerCAmelCase   ,					)

										# ensure backward compatibility for BART CNN models
										if self.forced_bos_token_id is None and kwargs.get('''force_bos_token_to_be_generated'''   ,					lowerCAmelCase   ):
																UpperCAmelCase         =			self.bos_token_id
																warnings.warn(
																    F"""Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. """
																    '''The config can simply be saved and uploaded again to be fixed.'''   )







class  UpperCamelCase__(     lowerCAmelCase ):




				@property
				def   a__(	self    :				Tuple   )->  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_(lowerCAmelCase   ,					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(lowerCAmelCase   ):
																												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    :				Any   )->  Mapping[str, Mapping[int, str]]:

										"""simple docstring"""




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




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

										"""simple docstring"""




										UpperCAmelCase         =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )

										# Generate decoder inputs
										UpperCAmelCase         =			seq_length if not self.use_past else 1
										UpperCAmelCase         =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   ,					lowerCAmelCase   )
										UpperCAmelCase         =			{F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
										UpperCAmelCase         =			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
																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(lowerCAmelCase   ,					lowerCAmelCase   )]   ,					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(lowerCAmelCase   ,					lowerCAmelCase   )
																UpperCAmelCase         =			max(lowerCAmelCase   ,					lowerCAmelCase   ) - min_num_layers
																UpperCAmelCase         =			'''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.
																UpperCAmelCase         =			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   a__(	self    :				int   ,					lowerCAmelCase    :				PreTrainedTokenizer   ,					lowerCAmelCase    :				int = -1   ,					lowerCAmelCase    :				int = -1   ,					lowerCAmelCase    :				bool = False   ,					lowerCAmelCase    :				Optional[TensorType] = None   ,					)->  Mapping[str, Any]:

										"""simple docstring"""




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




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

										"""simple docstring"""




										UpperCAmelCase         =			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
										UpperCAmelCase         =			tokenizer.num_special_tokens_to_add(lowerCAmelCase   )
										UpperCAmelCase         =			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
										UpperCAmelCase         =			[''' '''.join([tokenizer.unk_token]   ) * seq_length] * batch_size
										UpperCAmelCase         =			dict(tokenizer(lowerCAmelCase   ,					return_tensors=lowerCAmelCase   )   )
										return common_inputs




				def   a__(	self    :				Optional[Any]   ,					lowerCAmelCase    :				PreTrainedTokenizer   ,					lowerCAmelCase    :				int = -1   ,					lowerCAmelCase    :				int = -1   ,					lowerCAmelCase    :				bool = False   ,					lowerCAmelCase    :				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(
																    lowerCAmelCase   ,					batch_size=lowerCAmelCase   ,					seq_length=lowerCAmelCase   ,					is_pair=lowerCAmelCase   ,					framework=lowerCAmelCase   )

										elif self.task == "causal-lm":
																UpperCAmelCase         =			self._generate_dummy_inputs_for_causal_lm(
																    lowerCAmelCase   ,					batch_size=lowerCAmelCase   ,					seq_length=lowerCAmelCase   ,					is_pair=lowerCAmelCase   ,					framework=lowerCAmelCase   )
										else:
																UpperCAmelCase         =			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   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				List[Any]   ,					lowerCAmelCase    :				List[str]   ,					lowerCAmelCase    :				Optional[Any]   )->  List[Any]:

										"""simple docstring"""




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









'''simple docstring'''



import time
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
	import torch

	from transformers.generation import (
	    MaxLengthCriteria,
	    MaxNewTokensCriteria,
	    MaxTimeCriteria,
	    StoppingCriteriaList,
	    validate_stopping_criteria,
	)






@require_torch
class  UpperCamelCase__(     unittest.TestCase ):




				def   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				Optional[Any]   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase         =			3
										UpperCAmelCase         =			250

										UpperCAmelCase         =			ids_tensor((batch_size, length)   ,					lowerCAmelCase   )
										UpperCAmelCase         =			torch.ones((batch_size, length)   ,					device=lowerCAmelCase   ,					dtype=torch.float   ) / length
										return input_ids, scores




				def   a__(	self    :				Dict   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )

										UpperCAmelCase         =			StoppingCriteriaList(
										    [
										        MaxLengthCriteria(max_length=10   ),
										        MaxTimeCriteria(max_time=0.1   ),
										    ]   )

										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )




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

										"""simple docstring"""




										UpperCAmelCase         =			MaxLengthCriteria(max_length=10   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )




				def   a__(	self    :				Optional[Any]   )->  List[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			MaxNewTokensCriteria(start_length=5   ,					max_new_tokens=5   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase         =			StoppingCriteriaList([criteria]   )
										self.assertEqual(criteria_list.max_length   ,					10   )




				def   a__(	self    :				Tuple   )->  Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )

										UpperCAmelCase         =			MaxTimeCriteria(max_time=0.1   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase         =			MaxTimeCriteria(max_time=0.1   ,					initial_timestamp=time.time() - 0.2   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )





				def   a__(	self    :				int   )->  Any:

										"""simple docstring"""




										validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )   ,					10   )

										with self.assertWarns(lowerCAmelCase   ):
																validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )   ,					11   )

										UpperCAmelCase         =			validate_stopping_criteria(StoppingCriteriaList()   ,					11   )

										self.assertEqual(len(lowerCAmelCase   )   ,					1   )





