


"""simple docstring"""
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


SCREAMING_SNAKE_CASE :    List[Any]									=				logging.get_logger(__name__)

SCREAMING_SNAKE_CASE :    Union[str, Any]									=				{
    """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 			_UpperCAmelCase (							__snake_case				):






		'''simple docstring'''

		lowerCamelCase__				='blenderbot-small'
		lowerCamelCase__				=['past_key_values']
		lowerCamelCase__				={'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}



		def __init__(self				,       a_=5_02_65				,       a_=5_12				,       a_=8				,       a_=20_48				,       a_=16				,       a_=8				,       a_=20_48				,       a_=16				,       a_=0.0				,       a_=0.0				,       a_=True				,       a_=True				,       a_="gelu"				,       a_=5_12				,       a_=0.1				,       a_=0.0				,       a_=0.0				,       a_=0.02				,       a_=1				,       a_=False				,       a_=0				,       a_=1				,       a_=2				,       a_=2				,       **a_				,       ):





			'''simple docstring'''

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

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



class 			_UpperCAmelCase (							__snake_case				):






		'''simple docstring'''

		@property
		def 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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

				if self.use_past:
					__snake_case	:						Optional[Any]						=				{0: '''batch'''}
					__snake_case	:						Tuple						=				{0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
				else:
					__snake_case	:						Dict						=				{0: '''batch''', 1: '''decoder_sequence'''}
					__snake_case	:						Optional[Any]						=				{0: '''batch''', 1: '''decoder_sequence'''}

				if self.use_past:
					self.fill_with_past_key_values_(a_				,       direction='''inputs'''					)
			elif self.task == "causal-lm":
				# TODO: figure this case out.
				__snake_case	:						str						=				OrderedDict(
				    [
				        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
				        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
				    ]					)
				if self.use_past:
					__snake_case	,							__snake_case	:						Tuple						=				self.num_layers
					for i in range(a_					):
						__snake_case	:						int						=				{0: '''batch''', 2: '''past_sequence + sequence'''}
						__snake_case	:						int						=				{0: '''batch''', 2: '''past_sequence + sequence'''}
			else:
				__snake_case	:						Dict						=				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 				SCREAMING_SNAKE_CASE  (self					):





			'''simple docstring'''

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



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_ = -1				,       a_ = -1				,       a_ = False				,       a_ = None				,       ):





			'''simple docstring'''

			__snake_case	:						Optional[Any]						=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
			    a_				,       a_				,       a_				,       a_				,       a_					)

			# Generate decoder inputs
			__snake_case	:						str						=				seq_length if not self.use_past else 1
			__snake_case	:						Union[str, Any]						=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
			    a_				,       a_				,       a_				,       a_				,       a_					)
			__snake_case	:						List[Any]						=				{f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
			__snake_case	:						Dict						=				dict(**a_				,       **a_					)

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

				__snake_case	:						Optional[Any]						=				torch.cat(
				    [common_inputs['''decoder_attention_mask'''], torch.ones(a_				,       a_					)]				,       dim=1					)

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

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



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_ = -1				,       a_ = -1				,       a_ = False				,       a_ = None				,       ):





			'''simple docstring'''

			__snake_case	:						List[Any]						=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
			    a_				,       a_				,       a_				,       a_				,       a_					)

			if self.use_past:
				if not is_torch_available():
					raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''					)
				else:
					import torch
				__snake_case	,							__snake_case	:						Dict						=				common_inputs['''input_ids'''].shape
				# Not using the same length for past_key_values
				__snake_case	:						Optional[int]						=				seqlen + 2
				__snake_case	,							__snake_case	:						Dict						=				self.num_layers
				__snake_case	,							__snake_case	:						Optional[Any]						=				self.num_attention_heads
				__snake_case	:						Optional[int]						=				(
				    batch,
				    num_encoder_attention_heads,
				    past_key_values_length,
				    self._config.hidden_size // num_encoder_attention_heads,
				)

				__snake_case	:						str						=				common_inputs['''attention_mask'''].dtype
				__snake_case	:						Optional[int]						=				torch.cat(
				    [common_inputs['''attention_mask'''], torch.ones(a_				,       a_				,       dtype=a_					)]				,       dim=1					)
				__snake_case	:						List[str]						=				[
				    (torch.zeros(a_					), torch.zeros(a_					)) for _ in range(a_					)
				]
			return common_inputs



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_ = -1				,       a_ = -1				,       a_ = False				,       a_ = None				,       ):





			'''simple docstring'''

			__snake_case	:						Dict						=				compute_effective_axis_dimension(
			    a_				,       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
			__snake_case	:						List[Any]						=				tokenizer.num_special_tokens_to_add(a_					)
			__snake_case	:						Union[str, Any]						=				compute_effective_axis_dimension(
			    a_				,       fixed_dimension=OnnxConfig.default_fixed_sequence				,       num_token_to_add=a_					)

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



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_ = -1				,       a_ = -1				,       a_ = False				,       a_ = None				,       ):





			'''simple docstring'''

			if self.task in ["default", "seq2seq-lm"]:
				__snake_case	:						str						=				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
				    a_				,       batch_size=a_				,       seq_length=a_				,       is_pair=a_				,       framework=a_					)

			elif self.task == "causal-lm":
				__snake_case	:						str						=				self._generate_dummy_inputs_for_causal_lm(
				    a_				,       batch_size=a_				,       seq_length=a_				,       is_pair=a_				,       framework=a_					)
			else:
				__snake_case	:						int						=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
				    a_				,       batch_size=a_				,       seq_length=a_				,       is_pair=a_				,       framework=a_					)

			return common_inputs



		def 				SCREAMING_SNAKE_CASE  (self				,       a_				,       a_				,       a_				,       a_					):





			'''simple docstring'''

			if self.task in ["default", "seq2seq-lm"]:
				__snake_case	:						Tuple						=				super()._flatten_past_key_values_(a_				,       a_				,       a_				,       a_					)
			else:
				__snake_case	:						str						=				super(a_				,       self					)._flatten_past_key_values_(
				    a_				,       a_				,       a_				,       a_					)













import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFAutoModel, is_tensorflow_text_available, is_tf_available
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.testing_utils import require_tensorflow_text, require_tf, slow


if is_tf_available():
	import tensorflow as tf

if is_tensorflow_text_available():
	from transformers.models.bert import TFBertTokenizer


UpperCAmelCase_ :						Optional[Any]      =      ['''bert-base-uncased''', '''bert-base-cased''']
UpperCAmelCase_ :						List[str]      =      '''hf-internal-testing/tiny-bert-tf-only'''

if is_tf_available():

	class     _SCREAMING_SNAKE_CASE       (			tf.keras.Model      ):





		def __init__(		self				:					List[str] ,  __lowerCamelCase				:					Union[str, Any]	):
			super().__init__()
			UpperCamelCase     :Any          =	tokenizer
			UpperCamelCase     :List[str]          =	AutoConfig.from_pretrained(__lowerCamelCase	)
			UpperCamelCase     :List[str]          =	TFAutoModel.from_config(__lowerCamelCase	)





		def 						_A (		self				:					Tuple ,  __lowerCamelCase				:					str	):
			UpperCamelCase     :str          =	self.tokenizer(__lowerCamelCase	)
			UpperCamelCase     :Any          =	self.bert(**__lowerCamelCase	)
			return out["pooler_output"]

@require_tf
@require_tensorflow_text
class     _SCREAMING_SNAKE_CASE       (			unittest.TestCase      ):





	def 						_A (		self				:					Dict	):
		super().setUp()

		UpperCamelCase     :int          =	[
		    BertTokenizer.from_pretrained(__lowerCamelCase	) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
		]  # repeat for when fast_bert_tokenizer=false
		UpperCamelCase     :Any          =	[TFBertTokenizer.from_pretrained(__lowerCamelCase	) for checkpoint in TOKENIZER_CHECKPOINTS] + [
		    TFBertTokenizer.from_pretrained(__lowerCamelCase ,  use_fast_bert_tokenizer=__lowerCamelCase	)
		    for checkpoint in TOKENIZER_CHECKPOINTS
		]
		assert len(self.tokenizers	) == len(self.tf_tokenizers	)

		UpperCamelCase     :Any          =	[
		    """This is a straightforward English test sentence.""",
		    """This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.""",
		    """Now we're going to add some Chinese: 一 二 三 一二三""",
		    """And some much more rare Chinese: 齉 堃 齉堃""",
		    """Je vais aussi écrire en français pour tester les accents""",
		    """Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ""",
		]
		UpperCamelCase     :Union[str, Any]          =	list(zip(self.test_sentences ,  self.test_sentences[::-1]	)	)





	def 						_A (		self				:					Optional[int]	):
		for tokenizer, tf_tokenizer in zip(self.tokenizers ,  self.tf_tokenizers	):
			for test_inputs in (self.test_sentences, self.paired_sentences):
				UpperCamelCase     :Any          =	tokenizer(__lowerCamelCase ,  return_tensors="""tf""" ,  padding="""longest"""	)
				UpperCamelCase     :str          =	tf_tokenizer(__lowerCamelCase	)

				for key in python_outputs.keys():
					self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape	)	)
					self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key] ,  tf.intaa	) == tf_outputs[key]	)	)





	@slow
	def 						_A (		self				:					Dict	):
		for tf_tokenizer in self.tf_tokenizers:
			UpperCamelCase     :str          =	tf_tokenizer(self.paired_sentences	)
			UpperCamelCase     :Any          =	tf_tokenizer(
			    text=[sentence[0] for sentence in self.paired_sentences] ,  text_pair=[sentence[1] for sentence in self.paired_sentences] ,  )
			for key in merged_outputs.keys():
				self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key] ,  tf.intaa	) == separated_outputs[key]	)	)





	@slow
	def 						_A (		self				:					List[str]	):
		for tf_tokenizer in self.tf_tokenizers:
			UpperCamelCase     :List[Any]          =	tf.function(__lowerCamelCase	)
			for test_inputs in (self.test_sentences, self.paired_sentences):
				UpperCamelCase     :Any          =	tf.constant(__lowerCamelCase	)
				UpperCamelCase     :List[str]          =	compiled_tokenizer(__lowerCamelCase	)
				UpperCamelCase     :Optional[Any]          =	tf_tokenizer(__lowerCamelCase	)

				for key in eager_outputs.keys():
					self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]	)	)





	@slow
	def 						_A (		self				:					Tuple	):
		for tf_tokenizer in self.tf_tokenizers:
			UpperCamelCase     :List[str]          =	ModelToSave(tokenizer=__lowerCamelCase	)
			UpperCamelCase     :Union[str, Any]          =	tf.convert_to_tensor(self.test_sentences	)
			UpperCamelCase     :Union[str, Any]          =	model(__lowerCamelCase	)  # Build model with some sample inputs
			with TemporaryDirectory() as tempdir:
				UpperCamelCase     :List[str]          =	Path(__lowerCamelCase	) / """saved.model"""
				model.save(__lowerCamelCase	)
				UpperCamelCase     :List[Any]          =	tf.keras.models.load_model(__lowerCamelCase	)
			UpperCamelCase     :Dict          =	loaded_model(__lowerCamelCase	)
			# We may see small differences because the loaded model is compiled, so we need an epsilon for the test
			self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output	)	) ,  1E-5	)
