'''simple docstring'''
import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import BertGenerationDecoder, BertGenerationEncoder







class lowerCamelCase							:




  '''simple docstring'''



  def __init__(      self			:       Any  ,							UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       Optional[int]=13  ,							UpperCAmelCase__			:       int=7  ,							UpperCAmelCase__			:       Dict=True  ,							UpperCAmelCase__			:       Tuple=True  ,							UpperCAmelCase__			:       Any=99  ,							UpperCAmelCase__			:       Optional[int]=32  ,							UpperCAmelCase__			:       Union[str, Any]=5  ,							UpperCAmelCase__			:       List[str]=4  ,							UpperCAmelCase__			:       Union[str, Any]=37  ,							UpperCAmelCase__			:       List[Any]="gelu"  ,							UpperCAmelCase__			:       Union[str, Any]=0.1  ,							UpperCAmelCase__			:       Any=0.1  ,							UpperCAmelCase__			:       List[str]=50  ,							UpperCAmelCase__			:       List[Any]=0.02  ,							UpperCAmelCase__			:       List[Any]=True  ,							UpperCAmelCase__			:       str=None  ,							)				->List[str]:
        UpperCAmelCase_            =  parent
        UpperCAmelCase_            =  batch_size
        UpperCAmelCase_            =  seq_length
        UpperCAmelCase_            =  is_training
        UpperCAmelCase_            =  use_input_mask
        UpperCAmelCase_            =  vocab_size
        UpperCAmelCase_            =  hidden_size
        UpperCAmelCase_            =  num_hidden_layers
        UpperCAmelCase_            =  num_attention_heads
        UpperCAmelCase_            =  intermediate_size
        UpperCAmelCase_            =  hidden_act
        UpperCAmelCase_            =  hidden_dropout_prob
        UpperCAmelCase_            =  attention_probs_dropout_prob
        UpperCAmelCase_            =  max_position_embeddings
        UpperCAmelCase_            =  initializer_range
        UpperCAmelCase_            =  use_labels
        UpperCAmelCase_            =  scope






  def  lowerCAmelCase__ (      self			:       Dict  )				->str:
        UpperCAmelCase_            =  ids_tensor([self.batch_size, self.seq_length]  ,							self.vocab_size  )

        UpperCAmelCase_            =  None
        if self.use_input_mask:
              UpperCAmelCase_            =  random_attention_mask([self.batch_size, self.seq_length]  )

        if self.use_labels:
              UpperCAmelCase_            =  ids_tensor([self.batch_size, self.seq_length]  ,							self.vocab_size  )

        UpperCAmelCase_            =  self.get_config()

        return config, input_ids, input_mask, token_labels






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Any:
        return BertGenerationConfig(
            vocab_size=self.vocab_size  ,							hidden_size=self.hidden_size  ,							num_hidden_layers=self.num_hidden_layers  ,							num_attention_heads=self.num_attention_heads  ,							intermediate_size=self.intermediate_size  ,							hidden_act=self.hidden_act  ,							hidden_dropout_prob=self.hidden_dropout_prob  ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,							max_position_embeddings=self.max_position_embeddings  ,							is_decoder=A__  ,							initializer_range=self.initializer_range  ,							)






  def  lowerCAmelCase__ (      self			:       int  )				->Tuple:
        (
            UpperCAmelCase_     
        )            =  self.prepare_config_and_inputs()

        UpperCAmelCase_            =  True
        UpperCAmelCase_            =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
        UpperCAmelCase_            =  ids_tensor([self.batch_size, self.seq_length]  ,							vocab_size=2  )

        return (
            config,
            input_ids,
            input_mask,
            token_labels,
            encoder_hidden_states,
            encoder_attention_mask,
        )






  def  lowerCAmelCase__ (      self			:       Dict  ,							UpperCAmelCase__			:       Any  ,							UpperCAmelCase__			:       Optional[int]  ,							UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       List[str]  ,							**UpperCAmelCase__			:       Dict  ,							)				->List[Any]:
        UpperCAmelCase_            =  BertGenerationEncoder(config=A__  )
        model.to(A__  )
        model.eval()
        UpperCAmelCase_            =  model(A__  ,							attention_mask=A__  )
        UpperCAmelCase_            =  model(A__  )
        self.parent.assertEqual(result.last_hidden_state.shape  ,							(self.batch_size, self.seq_length, self.hidden_size)  )






  def  lowerCAmelCase__ (      self			:       int  ,							UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       Optional[int]  ,							UpperCAmelCase__			:       Optional[int]  ,							UpperCAmelCase__			:       Optional[int]  ,							UpperCAmelCase__			:       List[str]  ,							UpperCAmelCase__			:       int  ,							**UpperCAmelCase__			:       List[Any]  ,							)				->Optional[int]:
        UpperCAmelCase_            =  True
        UpperCAmelCase_            =  BertGenerationEncoder(config=A__  )
        model.to(A__  )
        model.eval()
        UpperCAmelCase_            =  model(
            A__  ,							attention_mask=A__  ,							encoder_hidden_states=A__  ,							encoder_attention_mask=A__  ,							)
        UpperCAmelCase_            =  model(
            A__  ,							attention_mask=A__  ,							encoder_hidden_states=A__  ,							)
        self.parent.assertEqual(result.last_hidden_state.shape  ,							(self.batch_size, self.seq_length, self.hidden_size)  )






  def  lowerCAmelCase__ (      self			:       Any  ,							UpperCAmelCase__			:       List[str]  ,							UpperCAmelCase__			:       List[str]  ,							UpperCAmelCase__			:       str  ,							UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       Dict  ,							UpperCAmelCase__			:       int  ,							**UpperCAmelCase__			:       int  ,							)				->List[Any]:
        UpperCAmelCase_            =  True
        UpperCAmelCase_            =  True
        UpperCAmelCase_            =  BertGenerationDecoder(config=A__  ).to(A__  ).eval()

        # first forward pass
        UpperCAmelCase_            =  model(
            A__  ,							attention_mask=A__  ,							encoder_hidden_states=A__  ,							encoder_attention_mask=A__  ,							use_cache=A__  ,							)
        UpperCAmelCase_            =  outputs.past_key_values

        # create hypothetical multiple next token and extent to next_input_ids
        UpperCAmelCase_            =  ids_tensor((self.batch_size, 3)  ,							config.vocab_size  )
        UpperCAmelCase_            =  ids_tensor((self.batch_size, 3)  ,							vocab_size=2  )

        # append to next input_ids and
        UpperCAmelCase_            =  torch.cat([input_ids, next_tokens]  ,							dim=-1  )
        UpperCAmelCase_            =  torch.cat([input_mask, next_mask]  ,							dim=-1  )

        UpperCAmelCase_            =  model(
            A__  ,							attention_mask=A__  ,							encoder_hidden_states=A__  ,							encoder_attention_mask=A__  ,							output_hidden_states=A__  ,							)["""hidden_states"""][0]
        UpperCAmelCase_            =  model(
            A__  ,							attention_mask=A__  ,							encoder_hidden_states=A__  ,							encoder_attention_mask=A__  ,							past_key_values=A__  ,							output_hidden_states=A__  ,							)["""hidden_states"""][0]

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

        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

        # test that outputs are equal for slice
        self.parent.assertTrue(torch.allclose(A__  ,							A__  ,							atol=1e-3  )  )






  def  lowerCAmelCase__ (      self			:       Tuple  ,							UpperCAmelCase__			:       Tuple  ,							UpperCAmelCase__			:       Any  ,							UpperCAmelCase__			:       List[str]  ,							UpperCAmelCase__			:       Union[str, Any]  ,							*UpperCAmelCase__			:       Tuple  ,							)				->List[str]:
        UpperCAmelCase_            =  BertGenerationDecoder(A__  )
        model.to(A__  )
        model.eval()
        UpperCAmelCase_            =  model(A__  ,							attention_mask=A__  ,							labels=A__  )
        self.parent.assertEqual(result.logits.shape  ,							(self.batch_size, self.seq_length, self.vocab_size)  )






  def  lowerCAmelCase__ (      self			:       List[str]  )				->str:
        UpperCAmelCase_            =  self.prepare_config_and_inputs()
        UpperCAmelCase_            =  {"""input_ids""": input_ids, """attention_mask""": input_mask}
        return config, inputs_dict







@require_torch
class lowerCamelCase							(	lowerCamelCase     ,		lowerCamelCase     ,		lowerCamelCase     ,		unittest.TestCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	(BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
  lowerCAmelCase__										=	(BertGenerationDecoder,) if is_torch_available() else ()
  lowerCAmelCase__										=	(
      {'''feature-extraction''': BertGenerationEncoder, '''text-generation''': BertGenerationDecoder}
      if is_torch_available()
      else {}
  )






  def  lowerCAmelCase__ (      self			:       Any  )				->Optional[Any]:
        UpperCAmelCase_            =  BertGenerationEncoderTester(self  )
        UpperCAmelCase_            =  ConfigTester(self  ,							config_class=A__  ,							hidden_size=37  )






  def  lowerCAmelCase__ (      self			:       Tuple  )				->Any:
        self.config_tester.run_common_tests()






  def  lowerCAmelCase__ (      self			:       Union[str, Any]  )				->Any:
        UpperCAmelCase_            =  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*A__  )






  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->List[str]:
        UpperCAmelCase_            =  self.model_tester.prepare_config_and_inputs()
        UpperCAmelCase_            =  """bert"""
        self.model_tester.create_and_check_model(A__  ,							A__  ,							A__  ,							A__  )






  def  lowerCAmelCase__ (      self			:       Optional[Any]  )				->Tuple:
        UpperCAmelCase_            =  self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*A__  )






  def  lowerCAmelCase__ (      self			:       Dict  )				->List[Any]:
        UpperCAmelCase_            =  self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_decoder_model_past_large_inputs(*A__  )






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Any:
        # This regression test was failing with PyTorch < 1.3
        (
            UpperCAmelCase_     
        )            =  self.model_tester.prepare_config_and_inputs_for_decoder()

        UpperCAmelCase_            =  None

        self.model_tester.create_and_check_model_as_decoder(
            A__  ,							A__  ,							A__  ,							A__  ,							A__  ,							A__  ,							)






  def  lowerCAmelCase__ (      self			:       Tuple  )				->List[str]:
        UpperCAmelCase_            =  self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_for_causal_lm(*A__  )






  @slow
  def  lowerCAmelCase__ (      self			:       str  )				->Tuple:
        UpperCAmelCase_            =  BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
        self.assertIsNotNone(A__  )







@require_torch
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  @slow
  def  lowerCAmelCase__ (      self			:       Optional[int]  )				->Optional[Any]:
        UpperCAmelCase_            =  BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
        UpperCAmelCase_            =  torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]  )
        with torch.no_grad():
              UpperCAmelCase_            =  model(A__  )[0]
        UpperCAmelCase_            =  torch.Size([1, 8, 1024]  )
        self.assertEqual(output.shape  ,							A__  )
        UpperCAmelCase_            =  torch.tensor(
            [[[0.1775, 0.0083, -0.0321], [1.6002, 0.1287, 0.3912], [2.1473, 0.5791, 0.6066]]]  )
        self.assertTrue(torch.allclose(output[:, :3, :3]  ,							A__  ,							atol=1e-4  )  )







@require_torch
class lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



  @slow
  def  lowerCAmelCase__ (      self			:       Any  )				->Dict:
        UpperCAmelCase_            =  BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
        UpperCAmelCase_            =  torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]  )
        with torch.no_grad():
              UpperCAmelCase_            =  model(A__  )[0]
        UpperCAmelCase_            =  torch.Size([1, 8, 5_0358]  )
        self.assertEqual(output.shape  ,							A__  )
        UpperCAmelCase_            =  torch.tensor(
            [[[-0.5788, -2.5994, -3.7054], [0.0438, 4.7997, 1.8795], [1.5862, 6.6409, 4.4638]]]  )
        self.assertTrue(torch.allclose(output[:, :3, :3]  ,							A__  ,							atol=1e-4  )  )






'''simple docstring'''
from datetime import datetime

import matplotlib.pyplot as plt
import torch



def  __lowerCamelCase       (      _UpperCamelCase  :  Union[str, Any]					):


      '''simple docstring'''


      for param in module.parameters():
            UpperCAmelCase_            =  False



def  __lowerCamelCase       (      ):


      '''simple docstring'''


      UpperCAmelCase_            =  '''cuda''' if torch.cuda.is_available() else '''cpu'''
      if torch.backends.mps.is_available() and torch.backends.mps.is_built():
            UpperCAmelCase_            =  '''mps'''
      if device == "mps":
            print(
                '''WARNING: MPS currently doesn\'t seem to work, and messes up backpropagation without any visible torch'''
                ''' errors. I recommend using CUDA on a colab notebook or CPU instead if you\'re facing inexplicable issues'''
                ''' with generations.'''					)
      return device



def  __lowerCamelCase       (      _UpperCamelCase  :  Union[str, Any]					):


      '''simple docstring'''


      UpperCAmelCase_            =  plt.imshow(_UpperCamelCase					)
      fig.axes.get_xaxis().set_visible(_UpperCamelCase					)
      fig.axes.get_yaxis().set_visible(_UpperCamelCase					)
      plt.show()





def  __lowerCamelCase       (      ):


      '''simple docstring'''


      UpperCAmelCase_            =  datetime.now()
      UpperCAmelCase_            =  current_time.strftime('''%H:%M:%S'''					)
      return timestamp





