




'''simple docstring'''





from __future__ import annotations

import unittest

from transformers import BlenderbotConfig, BlenderbotTokenizer, 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, TFBlenderbotForConditionalGeneration, TFBlenderbotModel






@require_tf
class  _lowerCAmelCase     :
      A__		     =		BlenderbotConfig
      A__		     =		{}
      A__		     =		'gelu'

      def __init__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase=13      ,				__UpperCAmelCase=7      ,				__UpperCAmelCase=True      ,				__UpperCAmelCase=False      ,				__UpperCAmelCase=99      ,				__UpperCAmelCase=32      ,				__UpperCAmelCase=2      ,				__UpperCAmelCase=4      ,				__UpperCAmelCase=37      ,				__UpperCAmelCase=0.1      ,				__UpperCAmelCase=0.1      ,				__UpperCAmelCase=20      ,				__UpperCAmelCase=2      ,				__UpperCAmelCase=1      ,				__UpperCAmelCase=0      ,				):
        lowerCAmelCase__				:  Any		       =     parent
        lowerCAmelCase__				:  List[Any]		       =     batch_size
        lowerCAmelCase__				:  List[str]		       =     seq_length
        lowerCAmelCase__				:  Union[str, Any]		       =     is_training
        lowerCAmelCase__				:  Dict		       =     use_labels
        lowerCAmelCase__				:  Dict		       =     vocab_size
        lowerCAmelCase__				:  Tuple		       =     hidden_size
        lowerCAmelCase__				:  List[Any]		       =     num_hidden_layers
        lowerCAmelCase__				:  Optional[Any]		       =     num_attention_heads
        lowerCAmelCase__				:  Dict		       =     intermediate_size

        lowerCAmelCase__				:  List[Any]		       =     hidden_dropout_prob
        lowerCAmelCase__				:  str		       =     attention_probs_dropout_prob
        lowerCAmelCase__				:  Optional[int]		       =     max_position_embeddings
        lowerCAmelCase__				:  str		       =     eos_token_id
        lowerCAmelCase__				:  List[Any]		       =     pad_token_id
        lowerCAmelCase__				:  Optional[int]		       =     bos_token_id

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  str		       =     ids_tensor([self.batch_size, self.seq_length - 1]      ,				self.vocab_size		)
        lowerCAmelCase__				:  Union[str, Any]		       =     tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)      ,				1		)
        lowerCAmelCase__				:  Tuple		       =     tf.concat([input_ids, eos_tensor]      ,				axis=1		)

        lowerCAmelCase__				:  Union[str, Any]		       =     ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size		)

        lowerCAmelCase__				:  List[Any]		       =     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      ,				)
        lowerCAmelCase__				:  List[Any]		       =     prepare_blenderbot_inputs_dict(__UpperCamelCase      ,				__UpperCamelCase      ,				__UpperCamelCase		)
        return config, inputs_dict







      def 	__magic_name__(   self      ,				__UpperCAmelCase      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  int		       =     TFBlenderbotModel(config=__UpperCamelCase		).get_decoder()
        lowerCAmelCase__				:  List[str]		       =     inputs_dict['''input_ids''']

        lowerCAmelCase__				:  Any		       =     input_ids[:1, :]
        lowerCAmelCase__				:  Dict		       =     inputs_dict['''attention_mask'''][:1, :]
        lowerCAmelCase__				:  int		       =     inputs_dict['''head_mask''']
        lowerCAmelCase__				:  List[str]		       =     1

        # first forward pass
        lowerCAmelCase__				:  List[str]		       =     model(__UpperCamelCase      ,				attention_mask=__UpperCamelCase      ,				head_mask=__UpperCamelCase      ,				use_cache=__UpperCamelCase		)

        lowerCAmelCase__       ,			lowerCAmelCase__				:  Any		       =     outputs.to_tuple()

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

        # append to next input_ids and
        lowerCAmelCase__				:  Tuple		       =     tf.concat([input_ids, next_tokens]      ,				axis=-1		)
        lowerCAmelCase__				:  List[Any]		       =     tf.concat([attention_mask, next_attn_mask]      ,				axis=-1		)

        lowerCAmelCase__				:  Tuple		       =     model(__UpperCamelCase      ,				attention_mask=__UpperCamelCase		)[0]
        lowerCAmelCase__				:  Tuple		       =     model(__UpperCamelCase      ,				attention_mask=__UpperCamelCase      ,				past_key_values=__UpperCamelCase		)[0]

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

        # select random slice
        lowerCAmelCase__				:  Dict		       =     int(ids_tensor((1,)      ,				output_from_past.shape[-1]		)		)
        lowerCAmelCase__				:  Union[str, Any]		       =     output_from_no_past[:, -3:, random_slice_idx]
        lowerCAmelCase__				:  List[Any]		       =     output_from_past[:, :, random_slice_idx]

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



def      __lowerCAmelCase			(					UpperCamelCase	,      UpperCamelCase	,      UpperCamelCase	,      UpperCamelCase=None	,      UpperCamelCase=None	,      UpperCamelCase=None	,      UpperCamelCase=None	,      UpperCamelCase=None	,      )     ->	Union[str, Any]:
  if attention_mask is None:
    lowerCAmelCase__				:  int		       =     tf.cast(tf.math.not_equal(UpperCAmelCase__	,      config.pad_token_id      )	,      tf.inta      )
  if decoder_attention_mask is None:
    lowerCAmelCase__				:  Optional[Any]		       =     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:
    lowerCAmelCase__				:  int		       =     tf.ones((config.encoder_layers, config.encoder_attention_heads)      )
  if decoder_head_mask is None:
    lowerCAmelCase__				:  str		       =     tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
  if cross_attn_head_mask is None:
    lowerCAmelCase__				:  List[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  _lowerCAmelCase     (  _lowercase    ,			_lowercase    ,			unittest.TestCase	):
      A__		     =		(TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
      A__		     =		(TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
      A__		     =		(
          {
              'conversational': TFBlenderbotForConditionalGeneration,
              'feature-extraction': TFBlenderbotModel,
              'summarization': TFBlenderbotForConditionalGeneration,
              'text2text-generation': TFBlenderbotForConditionalGeneration,
              'translation': TFBlenderbotForConditionalGeneration,
          }
          if is_tf_available()
          else {}
      )
      A__		     =		True
      A__		     =		False
      A__		     =		False

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     TFBlenderbotModelTester(self		)
        lowerCAmelCase__				:  Optional[Any]		       =     ConfigTester(self      ,				config_class=__UpperCamelCase		)

      def 	__magic_name__(   self		):
        self.config_tester.run_common_tests()







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Any		       =     self.model_tester.prepare_config_and_inputs_for_common()
        self.model_tester.check_decoder_model_past_large_inputs(*__UpperCamelCase		)






@require_tokenizers
@require_tf
class  _lowerCAmelCase     (  unittest.TestCase	):
      A__		     =		['My friends are cool but they eat too many carbs.']
      A__		     =		'facebook/blenderbot-400M-distill'

      @cached_property
      def 	__magic_name__(   self		):
        return BlenderbotTokenizer.from_pretrained(self.model_name		)

      @cached_property
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name		)
        return model







      @slow
      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     self.tokenizer(self.src_text      ,				return_tensors='''tf'''		)
        lowerCAmelCase__				:  str		       =     self.model.generate(
            model_inputs.input_ids      ,				)
        lowerCAmelCase__				:  List[str]		       =     self.tokenizer.batch_decode(generated_ids.numpy()      ,				skip_special_tokens=__UpperCamelCase		)[0]
        assert (
            generated_words
            == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?"
        )







from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch






@require_torch
@require_tf
class  _lowerCAmelCase     (  _lowercase	):

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  Union[str, Any]		       =     SMALL_MODEL_IDENTIFIER
        lowerCAmelCase__				:  List[str]		       =     '''pt'''
        lowerCAmelCase__				:  int		       =     '''tf'''

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  Optional[Any]		       =     AutoModel.from_pretrained(self.test_model		)
        model_pt.save_pretrained(__UpperCAmelCase		)

      def 	__magic_name__(   self      ,				__UpperCAmelCase		):
        lowerCAmelCase__				:  str		       =     TFAutoModel.from_pretrained(self.test_model      ,				from_pt=__UpperCAmelCase		)
        model_tf.save_pretrained(__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  int		       =     '''mock_framework'''

        # Framework provided - return whatever the user provides
        lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(self.test_model      ,				__UpperCAmelCase		)
        self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

        # Local checkpoint and framework provided - return provided framework
        # PyTorch checkpoint
        with TemporaryDirectory() as local_pt_ckpt:
          self._setup_pt_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(__UpperCAmelCase      ,				__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

        # TensorFlow checkpoint
        with TemporaryDirectory() as local_tf_ckpt:
          self._setup_tf_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  List[str]		       =     FeaturesManager.determine_framework(__UpperCAmelCase      ,				__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				__UpperCAmelCase		)

      def 	__magic_name__(   self		):
        # PyTorch checkpoint
        with TemporaryDirectory() as local_pt_ckpt:
          self._setup_pt_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # TensorFlow checkpoint
        with TemporaryDirectory() as local_tf_ckpt:
          self._setup_tf_ckpt(__UpperCAmelCase		)
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_tf		)

        # Invalid local checkpoint
        with TemporaryDirectory() as local_invalid_ckpt:
          with self.assertRaises(__UpperCAmelCase		):
            lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(__UpperCAmelCase		)







      def 	__magic_name__(   self		):
        lowerCAmelCase__				:  List[Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  int		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # PyTorch not in environment -> use TensorFlow
        lowerCAmelCase__				:  Union[str, Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  Optional[int]		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_tf		)

        # Both in environment -> use PyTorch
        lowerCAmelCase__				:  Union[str, Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        lowerCAmelCase__				:  List[Any]		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		), patch(
            '''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          lowerCAmelCase__				:  List[Any]		       =     FeaturesManager.determine_framework(self.test_model		)
          self.assertEqual(__UpperCAmelCase      ,				self.framework_pt		)

        # Both not in environment -> raise error
        lowerCAmelCase__				:  Optional[int]		       =     MagicMock(return_value=__UpperCAmelCase		)
        lowerCAmelCase__				:  Dict		       =     MagicMock(return_value=__UpperCAmelCase		)
        with patch('''transformers.onnx.features.is_tf_available'''      ,				__UpperCAmelCase		), patch(
            '''transformers.onnx.features.is_torch_available'''      ,				__UpperCAmelCase		):
          with self.assertRaises(__UpperCAmelCase		):
            lowerCAmelCase__				:  Tuple		       =     FeaturesManager.determine_framework(self.test_model		)

