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

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


_A  =	logging.get_logger(__name__)

_A  =	{
    "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json",
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}





class       _lowerCAmelCase   (      __a			):
				_lowercase     						='''blenderbot-small'''
				_lowercase     						=['''past_key_values''']
				_lowercase     						={'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}




				def __init__(   self							,							_UpperCamelCase=50_265							,							_UpperCamelCase=512							,							_UpperCamelCase=8							,							_UpperCamelCase=2_048							,							_UpperCamelCase=16							,							_UpperCamelCase=8							,							_UpperCamelCase=2_048							,							_UpperCamelCase=16							,							_UpperCamelCase=0.0							,							_UpperCamelCase=0.0							,							_UpperCamelCase=True							,							_UpperCamelCase=True							,							_UpperCamelCase="gelu"							,							_UpperCamelCase=512							,							_UpperCamelCase=0.1							,							_UpperCamelCase=0.0							,							_UpperCamelCase=0.0							,							_UpperCamelCase=0.02							,							_UpperCamelCase=1							,							_UpperCamelCase=False							,							_UpperCamelCase=0							,							_UpperCamelCase=1							,							_UpperCamelCase=2							,							_UpperCamelCase=2							,							**_UpperCamelCase							,							)							-> Optional[int]:
							lowerCAmelCase_         =      vocab_size
							lowerCAmelCase_         =      max_position_embeddings
							lowerCAmelCase_         =      d_model
							lowerCAmelCase_         =      encoder_ffn_dim
							lowerCAmelCase_         =      encoder_layers
							lowerCAmelCase_         =      encoder_attention_heads
							lowerCAmelCase_         =      decoder_ffn_dim
							lowerCAmelCase_         =      decoder_layers
							lowerCAmelCase_         =      decoder_attention_heads
							lowerCAmelCase_         =      dropout
							lowerCAmelCase_         =      attention_dropout
							lowerCAmelCase_         =      activation_dropout
							lowerCAmelCase_         =      activation_function
							lowerCAmelCase_         =      init_std
							lowerCAmelCase_         =      encoder_layerdrop
							lowerCAmelCase_         =      decoder_layerdrop
							lowerCAmelCase_         =      use_cache
							lowerCAmelCase_         =      encoder_layers
							lowerCAmelCase_         =      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							,							decoder_start_token_id=_UpperCamelCase							,							forced_eos_token_id=_UpperCamelCase							,							**_UpperCamelCase							,							)





class       _lowerCAmelCase   (      __a			):




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

										if self.use_past:
													lowerCAmelCase_         =      {0: "batch"}
													lowerCAmelCase_         =      {0: "batch", 1: "past_decoder_sequence + sequence"}
										else:
													lowerCAmelCase_         =      {0: "batch", 1: "decoder_sequence"}
													lowerCAmelCase_         =      {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.
										lowerCAmelCase_         =      OrderedDict(
										    [
										        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
										        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
										    ]	)
										if self.use_past:
													lowerCAmelCase_      ,    lowerCAmelCase_         =      self.num_layers
													for i in range(_UpperCamelCase	):
																lowerCAmelCase_         =      {0: "batch", 2: "past_sequence + sequence"}
																lowerCAmelCase_         =      {0: "batch", 2: "past_sequence + sequence"}
							else:
										lowerCAmelCase_         =      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]]:
							if self.task in ["default", "seq2seq-lm"]:
										lowerCAmelCase_         =      super().outputs
							else:
										lowerCAmelCase_         =      super(_UpperCamelCase							,							self	).outputs
										if self.use_past:
													lowerCAmelCase_      ,    lowerCAmelCase_         =      self.num_layers
													for i in range(_UpperCamelCase	):
																lowerCAmelCase_         =      {0: "batch", 2: "past_sequence + sequence"}
																lowerCAmelCase_         =      {0: "batch", 2: "past_sequence + sequence"}
							return common_outputs




				def 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase = -1							,							_UpperCamelCase = -1							,							_UpperCamelCase = False							,							_UpperCamelCase = None							,							)							-> Mapping[str, Any]:
							lowerCAmelCase_         =      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    _UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)

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

										lowerCAmelCase_         =      torch.cat(
										    [common_inputs["decoder_attention_mask"], torch.ones(_UpperCamelCase							,							_UpperCamelCase	)]							,							dim=1	)

										lowerCAmelCase_         =      []
										# If the number of encoder and decoder layers are present in the model configuration, both are considered
										lowerCAmelCase_      ,    lowerCAmelCase_         =      self.num_layers
										lowerCAmelCase_         =      min(_UpperCamelCase							,							_UpperCamelCase	)
										lowerCAmelCase_         =      max(_UpperCamelCase							,							_UpperCamelCase	) - min_num_layers
										lowerCAmelCase_         =      "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.
										lowerCAmelCase_         =      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]:
							lowerCAmelCase_         =      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
										lowerCAmelCase_      ,    lowerCAmelCase_         =      common_inputs["input_ids"].shape
										# Not using the same length for past_key_values
										lowerCAmelCase_         =      seqlen + 2
										lowerCAmelCase_      ,    lowerCAmelCase_         =      self.num_layers
										lowerCAmelCase_      ,    lowerCAmelCase_         =      self.num_attention_heads
										lowerCAmelCase_         =      (
										    batch,
										    num_encoder_attention_heads,
										    past_key_values_length,
										    self._config.hidden_size // num_encoder_attention_heads,
										)

										lowerCAmelCase_         =      common_inputs["attention_mask"].dtype
										lowerCAmelCase_         =      torch.cat(
										    [common_inputs["attention_mask"], torch.ones(_UpperCamelCase							,							_UpperCamelCase							,							dtype=_UpperCamelCase	)]							,							dim=1	)
										lowerCAmelCase_         =      [
										    (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]:
							# 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
							lowerCAmelCase_         =      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
							lowerCAmelCase_         =      tokenizer.num_special_tokens_to_add(_UpperCamelCase	)
							lowerCAmelCase_         =      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
							lowerCAmelCase_         =      [" ".join([tokenizer.unk_token]	) * seq_length] * batch_size
							lowerCAmelCase_         =      dict(tokenizer(_UpperCamelCase							,							return_tensors=_UpperCamelCase	)	)
							return common_inputs




				def 	__a		(   self							,							_UpperCamelCase							,							_UpperCamelCase = -1							,							_UpperCamelCase = -1							,							_UpperCamelCase = False							,							_UpperCamelCase = None							,							)							-> Mapping[str, Any]:
							if self.task in ["default", "seq2seq-lm"]:
										lowerCAmelCase_         =      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":
										lowerCAmelCase_         =      self._generate_dummy_inputs_for_causal_lm(
										    _UpperCamelCase							,							batch_size=_UpperCamelCase							,							seq_length=_UpperCamelCase							,							is_pair=_UpperCamelCase							,							framework=_UpperCamelCase	)
							else:
										lowerCAmelCase_         =      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	)							-> List[str]:
							if self.task in ["default", "seq2seq-lm"]:
										lowerCAmelCase_         =      super()._flatten_past_key_values_(_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)
							else:
										lowerCAmelCase_         =      super(_UpperCamelCase							,							self	)._flatten_past_key_values_(
										    _UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase							,							_UpperCamelCase	)





from math import ceil






def   lowerCamelCase__							(							__lowerCAmelCase    :    Dict       ,				__lowerCAmelCase    :    int							):




			"""simple docstring"""


			lowerCAmelCase_         =      list(range(0       ,				__lowerCAmelCase							)							)

			lowerCAmelCase_         =      [item for sublist in list(device_map.values()							) for item in sublist]

			# Duplicate check
			lowerCAmelCase_         =      []
			for i in device_map_blocks:
						if device_map_blocks.count(__lowerCAmelCase							) > 1 and i not in duplicate_blocks:
									duplicate_blocks.append(__lowerCAmelCase							)
    # Missing blocks
			lowerCAmelCase_         =      [i for i in blocks if i not in device_map_blocks]
			lowerCAmelCase_         =      [i for i in device_map_blocks if i not in blocks]

			if len(__lowerCAmelCase							) != 0:
						raise ValueError(
						    "Duplicate attention blocks specified in device_map. Attention blocks must be specified to one device."
						    " These attention blocks were specified more than once: " + str(__lowerCAmelCase							)							)
			if len(__lowerCAmelCase							) != 0:
						raise ValueError(
						    "There are attention blocks for this model that are not specified in the device_map. Add these attention "
						    "blocks to a device on the device_map: " + str(__lowerCAmelCase							)							)
			if len(__lowerCAmelCase							) != 0:
						raise ValueError(
						    "The device_map contains more attention blocks than this model has. Remove these from the device_map:"
						    + str(__lowerCAmelCase							)							)






def   lowerCamelCase__							(							__lowerCAmelCase    :    List[str]       ,				__lowerCAmelCase    :    Tuple							):




			"""simple docstring"""


			lowerCAmelCase_         =      list(range(__lowerCAmelCase							)							)
			lowerCAmelCase_         =      int(ceil(n_layers / len(__lowerCAmelCase							)							)							)
			lowerCAmelCase_         =      [layers[i : i + n_blocks] for i in range(0       ,				__lowerCAmelCase       ,				__lowerCAmelCase							)]

			return dict(zip(__lowerCAmelCase       ,				__lowerCAmelCase							)							)



