



"""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       UpperCAmelCase_    :
      __SCREAMING_SNAKE_CASE		:  Optional[int]						=					BlenderbotSmallConfig
      __SCREAMING_SNAKE_CASE		:  Optional[int]						=					{}
      __SCREAMING_SNAKE_CASE		:  Union[str, Any]						=					'gelu'



      def __init__(	self	:      str	,      A	:      List[Any]	,      A	:      Any=1_3	,      A	:      Optional[int]=7	,      A	:      List[str]=True	,      A	:      str=False	,      A	:      Optional[Any]=9_9	,      A	:      Union[str, Any]=3_2	,      A	:      List[str]=2	,      A	:      Optional[int]=4	,      A	:      Tuple=3_7	,      A	:      List[Any]=0.1	,      A	:      Dict=0.1	,      A	:      Any=2_0	,      A	:      str=2	,      A	:      Optional[Any]=1	,      A	:      Union[str, Any]=0	,      ):
           _UpperCAmelCase     :      Optional[int]														=       parent
           _UpperCAmelCase     :      Dict														=       batch_size
           _UpperCAmelCase     :      Dict														=       seq_length
           _UpperCAmelCase     :      Tuple														=       is_training
           _UpperCAmelCase     :      Optional[int]														=       use_labels
           _UpperCAmelCase     :      Union[str, Any]														=       vocab_size
           _UpperCAmelCase     :      Dict														=       hidden_size
           _UpperCAmelCase     :      Dict														=       num_hidden_layers
           _UpperCAmelCase     :      str														=       num_attention_heads
           _UpperCAmelCase     :      List[Any]														=       intermediate_size

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



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

           _UpperCAmelCase     :      Dict														=       ids_tensor([self.batch_size, self.seq_length]	,      self.vocab_size		)

           _UpperCAmelCase     :      Optional[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	,      )
           _UpperCAmelCase     :      int														=       prepare_blenderbot_small_inputs_dict(A	,      A	,      A		)
           return config, inputs_dict



      def 		snake_case_   (	self	:      Tuple	,      A	:      Optional[Any]	,      A	:      Union[str, Any]		):
           _UpperCAmelCase     :      str														=       TFBlenderbotSmallModel(config=A		).get_decoder()
           _UpperCAmelCase     :      Union[str, Any]														=       inputs_dict["input_ids"]

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

           # first forward pass
           _UpperCAmelCase     :      List[Any]														=       model(A	,      attention_mask=A	,      head_mask=A	,      use_cache=A		)

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

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

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

           _UpperCAmelCase     :      List[Any]														=       model(A	,      attention_mask=A		)[0]
           _UpperCAmelCase     :      Any														=       model(A	,      attention_mask=A	,      past_key_values=A		)[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[Any]														=       output_from_no_past[:, -3:, random_slice_idx]
           _UpperCAmelCase     :      List[str]														=       output_from_past[:, :, random_slice_idx]

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







def 		__snake_case		( SCREAMING_SNAKE_CASE__      :		str		,    SCREAMING_SNAKE_CASE__      :		Tuple		,    SCREAMING_SNAKE_CASE__      :		Dict		,    SCREAMING_SNAKE_CASE__      :		Optional[int]=None		,    SCREAMING_SNAKE_CASE__      :		Optional[Any]=None		,    SCREAMING_SNAKE_CASE__      :		List[str]=None		,    SCREAMING_SNAKE_CASE__      :		List[Any]=None		,    SCREAMING_SNAKE_CASE__      :		Optional[int]=None		,    )    ->  Optional[int]:





     '''simple docstring'''

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



      def 		snake_case_   (	self	:      int		):
           _UpperCAmelCase     :      Optional[int]														=       TFBlenderbotSmallModelTester(self		)
           _UpperCAmelCase     :      Optional[int]														=       ConfigTester(self	,      config_class=A		)



      def 		snake_case_   (	self	:      Dict		):
           self.config_tester.run_common_tests()



      def 		snake_case_   (	self	:      Union[str, Any]		):
           _UpperCAmelCase     :      Union[str, Any]														=       self.model_tester.prepare_config_and_inputs_for_common()
           self.model_tester.check_decoder_model_past_large_inputs(*A		)







@require_tokenizers
@require_tf
class       UpperCAmelCase_    (							unittest.TestCase		):
      __SCREAMING_SNAKE_CASE		:  Dict						=					[
          '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		:  Optional[int]						=					'facebook/blenderbot_small-90M'



      @cached_property
      def 		snake_case_   (	self	:      List[str]		):
           # use "old" tokenizer here because of bug when downloading new tokenizer
           return BlenderbotSmallTokenizer.from_pretrained("facebook/blenderbot-90M"		)



      @cached_property
      def 		snake_case_   (	self	:      Dict		):
           _UpperCAmelCase     :      Union[str, Any]														=       TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name		)
           return model



      @slow
      def 		snake_case_   (	self	:      Union[str, Any]		):
           _UpperCAmelCase     :      Tuple														=       self.tokenizer(self.src_text	,      return_tensors="tf"		)
           _UpperCAmelCase     :      int														=       self.model.generate(
               model_inputs.input_ids	,      attention_mask=model_inputs.attention_mask	,      num_beams=2	,      use_cache=A	,      )
           _UpperCAmelCase     :      Tuple														=       self.tokenizer.batch_decode(generated_ids.numpy()	,      skip_special_tokens=A		)[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.",
           )





"""simple docstring"""




from manim import *






class       UpperCAmelCase_    (							_UpperCamelCase		):



      def 		snake_case_   (	self	:      Tuple		):
           _UpperCAmelCase     :      str														=       Rectangle(height=0.5	,      width=0.5		)
           _UpperCAmelCase     :      List[Any]														=       Rectangle(height=0.25	,      width=0.25		)
           _UpperCAmelCase     :      Tuple														=       Rectangle(height=0.46	,      width=0.46		).set_stroke(width=0		)

           _UpperCAmelCase     :      List[Any]														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      List[str]														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      int														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Dict														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      int														=       VGroup(A	,      A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      str														=       Text("CPU"	,      font_size=2_4		)
           _UpperCAmelCase     :      Union[str, Any]														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           cpu.move_to([-2.5, -0.5, 0]		)
           self.add(A		)

           _UpperCAmelCase     :      Union[str, Any]														=       [mem.copy() for i in range(4		)]
           _UpperCAmelCase     :      Optional[int]														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Any														=       Text("GPU"	,      font_size=2_4		)
           _UpperCAmelCase     :      Dict														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           gpu.move_to([-1, -1, 0]		)
           self.add(A		)

           _UpperCAmelCase     :      Tuple														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      str														=       VGroup(*A		).arrange(A	,      buff=0		)

           _UpperCAmelCase     :      Any														=       Text("Model"	,      font_size=2_4		)
           _UpperCAmelCase     :      Any														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           model.move_to([3, -1.0, 0]		)
           self.add(A		)

           _UpperCAmelCase     :      List[str]														=       []
           _UpperCAmelCase     :      Optional[int]														=       []
           _UpperCAmelCase     :      str														=       []

           for i, rect in enumerate(A		):
                rect.set_stroke(A		)

                _UpperCAmelCase     :      Dict														=       Rectangle(height=0.46 / 4	,      width=0.46 / 3		).set_stroke(width=0.0		).set_fill(A	,      opacity=0.7		)

                if i == 0:
                     cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT		)	,      buff=0.02	,      direction=A		)
                     cpu_target.set_x(cpu_target.get_x() + 0.1		)
                elif i == 3:
                     cpu_target.next_to(model_cpu_arr[0]	,      direction=A	,      buff=0.0		)
                else:
                     cpu_target.next_to(model_cpu_arr[i - 1]	,      direction=A	,      buff=0.0		)
                self.add(A		)
                model_cpu_arr.append(A		)

           self.add(*A	,      *A	,      *A		)

           _UpperCAmelCase     :      Dict														=       [mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      List[str]														=       VGroup(*A		).arrange(A	,      buff=0		)

           _UpperCAmelCase     :      List[str]														=       Text("Loaded Checkpoint"	,      font_size=2_4		)
           _UpperCAmelCase     :      Optional[int]														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           checkpoint.move_to([3, 0.5, 0]		)

           self.add(A		)

           _UpperCAmelCase     :      Union[str, Any]														=       []
           _UpperCAmelCase     :      Union[str, Any]														=       []

           for i, rect in enumerate(A		):
                _UpperCAmelCase     :      str														=       fill.copy().set_fill(A	,      opacity=0.7		)
                target.move_to(A		)
                ckpt_arr.append(A		)

                _UpperCAmelCase     :      Dict														=       target.copy()
                if i < 5:
                     cpu_target.move_to(cpu_left_col_base[i + 1]		)
                else:
                     cpu_target.move_to(cpu_right_col_base[i - 5]		)
                ckpt_cpu_arr.append(A		)
           self.add(*A	,      *A		)

           _UpperCAmelCase     :      List[str]														=       Square(side_length=2.2		)
           key.move_to([-5, 2, 0]		)

           _UpperCAmelCase     :      List[Any]														=       MarkupText(
               f'<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'	,      font_size=1_8	,      )

           key_text.move_to([-5, 2.4, 0]		)

           self.add(A	,      A		)

           _UpperCAmelCase     :      Tuple														=       MarkupText(
               f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'	,      font_size=1_8	,      )

           blue_text.next_to(A	,      DOWN * 2.4	,      aligned_edge=key_text.get_left()		)
           self.add(A		)

           _UpperCAmelCase     :      Any														=       MarkupText(
               f'Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.'	,      font_size=2_4	,      )
           step_a.move_to([2, 2, 0]		)

           _UpperCAmelCase     :      int														=       [meta_mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      Optional[int]														=       [meta_mem.copy() for i in range(6		)]
           _UpperCAmelCase     :      List[str]														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Optional[int]														=       VGroup(*A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      List[Any]														=       VGroup(A	,      A		).arrange(A	,      buff=0		)
           _UpperCAmelCase     :      Union[str, Any]														=       Text("Disk"	,      font_size=2_4		)
           _UpperCAmelCase     :      int														=       Group(A	,      A		).arrange(A	,      buff=0.5	,      aligned_edge=A		)
           disk.move_to([-4.0, -1.25, 0]		)
           self.play(Write(A	,      run_time=3		)	,      Write(A	,      run_time=1		)	,      Create(A	,      run_time=1		)		)

           _UpperCAmelCase     :      List[Any]														=       []
           for i, rect in enumerate(A		):
                _UpperCAmelCase     :      Any														=       rect.copy()
                target.generate_target()
                target.target.move_to(disk_left_col_base[i]		).scale(0.5		)
                animations.append(MoveToTarget(A	,      run_time=1.5		)		)
           self.play(*A		)

           self.play(FadeOut(A		)		)

           _UpperCAmelCase     :      int														=       MarkupText(f'Then, the checkpoint is removed from memory\nthrough garbage collection.'	,      font_size=2_4		)
           step_a.move_to([2, 2, 0]		)

           self.play(Write(A	,      run_time=3		)		)

           self.play(
               FadeOut(A	,      A	,      *A	,      *A		)	,      )

           self.wait()
