

"""simple docstring"""






from typing import Tuple, Union

from ...modeling_outputs import BackboneOutput
from ...modeling_utils import PreTrainedModel
from ...utils import is_timm_available, is_torch_available, requires_backends
from ...utils.backbone_utils import BackboneMixin
from .configuration_timm_backbone import TimmBackboneConfig


if is_timm_available():
	import timm


if is_torch_available():
	from torch import Tensor







class       snake_case__  ( lowerCAmelCase_					,  lowerCAmelCase_     ):
			SCREAMING_SNAKE_CASE__     				=						'''pixel_values'''
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						TimmBackboneConfig







			def __init__(			self			: Optional[int]			,  lowercase			: Optional[Any]			,  **lowercase			: List[str]  ):





								'''simple docstring'''



								requires_backends(self			,  "timm"  )
								super().__init__(lowercase  )
								UpperCAmelCase			:   Optional[int]      =				config

								if config.backbone is None:
													raise ValueError("backbone is not set in the config. Please set it to a timm model name."  )

								if config.backbone not in timm.list_models():
													raise ValueError(f"""backbone {config.backbone} is not supported by timm."""  )

								if hasattr(lowercase			,  "out_features"  ) and config.out_features is not None:
													raise ValueError("out_features is not supported by TimmBackbone. Please use out_indices instead."  )

								UpperCAmelCase			:   Dict      =				getattr(lowercase			,  "use_pretrained_backbone"			,  lowercase  )
								if pretrained is None:
													raise ValueError("use_pretrained_backbone is not set in the config. Please set it to True or False."  )

								# We just take the final layer by default. This matches the default for the transformers models.
								UpperCAmelCase			:   Optional[Any]      =				config.out_indices if getattr(lowercase			,  "out_indices"			,  lowercase  ) is not None else (-1,)

								UpperCAmelCase			:   Any      =				timm.create_model(
								    config.backbone			,  pretrained=lowercase			,  features_only=config.features_only			,  in_chans=config.num_channels			,  out_indices=lowercase			,  **lowercase			,  )
								# These are used to control the output of the model when called. If output_hidden_states is True, then
								# return_layers is modified to include all layers.
								UpperCAmelCase			:   Optional[int]      =				self._backbone.return_layers
								UpperCAmelCase			:   str      =				{layer["module"]: str(lowercase  ) for i, layer in enumerate(self._backbone.feature_info.info  )}
								super()._init_backbone(lowercase  )







			@classmethod
			def 			__lowerCAmelCase				(			cls			: Dict			,  lowercase			: Optional[int]			,  *lowercase			: Any			,  **lowercase			: Optional[Any]  ):





								'''simple docstring'''



								requires_backends(cls			,  ["vision", "timm"]  )
								from ...models.timm_backbone import TimmBackboneConfig

								UpperCAmelCase			:   Union[str, Any]      =				kwargs.pop("config"			,  TimmBackboneConfig()  )

								UpperCAmelCase			:   Tuple      =				kwargs.pop("use_timm_backbone"			,  lowercase  )
								if not use_timm:
													raise ValueError("use_timm_backbone must be True for timm backbones"  )

								UpperCAmelCase			:   Tuple      =				kwargs.pop("num_channels"			,  config.num_channels  )
								UpperCAmelCase			:   List[Any]      =				kwargs.pop("features_only"			,  config.features_only  )
								UpperCAmelCase			:   str      =				kwargs.pop("use_pretrained_backbone"			,  config.use_pretrained_backbone  )
								UpperCAmelCase			:   int      =				kwargs.pop("out_indices"			,  config.out_indices  )
								UpperCAmelCase			:   Union[str, Any]      =				TimmBackboneConfig(
								    backbone=lowercase			,  num_channels=lowercase			,  features_only=lowercase			,  use_pretrained_backbone=lowercase			,  out_indices=lowercase			,  )
								return super()._from_config(lowercase			,  **lowercase  )







			def 			__lowerCAmelCase				(			self			: Any			,  lowercase			: Dict  ):





								'''simple docstring'''



								pass







			def 			__lowerCAmelCase				(			self			: Any			,  lowercase			: List[Any]			,  lowercase			: Any=None			,  lowercase			: Union[str, Any]=None			,  lowercase			: str=None			,  **lowercase			: Union[str, Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   str      =				return_dict if return_dict is not None else self.config.use_return_dict
								UpperCAmelCase			:   Optional[int]      =				(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								UpperCAmelCase			:   str      =				output_attentions if output_attentions is not None else self.config.output_attentions

								if output_attentions:
													raise ValueError("Cannot output attentions for timm backbones at the moment"  )

								if output_hidden_states:
													# We modify the return layers to include all the stages of the backbone
													UpperCAmelCase			:   int      =				self._all_layers
													UpperCAmelCase			:   Optional[int]      =				self._backbone(lowercase			,  **lowercase  )
													UpperCAmelCase			:   Any      =				self._return_layers
													UpperCAmelCase			:   str      =				tuple(hidden_states[i] for i in self.out_indices  )
								else:
													UpperCAmelCase			:   List[Any]      =				self._backbone(lowercase			,  **lowercase  )
													UpperCAmelCase			:   Any      =				None

								UpperCAmelCase			:   List[str]      =				tuple(lowercase  )
								UpperCAmelCase			:   List[Any]      =				tuple(lowercase  ) if hidden_states is not None else None

								if not return_dict:
													UpperCAmelCase			:   Optional[int]      =				(feature_maps,)
													if output_hidden_states:
																		UpperCAmelCase			:   Dict      =				output + (hidden_states,)
													return output

								return BackboneOutput(feature_maps=lowercase			,  hidden_states=lowercase			,  attentions=lowercase  )





"""simple docstring"""






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       snake_case__  :
			SCREAMING_SNAKE_CASE__     				=						BlenderbotSmallConfig
			SCREAMING_SNAKE_CASE__     				=						{}
			SCREAMING_SNAKE_CASE__     				=						'''gelu'''







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





								'''simple docstring'''



								UpperCAmelCase			:   Dict      =				parent
								UpperCAmelCase			:   Union[str, Any]      =				batch_size
								UpperCAmelCase			:   Tuple      =				seq_length
								UpperCAmelCase			:   Dict      =				is_training
								UpperCAmelCase			:   str      =				use_labels
								UpperCAmelCase			:   Union[str, Any]      =				vocab_size
								UpperCAmelCase			:   Tuple      =				hidden_size
								UpperCAmelCase			:   List[Any]      =				num_hidden_layers
								UpperCAmelCase			:   Union[str, Any]      =				num_attention_heads
								UpperCAmelCase			:   Dict      =				intermediate_size

								UpperCAmelCase			:   Optional[int]      =				hidden_dropout_prob
								UpperCAmelCase			:   Optional[Any]      =				attention_probs_dropout_prob
								UpperCAmelCase			:   Dict      =				max_position_embeddings
								UpperCAmelCase			:   List[Any]      =				eos_token_id
								UpperCAmelCase			:   List[Any]      =				pad_token_id
								UpperCAmelCase			:   Union[str, Any]      =				bos_token_id







			def 			__lowerCAmelCase				(			self			: Union[str, Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[str]      =				ids_tensor([self.batch_size, self.seq_length - 1]			,  self.vocab_size  )
								UpperCAmelCase			:   Optional[int]      =				tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size  )			,  1  )
								UpperCAmelCase			:   Optional[int]      =				tf.concat([input_ids, eos_tensor]			,  axis=1  )

								UpperCAmelCase			:   str      =				ids_tensor([self.batch_size, self.seq_length]			,  self.vocab_size  )

								UpperCAmelCase			:   int      =				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			,  )
								UpperCAmelCase			:   Optional[Any]      =				prepare_blenderbot_small_inputs_dict(lowercase			,  lowercase			,  lowercase  )
								return config, inputs_dict







			def 			__lowerCAmelCase				(			self			: int			,  lowercase			: List[Any]			,  lowercase			: Optional[Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[Any]      =				TFBlenderbotSmallModel(config=lowercase  ).get_decoder()
								UpperCAmelCase			:   List[Any]      =				inputs_dict["input_ids"]

								UpperCAmelCase			:   int      =				input_ids[:1, :]
								UpperCAmelCase			:   Union[str, Any]      =				inputs_dict["attention_mask"][:1, :]
								UpperCAmelCase			:   Union[str, Any]      =				inputs_dict["head_mask"]
								UpperCAmelCase			:   List[Any]      =				1

								# first forward pass
								UpperCAmelCase			:   Optional[Any]      =				model(lowercase			,  attention_mask=lowercase			,  head_mask=lowercase			,  use_cache=lowercase  )

								UpperCAmelCase      ,							UpperCAmelCase			:   List[str]      =				outputs.to_tuple()

								# create hypothetical next token and extent to next_input_ids
								UpperCAmelCase			:   Any      =				ids_tensor((self.batch_size, 3)			,  config.vocab_size  )
								UpperCAmelCase			:   Optional[int]      =				tf.cast(ids_tensor((self.batch_size, 3)			,  2  )			,  tf.inta  )

								# append to next input_ids and
								UpperCAmelCase			:   Union[str, Any]      =				tf.concat([input_ids, next_tokens]			,  axis=-1  )
								UpperCAmelCase			:   str      =				tf.concat([attention_mask, next_attn_mask]			,  axis=-1  )

								UpperCAmelCase			:   Optional[Any]      =				model(lowercase			,  attention_mask=lowercase  )[0]
								UpperCAmelCase			:   List[str]      =				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
								UpperCAmelCase			:   List[str]      =				int(ids_tensor((1,)			,  output_from_past.shape[-1]  )  )
								UpperCAmelCase			:   Optional[int]      =				output_from_no_past[:, -3:, random_slice_idx]
								UpperCAmelCase			:   Optional[Any]      =				output_from_past[:, :, random_slice_idx]

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



def   lowercase_	(    _lowercase       :		Optional[Any]							,      _lowercase       :		Optional[int]							,      _lowercase       :		Tuple							,      _lowercase       :		Tuple=None							,      _lowercase       :		Any=None							,      _lowercase       :		str=None							,      _lowercase       :		Optional[Any]=None							,      _lowercase       :		Dict=None							,      ):




					'''simple docstring'''






					if attention_mask is None:
										UpperCAmelCase			:   List[str]      =				tf.cast(tf.math.not_equal(_lowercase							,      config.pad_token_id	)							,      tf.inta	)
					if decoder_attention_mask is None:
										UpperCAmelCase			:   str      =				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:
										UpperCAmelCase			:   Dict      =				tf.ones((config.encoder_layers, config.encoder_attention_heads)	)
					if decoder_head_mask is None:
										UpperCAmelCase			:   Optional[Any]      =				tf.ones((config.decoder_layers, config.decoder_attention_heads)	)
					if cross_attn_head_mask is None:
										UpperCAmelCase			:   str      =				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       snake_case__  ( lowerCAmelCase_					,  lowerCAmelCase_					,  unittest.TestCase     ):
			SCREAMING_SNAKE_CASE__     				=						(
			    (TFBlenderbotSmallForConditionalGeneration, TFBlenderbotSmallModel) if is_tf_available() else ()
			)
			SCREAMING_SNAKE_CASE__     				=						(TFBlenderbotSmallForConditionalGeneration,) if is_tf_available() else ()
			SCREAMING_SNAKE_CASE__     				=						(
			    {
			        '''conversational''': TFBlenderbotSmallForConditionalGeneration,
			        '''feature-extraction''': TFBlenderbotSmallModel,
			        '''summarization''': TFBlenderbotSmallForConditionalGeneration,
			        '''text2text-generation''': TFBlenderbotSmallForConditionalGeneration,
			        '''translation''': TFBlenderbotSmallForConditionalGeneration,
			    }
			    if is_tf_available()
			    else {}
			)
			SCREAMING_SNAKE_CASE__     				=						True
			SCREAMING_SNAKE_CASE__     				=						False
			SCREAMING_SNAKE_CASE__     				=						False







			def 			__lowerCAmelCase				(			self			: Optional[Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Dict      =				TFBlenderbotSmallModelTester(self  )
								UpperCAmelCase			:   Dict      =				ConfigTester(self			,  config_class=lowercase  )







			def 			__lowerCAmelCase				(			self			: Optional[Any]  ):





								'''simple docstring'''



								self.config_tester.run_common_tests()







			def 			__lowerCAmelCase				(			self			: Union[str, Any]  ):





								'''simple docstring'''



								UpperCAmelCase			:   Dict      =				self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.check_decoder_model_past_large_inputs(*lowercase  )







@require_tokenizers
@require_tf
class       snake_case__  ( unittest.TestCase     ):
			SCREAMING_SNAKE_CASE__     				=						[
			    '''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?'''
			]
			SCREAMING_SNAKE_CASE__     				=						'''facebook/blenderbot_small-90M'''







			@cached_property
			def 			__lowerCAmelCase				(			self			: Any  ):





								'''simple docstring'''



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







			@cached_property
			def 			__lowerCAmelCase				(			self			: Any  ):





								'''simple docstring'''



								UpperCAmelCase			:   Union[str, Any]      =				TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name  )
								return model







			@slow
			def 			__lowerCAmelCase				(			self			: Optional[int]  ):





								'''simple docstring'''



								UpperCAmelCase			:   List[Any]      =				self.tokenizer(self.src_text			,  return_tensors="tf"  )
								UpperCAmelCase			:   List[Any]      =				self.model.generate(
								    model_inputs.input_ids			,  attention_mask=model_inputs.attention_mask			,  num_beams=2			,  use_cache=lowercase			,  )
								UpperCAmelCase			:   Dict      =				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.",
								)


