



'''simple docstring'''
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									=      logging.get_logger(__name__)

_UpperCamelCase									=      {
    '''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	(     __SCREAMING_SNAKE_CASE     ):
		_SCREAMING_SNAKE_CASE       :     Any       =      "mbart"
		_SCREAMING_SNAKE_CASE       :     Optional[int]       =      ["past_key_values"]
		_SCREAMING_SNAKE_CASE       :     Optional[Any]       =      {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}





		def __init__(  self			,							__UpperCAmelCase=50_265			,							__UpperCAmelCase=1_024			,							__UpperCAmelCase=12			,							__UpperCAmelCase=4_096			,							__UpperCAmelCase=16			,							__UpperCAmelCase=12			,							__UpperCAmelCase=4_096			,							__UpperCAmelCase=16			,							__UpperCAmelCase=0.0			,							__UpperCAmelCase=0.0			,							__UpperCAmelCase=True			,							__UpperCAmelCase=True			,							__UpperCAmelCase="gelu"			,							__UpperCAmelCase=1_024			,							__UpperCAmelCase=0.1			,							__UpperCAmelCase=0.0			,							__UpperCAmelCase=0.0			,							__UpperCAmelCase=0.02			,							__UpperCAmelCase=0.0			,							__UpperCAmelCase=False			,							__UpperCAmelCase=1			,							__UpperCAmelCase=0			,							__UpperCAmelCase=2			,							__UpperCAmelCase=2			,							**__UpperCAmelCase			,							)							->  List[str]:

					'''simple docstring'''
					__UpperCAmelCase :   Tuple    					=  vocab_size
					__UpperCAmelCase :   Optional[Any]    					=  max_position_embeddings
					__UpperCAmelCase :   Optional[Any]    					=  d_model
					__UpperCAmelCase :   Optional[int]    					=  encoder_ffn_dim
					__UpperCAmelCase :   Tuple    					=  encoder_layers
					__UpperCAmelCase :   Union[str, Any]    					=  encoder_attention_heads
					__UpperCAmelCase :   str    					=  decoder_ffn_dim
					__UpperCAmelCase :   int    					=  decoder_layers
					__UpperCAmelCase :   str    					=  decoder_attention_heads
					__UpperCAmelCase :   Any    					=  dropout
					__UpperCAmelCase :   Union[str, Any]    					=  attention_dropout
					__UpperCAmelCase :   Dict    					=  activation_dropout
					__UpperCAmelCase :   List[Any]    					=  activation_function
					__UpperCAmelCase :   str    					=  init_std
					__UpperCAmelCase :   Dict    					=  encoder_layerdrop
					__UpperCAmelCase :   Dict    					=  decoder_layerdrop
					__UpperCAmelCase :   int    					=  classifier_dropout
					__UpperCAmelCase :   Dict    					=  use_cache
					__UpperCAmelCase :   Any    					=  encoder_layers
					__UpperCAmelCase :   Optional[int]    					=  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	(     __SCREAMING_SNAKE_CASE     ):





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

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

								if self.use_past:
											__UpperCAmelCase :   Any    					=  {0: """batch"""}
											__UpperCAmelCase :   Union[str, Any]    					=  {0: """batch""", 1: """past_decoder_sequence + sequence"""}
								else:
											__UpperCAmelCase :   Optional[Any]    					=  {0: """batch""", 1: """decoder_sequence"""}
											__UpperCAmelCase :   str    					=  {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 :   int    					=  OrderedDict(
								    [
								        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
								        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
								    ]					)
								if self.use_past:
											__UpperCAmelCase					,	__UpperCAmelCase :   List[str]    					=  self.num_layers
											for i in range(__UpperCAmelCase					):
														__UpperCAmelCase :   Dict    					=  {0: """batch""", 2: """past_sequence + sequence"""}
														__UpperCAmelCase :   int    					=  {0: """batch""", 2: """past_sequence + sequence"""}
					else:
								__UpperCAmelCase :   Optional[int]    					=  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					)							->  Mapping[str, Mapping[int, str]]:

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





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

					'''simple docstring'''
					__UpperCAmelCase :   Union[str, Any]    					=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
					    __UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase					)

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

								__UpperCAmelCase :   List[str]    					=  torch.cat(
								    [common_inputs["""decoder_attention_mask"""], torch.ones(__UpperCAmelCase			,							__UpperCAmelCase					)]			,							dim=1					)

								__UpperCAmelCase :   Any    					=  []
								# If the number of encoder and decoder layers are present in the model configuration, both are considered
								__UpperCAmelCase					,	__UpperCAmelCase :   Dict    					=  self.num_layers
								__UpperCAmelCase :   Optional[int]    					=  min(__UpperCAmelCase			,							__UpperCAmelCase					)
								__UpperCAmelCase :   Any    					=  max(__UpperCAmelCase			,							__UpperCAmelCase					) - min_num_layers
								__UpperCAmelCase :   Optional[Any]    					=  """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 :   Dict    					=  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 				__A							(  self			,							__UpperCAmelCase			,							__UpperCAmelCase = -1			,							__UpperCAmelCase = -1			,							__UpperCAmelCase = False			,							__UpperCAmelCase = None			,							)							->  Mapping[str, Any]:

					'''simple docstring'''
					__UpperCAmelCase :   Tuple    					=  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 :   Optional[int]    					=  common_inputs["""input_ids"""].shape
								# Not using the same length for past_key_values
								__UpperCAmelCase :   List[Any]    					=  seqlen + 2
								__UpperCAmelCase					,	__UpperCAmelCase :   Optional[Any]    					=  self.num_layers
								__UpperCAmelCase					,	__UpperCAmelCase :   Optional[int]    					=  self.num_attention_heads
								__UpperCAmelCase :   Dict    					=  (
								    batch,
								    num_encoder_attention_heads,
								    past_key_values_length,
								    self._config.hidden_size // num_encoder_attention_heads,
								)

								__UpperCAmelCase :   List[Any]    					=  common_inputs["""attention_mask"""].dtype
								__UpperCAmelCase :   List[Any]    					=  torch.cat(
								    [common_inputs["""attention_mask"""], torch.ones(__UpperCAmelCase			,							__UpperCAmelCase			,							dtype=__UpperCAmelCase					)]			,							dim=1					)
								__UpperCAmelCase :   List[str]    					=  [
								    (torch.zeros(__UpperCAmelCase					), torch.zeros(__UpperCAmelCase					)) for _ in range(__UpperCAmelCase					)
								]
					return common_inputs





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

					'''simple docstring'''
					# 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
					__UpperCAmelCase :   str    					=  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 :   int    					=  tokenizer.num_special_tokens_to_add(__UpperCAmelCase					)
					__UpperCAmelCase :   Any    					=  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 :   int    					=  [""" """.join([tokenizer.unk_token]					) * seq_length] * batch_size
					__UpperCAmelCase :   Optional[Any]    					=  dict(tokenizer(__UpperCAmelCase			,							return_tensors=__UpperCAmelCase					)					)
					return common_inputs





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

					'''simple docstring'''
					if self.task in ["default", "seq2seq-lm"]:
								__UpperCAmelCase :   str    					=  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 :   str    					=  self._generate_dummy_inputs_for_causal_lm(
								    __UpperCAmelCase			,							batch_size=__UpperCAmelCase			,							seq_length=__UpperCAmelCase			,							is_pair=__UpperCAmelCase			,							framework=__UpperCAmelCase					)
					else:
								__UpperCAmelCase :   int    					=  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 				__A							(  self			,							__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase					)							->  Optional[int]:

					'''simple docstring'''
					if self.task in ["default", "seq2seq-lm"]:
								__UpperCAmelCase :   str    					=  super()._flatten_past_key_values_(__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase					)
					else:
								__UpperCAmelCase :   int    					=  super(__UpperCAmelCase			,							self					)._flatten_past_key_values_(
								    __UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase					)









'''simple docstring'''
from transformers import (
    DPRContextEncoderTokenizer,
    DPRContextEncoderTokenizerFast,
    DPRQuestionEncoderTokenizer,
    DPRQuestionEncoderTokenizerFast,
    DPRReaderOutput,
    DPRReaderTokenizer,
    DPRReaderTokenizerFast,
)
from transformers.testing_utils import require_tokenizers, slow
from transformers.tokenization_utils_base import BatchEncoding

from ..bert.test_tokenization_bert import BertTokenizationTest



@require_tokenizers
class _A	(     __SCREAMING_SNAKE_CASE     ):
		_SCREAMING_SNAKE_CASE       :     List[str]       =      DPRContextEncoderTokenizer
		_SCREAMING_SNAKE_CASE       :     Optional[int]       =      DPRContextEncoderTokenizerFast
		_SCREAMING_SNAKE_CASE       :     Optional[Any]       =      True



@require_tokenizers
class _A	(     __SCREAMING_SNAKE_CASE     ):
		_SCREAMING_SNAKE_CASE       :     str       =      DPRQuestionEncoderTokenizer
		_SCREAMING_SNAKE_CASE       :     List[str]       =      DPRQuestionEncoderTokenizerFast
		_SCREAMING_SNAKE_CASE       :     Union[str, Any]       =      True



@require_tokenizers
class _A	(     __SCREAMING_SNAKE_CASE     ):
		_SCREAMING_SNAKE_CASE       :     List[str]       =      DPRReaderTokenizer
		_SCREAMING_SNAKE_CASE       :     Optional[Any]       =      DPRReaderTokenizerFast
		_SCREAMING_SNAKE_CASE       :     Dict       =      True





		@slow
		def 				__A							(  self					)							->  List[Any]:

					'''simple docstring'''
					__UpperCAmelCase :   Dict    					=  self.tokenizer_class.from_pretrained("""bert-base-uncased"""					)

					__UpperCAmelCase :   int    					=  tokenizer.encode("""question sequence"""			,							add_special_tokens=__UpperCAmelCase					)
					__UpperCAmelCase :   Optional[Any]    					=  tokenizer.encode("""title sequence"""			,							add_special_tokens=__UpperCAmelCase					)
					__UpperCAmelCase :   Any    					=  tokenizer.encode("""text sequence """ * 4			,							add_special_tokens=__UpperCAmelCase					)
					__UpperCAmelCase :   Tuple    					=  [[101] + text_a + [102] + text_a + [102] + text_a]
					__UpperCAmelCase :   Optional[int]    					=  BatchEncoding({"""input_ids""": input_ids}					)

					__UpperCAmelCase :   str    					=  [[0] * len(input_ids[0]					)]
					__UpperCAmelCase :   str    					=  [[0] * len(input_ids[0]					)]
					__UpperCAmelCase :   Optional[int]    					=  [0]
					__UpperCAmelCase :   List[str]    					=  DPRReaderOutput(__UpperCAmelCase			,							__UpperCAmelCase			,							__UpperCAmelCase					)

					__UpperCAmelCase					,	__UpperCAmelCase :   Any    					=  8, 9
					__UpperCAmelCase :   Optional[int]    					=  10
					__UpperCAmelCase :   int    					=  10
					__UpperCAmelCase :   str    					=  tokenizer.decode_best_spans(__UpperCAmelCase			,							__UpperCAmelCase					)
					self.assertEqual(predicted_spans[0].start_index			,							__UpperCAmelCase					)
					self.assertEqual(predicted_spans[0].end_index			,							__UpperCAmelCase					)
					self.assertEqual(predicted_spans[0].doc_id			,							0					)





		@slow
		def 				__A							(  self					)							->  int:

					'''simple docstring'''
					__UpperCAmelCase :   Any    					=  self.tokenizer_class.from_pretrained("""bert-base-uncased"""					)

					__UpperCAmelCase :   List[Any]    					=  tokenizer.encode("""question sequence"""			,							add_special_tokens=__UpperCAmelCase					)
					__UpperCAmelCase :   List[str]    					=  tokenizer.encode("""title sequence"""			,							add_special_tokens=__UpperCAmelCase					)
					__UpperCAmelCase :   Union[str, Any]    					=  tokenizer.encode("""text sequence"""			,							add_special_tokens=__UpperCAmelCase					)
					__UpperCAmelCase :   Optional[int]    					=  [101] + text_a + [102] + text_a + [102] + text_a
					__UpperCAmelCase :   int    					=  tokenizer(questions=["""question sequence"""]			,							titles=["""title sequence"""]			,							texts=["""text sequence"""]					)
					self.assertIn("""input_ids"""			,							__UpperCAmelCase					)
					self.assertIn("""attention_mask"""			,							__UpperCAmelCase					)
					self.assertListEqual(encoded_input["""input_ids"""][0]			,							__UpperCAmelCase					)




