






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 							__SCREAMING_SNAKE_CASE       :
      __a  					=BlenderbotConfig
      __a  					={}
      __a  					="""gelu"""




      def __init__(  self ,   lowerCamelCase ,   lowerCamelCase=13 ,   lowerCamelCase=7 ,   lowerCamelCase=True ,   lowerCamelCase=False ,   lowerCamelCase=99 ,   lowerCamelCase=32 ,   lowerCamelCase=2 ,   lowerCamelCase=4 ,   lowerCamelCase=37 ,   lowerCamelCase=0.1 ,   lowerCamelCase=0.1 ,   lowerCamelCase=20 ,   lowerCamelCase=2 ,   lowerCamelCase=1 ,   lowerCamelCase=0 ,   )							->Dict:



         '''simple docstring'''
         __a				     =    parent
         __a				     =    batch_size
         __a				     =    seq_length
         __a				     =    is_training
         __a				     =    use_labels
         __a				     =    vocab_size
         __a				     =    hidden_size
         __a				     =    num_hidden_layers
         __a				     =    num_attention_heads
         __a				     =    intermediate_size

         __a				     =    hidden_dropout_prob
         __a				     =    attention_probs_dropout_prob
         __a				     =    max_position_embeddings
         __a				     =    eos_token_id
         __a				     =    pad_token_id
         __a				     =    bos_token_id




      def  __UpperCamelCase							(  self      )							->Any:



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

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

         __a				     =    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 ,   )
         __a				     =    prepare_blenderbot_inputs_dict(lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__      )
         return config, inputs_dict






      def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase      )							->Tuple:



         '''simple docstring'''
         __a				     =    TFBlenderbotModel(config=lowerCAmelCase__      ).get_decoder()
         __a				     =    inputs_dict["input_ids"]

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

         # first forward pass
         __a				     =    model(lowerCAmelCase__ ,   attention_mask=lowerCAmelCase__ ,   head_mask=lowerCAmelCase__ ,   use_cache=lowerCAmelCase__      )

         __a				     =    outputs.to_tuple()

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

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

         __a				     =    model(lowerCAmelCase__ ,   attention_mask=lowerCAmelCase__      )[0]
         __a				     =    model(lowerCAmelCase__ ,   attention_mask=lowerCAmelCase__ ,   past_key_values=lowerCAmelCase__      )[0]

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

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

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







def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Dict,					SCREAMING_SNAKE_CASE__:	Any,					SCREAMING_SNAKE_CASE__:	Optional[Any],					SCREAMING_SNAKE_CASE__:	Any=None,					SCREAMING_SNAKE_CASE__:	Union[str, Any]=None,					SCREAMING_SNAKE_CASE__:	Union[str, Any]=None,					SCREAMING_SNAKE_CASE__:	Tuple=None,					SCREAMING_SNAKE_CASE__:	Dict=None,					)  ->			Dict:

   """simple docstring"""






   if attention_mask is None:
      __a				     =    tf.cast(tf.math.not_equal(_lowercase,					config.pad_token_id    ),					tf.inta    )
   if decoder_attention_mask is None:
      __a				     =    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:
      __a				     =    tf.ones((config.encoder_layers, config.encoder_attention_heads)    )
   if decoder_head_mask is None:
      __a				     =    tf.ones((config.decoder_layers, config.decoder_attention_heads)    )
   if cross_attn_head_mask is None:
      __a				     =    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 							__SCREAMING_SNAKE_CASE       (							a__      ,      a__      ,      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  __UpperCamelCase							(  self      )							->Optional[Any]:



         '''simple docstring'''
         __a				     =    TFBlenderbotModelTester(self      )
         __a				     =    ConfigTester(self ,   config_class=lowerCAmelCase__      )




      def  __UpperCamelCase							(  self      )							->List[str]:



         '''simple docstring'''
         self.config_tester.run_common_tests()






      def  __UpperCamelCase							(  self      )							->int:



         '''simple docstring'''
         __a				     =    self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.check_decoder_model_past_large_inputs(*lowerCAmelCase__      )







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




      @cached_property
      def  __UpperCamelCase							(  self      )							->Any:



         '''simple docstring'''
         return BlenderbotTokenizer.from_pretrained(self.model_name      )




      @cached_property
      def  __UpperCamelCase							(  self      )							->List[Any]:



         '''simple docstring'''
         __a				     =    TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name      )
         return model






      @slow
      def  __UpperCamelCase							(  self      )							->Union[str, Any]:



         '''simple docstring'''
         __a				     =    self.tokenizer(self.src_text ,   return_tensors='tf'      )
         __a				     =    self.model.generate(
             model_inputs.input_ids ,   )
         __a				     =    self.tokenizer.batch_decode(generated_ids.numpy() ,   skip_special_tokens=lowerCAmelCase__      )[0]
         assert (
             generated_words
             == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?"
         )







'''simple docstring'''



from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_torch_available, is_torch_tpu_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_torch_available():
 import torch

if is_torch_tpu_available(check_device=False):
 import torch_xla.core.xla_model as xm


__UpperCamelCase      :       Optional[Any]      			=      logging.get_logger(__name__)
@dataclass
class 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):
      __a  					=[
          "no_inference",
          "no_cuda",
          "no_tpu",
          "no_speed",
          "no_memory",
          "no_env_print",
          "no_multi_process",
      ]




      def __init__(  self ,   **lowerCamelCase      )							->List[Any]:



         '''simple docstring'''
         for deprecated_arg in self.deprecated_args:
            if deprecated_arg in kwargs:
               __a				     =    deprecated_arg[3:]
               setattr(self ,   lowerCamelCase ,   not kwargs.pop(lowerCamelCase      )      )
               logger.warning(
                   F"""{deprecated_arg} is depreciated. Please use --no_{positive_arg} or"""
                   F""" {positive_arg}={kwargs[positive_arg]}"""      )

         __a				     =    kwargs.pop('torchscript' ,   self.torchscript      )
         __a				     =    kwargs.pop('torch_xla_tpu_print_metrics' ,   self.torch_xla_tpu_print_metrics      )
         __a				     =    kwargs.pop('fp16_opt_level' ,   self.fpaa_opt_level      )
         super().__init__(**lowerCamelCase      )

      __a  					=field(default=_lowerCAmelCase      ,      metadata={"help": "Trace the models using torchscript"}     )
      __a  					=field(default=_lowerCAmelCase      ,      metadata={"help": "Print Xla/PyTorch tpu metrics"}     )
      __a  					=field(
          default="O1"      ,      metadata={
              "help": (
                  "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. "
                  "See details at https://nvidia.github.io/apex/amp.html"
              )
          }      ,      )




      @cached_property
      def  __UpperCamelCase							(  self      )							->Tuple["torch.device", int]:



         '''simple docstring'''
         requires_backends(self ,   ['torch']      )
         logger.info('PyTorch: setting up devices'      )
         if not self.cuda:
            __a				     =    torch.device('cpu'      )
            __a				     =    0
         elif is_torch_tpu_available():
            __a				     =    xm.xla_device()
            __a				     =    0
         else:
            __a				     =    torch.device('cuda' if torch.cuda.is_available() else 'cpu'      )
            __a				     =    torch.cuda.device_count()
         return device, n_gpu




      @property
      def  __UpperCamelCase							(  self      )							->Optional[Any]:



         '''simple docstring'''
         return is_torch_tpu_available() and self.tpu




      @property
      def  __UpperCamelCase							(  self      )							->int:



         '''simple docstring'''
         requires_backends(self ,   ['torch']      )
         # TODO(PVP): currently only single GPU is supported
         return torch.cuda.current_device()




      @property
      def  __UpperCamelCase							(  self      )							->"torch.device":



         '''simple docstring'''
         requires_backends(self ,   ['torch']      )
         return self._setup_devices[0]




      @property
      def  __UpperCamelCase							(  self      )							->List[Any]:



         '''simple docstring'''
         requires_backends(self ,   ['torch']      )
         return self._setup_devices[1]






      @property
      def  __UpperCamelCase							(  self      )							->List[str]:



         '''simple docstring'''
         return self.n_gpu > 0