



from __future__ import annotations

import unittest

from transformers import BlenderbotSmallConfig, BlenderbotSmallTokenizer, is_tf_available
from transformers.testing_utils import require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import tensorflow as tf

			from transformers import TFAutoModelForSeqaSeqLM, TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel
@require_tf
class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  Union[str, Any]      =			BlenderbotSmallConfig
			lowercase__      :  List[Any]      =			{}
			lowercase__      :  int      =			"""gelu"""





			def __init__(				self	:       Optional[Any]			,     lowercase	:       List[str]			,     lowercase	:       List[Any]=13			,     lowercase	:       List[str]=7			,     lowercase	:       str=True			,     lowercase	:       List[Any]=False			,     lowercase	:       Optional[Any]=99			,     lowercase	:       Optional[Any]=32			,     lowercase	:       Any=2			,     lowercase	:       Union[str, Any]=4			,     lowercase	:       Tuple=37			,     lowercase	:       Optional[int]=0.1			,     lowercase	:       List[Any]=0.1			,     lowercase	:       Any=20			,     lowercase	:       Dict=2			,     lowercase	:       Optional[int]=1			,     lowercase	:       str=0			,     )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase					=				parent
									__lowercase					=				batch_size
									__lowercase					=				seq_length
									__lowercase					=				is_training
									__lowercase					=				use_labels
									__lowercase					=				vocab_size
									__lowercase					=				hidden_size
									__lowercase					=				num_hidden_layers
									__lowercase					=				num_attention_heads
									__lowercase					=				intermediate_size

									__lowercase					=				hidden_dropout_prob
									__lowercase					=				attention_probs_dropout_prob
									__lowercase					=				max_position_embeddings
									__lowercase					=				eos_token_id
									__lowercase					=				pad_token_id
									__lowercase					=				bos_token_id





			def 		snake_case__			(				self	:       Optional[Any]   )						->					Any:
									"""simple docstring"""







									__lowercase					=				ids_tensor([self.batch_size, self.seq_length - 1]			,     self.vocab_size   )
									__lowercase					=				tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size   )			,     1   )
									__lowercase					=				tf.concat([input_ids, eos_tensor]			,     axis=1   )

									__lowercase					=				ids_tensor([self.batch_size, self.seq_length]			,     self.vocab_size   )

									__lowercase					=				self.config_cls(
									    vocab_size=self.vocab_size			,     d_model=self.hidden_size			,     encoder_layers=self.num_hidden_layers			,     decoder_layers=self.num_hidden_layers			,     encoder_attention_heads=self.num_attention_heads			,     decoder_attention_heads=self.num_attention_heads			,     encoder_ffn_dim=self.intermediate_size			,     decoder_ffn_dim=self.intermediate_size			,     dropout=self.hidden_dropout_prob			,     attention_dropout=self.attention_probs_dropout_prob			,     max_position_embeddings=self.max_position_embeddings			,     eos_token_ids=[2]			,     bos_token_id=self.bos_token_id			,     pad_token_id=self.pad_token_id			,     decoder_start_token_id=self.pad_token_id			,     **self.config_updates			,     )
									__lowercase					=				prepare_blenderbot_small_inputs_dict(lowercase			,     lowercase			,     lowercase   )
									return config, inputs_dict







			def 		snake_case__			(				self	:       Optional[int]			,     lowercase	:       Optional[Any]			,     lowercase	:       Any   )						->					Dict:
									"""simple docstring"""







									__lowercase					=				TFBlenderbotSmallModel(config=lowercase   ).get_decoder()
									__lowercase					=				inputs_dict["""input_ids"""]

									__lowercase					=				input_ids[:1, :]
									__lowercase					=				inputs_dict["""attention_mask"""][:1, :]
									__lowercase					=				inputs_dict["""head_mask"""]
									__lowercase					=				1

									# first forward pass
									__lowercase					=				model(lowercase			,     attention_mask=lowercase			,     head_mask=lowercase			,     use_cache=lowercase   )

									__lowercase	,			__lowercase					=				outputs.to_tuple()

									# create hypothetical next token and extent to next_input_ids
									__lowercase					=				ids_tensor((self.batch_size, 3)			,     config.vocab_size   )
									__lowercase					=				tf.cast(ids_tensor((self.batch_size, 3)			,     2   )			,     tf.inta   )

									# append to next input_ids and
									__lowercase					=				tf.concat([input_ids, next_tokens]			,     axis=-1   )
									__lowercase					=				tf.concat([attention_mask, next_attn_mask]			,     axis=-1   )

									__lowercase					=				model(lowercase			,     attention_mask=lowercase   )[0]
									__lowercase					=				model(lowercase			,     attention_mask=lowercase			,     past_key_values=lowercase   )[0]

									self.parent.assertEqual(next_tokens.shape[1]			,     output_from_past.shape[1]   )

									# select random slice
									__lowercase					=				int(ids_tensor((1,)			,     output_from_past.shape[-1]   )   )
									__lowercase					=				output_from_no_past[:, -3:, random_slice_idx]
									__lowercase					=				output_from_past[:, :, random_slice_idx]

									# test that outputs are equal for slice
									tf.debugging.assert_near(lowercase			,     lowercase			,     rtol=1E-3   )






def  UpperCAmelCase__    (    lowercase__						,       lowercase__						,       lowercase__						,       lowercase__=None						,       lowercase__=None						,       lowercase__=None						,       lowercase__=None						,       lowercase__=None						,       ) ->      int:
						if attention_mask is None:
												__lowercase					=				tf.cast(tf.math.not_equal(lowercase__						,       config.pad_token_id   )						,       tf.inta   )
						if decoder_attention_mask is None:
												__lowercase					=				tf.concat(
												    [
												        tf.ones(decoder_input_ids[:, :1].shape						,       dtype=tf.inta   ),
												        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]						,       config.pad_token_id   )						,       tf.inta   ),
												    ]						,       axis=-1						,       )
						if head_mask is None:
												__lowercase					=				tf.ones((config.encoder_layers, config.encoder_attention_heads)   )
						if decoder_head_mask is None:
												__lowercase					=				tf.ones((config.decoder_layers, config.decoder_attention_heads)   )
						if cross_attn_head_mask is None:
												__lowercase					=				tf.ones((config.decoder_layers, config.decoder_attention_heads)   )
						return {
						    "input_ids": input_ids,
						    "decoder_input_ids": decoder_input_ids,
						    "attention_mask": attention_mask,
						    "decoder_attention_mask": decoder_attention_mask,
						    "head_mask": head_mask,
						    "decoder_head_mask": decoder_head_mask,
						    "cross_attn_head_mask": cross_attn_head_mask,
						}
@require_tf
class 			_lowerCAmelCase   ( _UpperCAmelCase    ,					_UpperCAmelCase    ,					unittest.TestCase			):






			"""simple docstring"""





			lowercase__      :  List[str]      =			(
			    (TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel) if is_tf_available() else ()
			)
			lowercase__      :  List[str]      =			(TFBlenderbotSmallForConditionalGeneration,) if is_tf_available() else ()
			lowercase__      :  Optional[Any]      =			(
			    {
			        """conversational""": TFBlenderbotSmallForConditionalGeneration,
			        """feature-extraction""": TFBlenderbotSmallModel,
			        """summarization""": TFBlenderbotSmallForConditionalGeneration,
			        """text2text-generation""": TFBlenderbotSmallForConditionalGeneration,
			        """translation""": TFBlenderbotSmallForConditionalGeneration,
			    }
			    if is_tf_available()
			    else {}
			)
			lowercase__      :  str      =			True
			lowercase__      :  Any      =			False
			lowercase__      :  int      =			False





			def 		snake_case__			(				self	:       Optional[int]   )						->					Optional[int]:
									"""simple docstring"""







									__lowercase					=				TFBlenderbotSmallModelTester(self   )
									__lowercase					=				ConfigTester(self			,     config_class=lowercase   )





			def 		snake_case__			(				self	:       int   )						->					Tuple:
									"""simple docstring"""







									self.config_tester.run_common_tests()







			def 		snake_case__			(				self	:       Union[str, Any]   )						->					Any:
									"""simple docstring"""







									__lowercase					=				self.model_tester.prepare_config_and_inputs_for_common()
									self.model_tester.check_decoder_model_past_large_inputs(*lowercase   )





@require_tokenizers
@require_tf
class 			_lowerCAmelCase   ( unittest.TestCase			):






			"""simple docstring"""





			lowercase__      :  List[str]      =			[
			    """Social anxiety\nWow, I am never shy. Do you have anxiety?\nYes. I end up sweating and blushing and feel like """
			    """  i'm going to throw up.\nand why is that?"""
			]
			lowercase__      :  Optional[Any]      =			"""facebook/blenderbot_small-90M"""





			@cached_property
			def 		snake_case__			(				self	:       str   )						->					Union[str, Any]:
									"""simple docstring"""







									return BlenderbotSmallTokenizer.from_pretrained("""facebook/blenderbot-90M"""   )





			@cached_property
			def 		snake_case__			(				self	:       Any   )						->					List[str]:
									"""simple docstring"""







									__lowercase					=				TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name   )
									return model







			@slow
			def 		snake_case__			(				self	:       Any   )						->					List[str]:
									"""simple docstring"""







									__lowercase					=				self.tokenizer(self.src_text			,     return_tensors="""tf"""   )
									__lowercase					=				self.model.generate(
									    model_inputs.input_ids			,     attention_mask=model_inputs.attention_mask			,     num_beams=2			,     use_cache=lowercase			,     )
									__lowercase					=				self.tokenizer.batch_decode(generated_ids.numpy()			,     skip_special_tokens=lowercase   )[0]
									assert generated_words in (
									    "i don't know. i just feel like i'm going to throw up. it's not fun.",
									    "i'm not sure. i just feel like i've been feeling like i have to be in a certain place",
									    "i'm not sure. i just feel like i've been in a bad situation.",
									)










from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__ =						logging.get_logger(__name__)

# TODO Update this
UpperCamelCase__ =						{
    "facebook/esm-1b": "https://huggingface.co/facebook/esm-1b/resolve/main/config.json",
    # See all ESM models at https://huggingface.co/models?filter=esm
}
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  Union[str, Any]      =			"""esm"""





			def __init__(				self	:       Any			,     lowercase	:       Optional[Any]=None			,     lowercase	:       Optional[int]=None			,     lowercase	:       List[Any]=None			,     lowercase	:       Optional[int]=768			,     lowercase	:       str=12			,     lowercase	:       Union[str, Any]=12			,     lowercase	:       Dict=3_072			,     lowercase	:       Optional[int]=0.1			,     lowercase	:       str=0.1			,     lowercase	:       Dict=1_026			,     lowercase	:       Tuple=0.02			,     lowercase	:       str=1E-1_2			,     lowercase	:       Dict="absolute"			,     lowercase	:       Optional[Any]=True			,     lowercase	:       int=None			,     lowercase	:       int=False			,     lowercase	:       List[str]=False			,     lowercase	:       Tuple=None			,     lowercase	:       Tuple=None			,     **lowercase	:       Union[str, Any]			,     )						->					int:
									"""simple docstring"""







									super().__init__(pad_token_id=lowercase			,     mask_token_id=lowercase			,     **lowercase   )

									__lowercase					=				vocab_size
									__lowercase					=				hidden_size
									__lowercase					=				num_hidden_layers
									__lowercase					=				num_attention_heads
									__lowercase					=				intermediate_size
									__lowercase					=				hidden_dropout_prob
									__lowercase					=				attention_probs_dropout_prob
									__lowercase					=				max_position_embeddings
									__lowercase					=				initializer_range
									__lowercase					=				layer_norm_eps
									__lowercase					=				position_embedding_type
									__lowercase					=				use_cache
									__lowercase					=				emb_layer_norm_before
									__lowercase					=				token_dropout
									__lowercase					=				is_folding_model
									if is_folding_model:
															if esmfold_config is None:
																					logger.info("""No esmfold_config supplied for folding model, using default values."""   )
																					__lowercase					=				EsmFoldConfig()
															elif isinstance(lowercase			,     lowercase   ):
																					__lowercase					=				EsmFoldConfig(**lowercase   )
															__lowercase					=				esmfold_config
															if vocab_list is None:
																					logger.warning("""No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"""   )
																					__lowercase					=				get_default_vocab_list()
															else:
																					__lowercase					=				vocab_list
									else:
															__lowercase					=				None
															__lowercase					=				None
									if self.esmfold_config is not None and getattr(self.esmfold_config			,     """use_esm_attn_map"""			,     lowercase   ):
															raise ValueError("""The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"""   )







			def 		snake_case__			(				self	:       Optional[Any]   )						->					Tuple:
									"""simple docstring"""







									__lowercase					=				super().to_dict()
									if isinstance(self.esmfold_config			,     lowercase   ):
															__lowercase					=				self.esmfold_config.to_dict()
									return output





@dataclass
class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  str      =			None
			lowercase__      :  bool      =			True
			lowercase__      :  bool      =			False
			lowercase__      :  bool      =			False
			lowercase__      :  bool      =			False
			lowercase__      :  float      =			0

			lowercase__      :  bool      =			True
			lowercase__      :  bool      =			False

			lowercase__      :  int      =			128
			lowercase__      :  "TrunkConfig"      =			None





			def 		snake_case__			(				self	:       List[str]   )						->					Any:
									"""simple docstring"""







									if self.trunk is None:
															__lowercase					=				TrunkConfig()
									elif isinstance(self.trunk			,     lowercase   ):
															__lowercase					=				TrunkConfig(**self.trunk   )







			def 		snake_case__			(				self	:       Dict   )						->					Any:
									"""simple docstring"""







									__lowercase					=				asdict(self   )
									__lowercase					=				self.trunk.to_dict()
									return output





@dataclass
class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  int      =			48
			lowercase__      :  int      =			1_024
			lowercase__      :  int      =			128
			lowercase__      :  int      =			32
			lowercase__      :  int      =			32
			lowercase__      :  int      =			32
			lowercase__      :  float      =			0
			lowercase__      :  float      =			0
			lowercase__      :  bool      =			False
			lowercase__      :  int      =			4
			lowercase__      :  Optional[int]      =			128
			lowercase__      :  "StructureModuleConfig"      =			None





			def 		snake_case__			(				self	:       Tuple   )						->					str:
									"""simple docstring"""







									if self.structure_module is None:
															__lowercase					=				StructureModuleConfig()
									elif isinstance(self.structure_module			,     lowercase   ):
															__lowercase					=				StructureModuleConfig(**self.structure_module   )

									if self.max_recycles <= 0:
															raise ValueError(F"`max_recycles` should be positive, got {self.max_recycles}."   )
									if self.sequence_state_dim % self.sequence_state_dim != 0:
															raise ValueError(
															    """`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"""
															    F" {self.sequence_state_dim} and {self.sequence_state_dim}."   )
									if self.pairwise_state_dim % self.pairwise_state_dim != 0:
															raise ValueError(
															    """`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"""
															    F" {self.pairwise_state_dim} and {self.pairwise_state_dim}."   )

									__lowercase					=				self.sequence_state_dim // self.sequence_head_width
									__lowercase					=				self.pairwise_state_dim // self.pairwise_head_width

									if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
															raise ValueError(
															    """`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"""
															    F" {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}."   )
									if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
															raise ValueError(
															    """`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"""
															    F" {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}."   )
									if self.pairwise_state_dim % 2 != 0:
															raise ValueError(F"`pairwise_state_dim` should be even, got {self.pairwise_state_dim}."   )

									if self.dropout >= 0.4:
															raise ValueError(F"`dropout` should not be greater than 0.4, got {self.dropout}."   )







			def 		snake_case__			(				self	:       List[str]   )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase					=				asdict(self   )
									__lowercase					=				self.structure_module.to_dict()
									return output





@dataclass
class 			_lowerCAmelCase   :






			"""simple docstring"""





			lowercase__      :  int      =			384
			lowercase__      :  int      =			128
			lowercase__      :  int      =			16
			lowercase__      :  int      =			128
			lowercase__      :  int      =			12
			lowercase__      :  int      =			4
			lowercase__      :  int      =			8
			lowercase__      :  float      =			0.1
			lowercase__      :  int      =			8
			lowercase__      :  int      =			1
			lowercase__      :  int      =			2
			lowercase__      :  int      =			7
			lowercase__      :  int      =			10
			lowercase__      :  float      =			1E-8
			lowercase__      :  float      =			1E5





			def 		snake_case__			(				self	:       Tuple   )						->					List[str]:
									"""simple docstring"""







									return asdict(self   )






def  UpperCAmelCase__    (    ) ->      List[Any]:
						return (
						    "<cls>",
						    "<pad>",
						    "<eos>",
						    "<unk>",
						    "L",
						    "A",
						    "G",
						    "V",
						    "S",
						    "E",
						    "R",
						    "T",
						    "I",
						    "D",
						    "P",
						    "K",
						    "Q",
						    "N",
						    "F",
						    "Y",
						    "M",
						    "H",
						    "W",
						    "C",
						    "X",
						    "B",
						    "U",
						    "Z",
						    "O",
						    ".",
						    "-",
						    "<null_1>",
						    "<mask>",
						)





