


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 			UpperCamelCase__			:






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






      def 						__lowercase(  self       : Dict )	->			Optional[int]:
           UpperCamelCase__				:	Tuple										=  ids_tensor([self.batch_size, self.seq_length],		self.vocab_size )

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

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

           UpperCamelCase__				:	Union[str, Any]										=  self.get_config()

           return config, input_ids, input_mask, token_labels






      def 						__lowercase(  self       : List[str] )	->			Tuple:
           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 						__lowercase(  self       : int )	->			int:
           (
               UpperCamelCase__       
           )				:	str										=  self.prepare_config_and_inputs()

           UpperCamelCase__				:	int										=  True
           UpperCamelCase__				:	str										=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
           UpperCamelCase__				:	Union[str, Any]										=  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 						__lowercase(  self       : List[str],		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : List[str],		__lowerCamelCase       : str,		__lowerCamelCase       : int,		**__lowerCamelCase       : int,		)	->			Any:
           UpperCamelCase__				:	Tuple										=  BertGenerationEncoder(config=A_ )
           model.to(A_ )
           model.eval()
           UpperCamelCase__				:	Dict										=  model(A_,		attention_mask=A_ )
           UpperCamelCase__				:	Optional[Any]										=  model(A_ )
           self.parent.assertEqual(result.last_hidden_state.shape,		(self.batch_size, self.seq_length, self.hidden_size) )






      def 						__lowercase(  self       : Any,		__lowerCamelCase       : Dict,		__lowerCamelCase       : Dict,		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : Dict,		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : Tuple,		**__lowerCamelCase       : Optional[int],		)	->			str:
           UpperCamelCase__				:	Any										=  True
           UpperCamelCase__				:	List[str]										=  BertGenerationEncoder(config=A_ )
           model.to(A_ )
           model.eval()
           UpperCamelCase__				:	List[str]										=  model(
               A_,		attention_mask=A_,		encoder_hidden_states=A_,		encoder_attention_mask=A_,		)
           UpperCamelCase__				:	int										=  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 						__lowercase(  self       : Optional[Any],		__lowerCamelCase       : Optional[Any],		__lowerCamelCase       : Dict,		__lowerCamelCase       : Optional[int],		__lowerCamelCase       : List[Any],		__lowerCamelCase       : List[Any],		__lowerCamelCase       : Optional[int],		**__lowerCamelCase       : str,		)	->			Dict:
           UpperCamelCase__				:	Any										=  True
           UpperCamelCase__				:	Union[str, Any]										=  True
           UpperCamelCase__				:	int										=  BertGenerationDecoder(config=A_ ).to(A_ ).eval()

           # first forward pass
           UpperCamelCase__				:	Optional[Any]										=  model(
               A_,		attention_mask=A_,		encoder_hidden_states=A_,		encoder_attention_mask=A_,		use_cache=A_,		)
           UpperCamelCase__				:	Any										=  outputs.past_key_values

           # create hypothetical multiple next token and extent to next_input_ids
           UpperCamelCase__				:	Any										=  ids_tensor((self.batch_size, 3),		config.vocab_size )
           UpperCamelCase__				:	Optional[Any]										=  ids_tensor((self.batch_size, 3),		vocab_size=2 )

           # append to next input_ids and
           UpperCamelCase__				:	int										=  torch.cat([input_ids, next_tokens],		dim=-1 )
           UpperCamelCase__				:	Union[str, Any]										=  torch.cat([input_mask, next_mask],		dim=-1 )

           UpperCamelCase__				:	Dict										=  model(
               A_,		attention_mask=A_,		encoder_hidden_states=A_,		encoder_attention_mask=A_,		output_hidden_states=A_,		)["hidden_states"][0]
           UpperCamelCase__				:	Dict										=  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__				:	Any										=  ids_tensor((1,),		output_from_past.shape[-1] ).item()
           UpperCamelCase__				:	int										=  output_from_no_past[:, -3:, random_slice_idx].detach()
           UpperCamelCase__				:	Any										=  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 						__lowercase(  self       : Any,		__lowerCamelCase       : Any,		__lowerCamelCase       : Optional[Any],		__lowerCamelCase       : Union[str, Any],		__lowerCamelCase       : Optional[int],		*__lowerCamelCase       : Union[str, Any],		)	->			List[Any]:
           UpperCamelCase__				:	Optional[int]										=  BertGenerationDecoder(A_ )
           model.to(A_ )
           model.eval()
           UpperCamelCase__				:	Any										=  model(A_,		attention_mask=A_,		labels=A_ )
           self.parent.assertEqual(result.logits.shape,		(self.batch_size, self.seq_length, self.vocab_size) )






      def 						__lowercase(  self       : Union[str, Any] )	->			Dict:
           UpperCamelCase__				:	Dict										=  self.prepare_config_and_inputs()
           UpperCamelCase__				:	int										=  {"input_ids": input_ids, "attention_mask": input_mask}
           return config, inputs_dict

@require_torch
class 			UpperCamelCase__			(			__snake_case				,    __snake_case				,    __snake_case				,    unittest.TestCase      ):
      a__ :  List[Any]		    =    (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
      a__ :  Any		    =    (BertGenerationDecoder,) if is_torch_available() else ()
      a__ :  Tuple		    =    (
          {'feature-extraction': BertGenerationEncoder, 'text-generation': BertGenerationDecoder}
          if is_torch_available()
          else {}
      )






      def 						__lowercase(  self       : str )	->			Any:
           UpperCamelCase__				:	Optional[Any]										=  BertGenerationEncoderTester(self )
           UpperCamelCase__				:	Dict										=  ConfigTester(self,		config_class=A_,		hidden_size=37 )






      def 						__lowercase(  self       : Tuple )	->			Optional[int]:
           self.config_tester.run_common_tests()






      def 						__lowercase(  self       : int )	->			Union[str, Any]:
           UpperCamelCase__				:	List[Any]										=  self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*A_ )






      def 						__lowercase(  self       : int )	->			List[Any]:
           UpperCamelCase__				:	Dict										=  self.model_tester.prepare_config_and_inputs()
           UpperCamelCase__				:	Optional[Any]										=  "bert"
           self.model_tester.create_and_check_model(A_,		A_,		A_,		A_ )






      def 						__lowercase(  self       : List[str] )	->			List[str]:
           UpperCamelCase__				:	Optional[int]										=  self.model_tester.prepare_config_and_inputs_for_decoder()
           self.model_tester.create_and_check_model_as_decoder(*A_ )






      def 						__lowercase(  self       : Union[str, Any] )	->			Any:
           UpperCamelCase__				:	str										=  self.model_tester.prepare_config_and_inputs_for_decoder()
           self.model_tester.create_and_check_decoder_model_past_large_inputs(*A_ )






      def 						__lowercase(  self       : List[str] )	->			Optional[Any]:
           (
               UpperCamelCase__       
           )				:	List[str]										=  self.model_tester.prepare_config_and_inputs_for_decoder()

           UpperCamelCase__				:	Optional[Any]										=  None

           self.model_tester.create_and_check_model_as_decoder(
               A_,		A_,		A_,		A_,		A_,		A_,		)






      def 						__lowercase(  self       : Union[str, Any] )	->			Optional[int]:
           UpperCamelCase__				:	List[str]										=  self.model_tester.prepare_config_and_inputs_for_decoder()
           self.model_tester.create_and_check_for_causal_lm(*A_ )






      @slow
      def 						__lowercase(  self       : Any )	->			int:
           UpperCamelCase__				:	Union[str, Any]										=  BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''' )
           self.assertIsNotNone(A_ )

@require_torch
class 			UpperCamelCase__			(			unittest.TestCase      ):






      @slow
      def 						__lowercase(  self       : int )	->			Tuple:
           UpperCamelCase__				:	Any										=  BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''' )
           UpperCamelCase__				:	int										=  torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]] )
           with torch.no_grad():
                UpperCamelCase__				:	int										=  model(A_ )[0]
           UpperCamelCase__				:	Tuple										=  torch.Size([1, 8, 10_24] )
           self.assertEqual(output.shape,		A_ )
           UpperCamelCase__				:	Tuple										=  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 			UpperCamelCase__			(			unittest.TestCase      ):






      @slow
      def 						__lowercase(  self       : Union[str, Any] )	->			str:
           UpperCamelCase__				:	int										=  BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder''' )
           UpperCamelCase__				:	Optional[Any]										=  torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]] )
           with torch.no_grad():
                UpperCamelCase__				:	List[str]										=  model(A_ )[0]
           UpperCamelCase__				:	List[str]										=  torch.Size([1, 8, 5_03_58] )
           self.assertEqual(output.shape,		A_ )
           UpperCamelCase__				:	Optional[int]										=  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 ) )











from typing import TYPE_CHECKING

from ...utils import _LazyModule


__lowerCamelCase			:	str					    =	{"""tokenization_byt5""": ["""ByT5Tokenizer"""]}


if TYPE_CHECKING:
    from .tokenization_byta import ByTaTokenizer
else:
    import sys

    __lowerCamelCase			:	List[str]					    =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
