



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





 """simple docstring"""



 def __init__(    self						,     a						,     a=1_3						,     a=7						,     a=True						,     a=True						,     a=9_9						,     a=3_2						,     a=5						,     a=4						,     a=3_7						,     a="gelu"						,     a=0.1						,     a=0.1						,     a=5_0						,     a=0.02						,     a=True						,     a=None						,     )						->				Union[str, Any]:




   """simple docstring"""







   _A							     =     parent
   _A							     =     batch_size
   _A							     =     seq_length
   _A							     =     is_training
   _A							     =     use_input_mask
   _A							     =     vocab_size
   _A							     =     hidden_size
   _A							     =     num_hidden_layers
   _A							     =     num_attention_heads
   _A							     =     intermediate_size
   _A							     =     hidden_act
   _A							     =     hidden_dropout_prob
   _A							     =     attention_probs_dropout_prob
   _A							     =     max_position_embeddings
   _A							     =     initializer_range
   _A							     =     use_labels
   _A							     =     scope
 def lowercase_		(    self       )						->				Any:




   """simple docstring"""







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

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

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

   _A							     =     self.get_config()

   return config, input_ids, input_mask, token_labels
 def lowercase_		(    self       )						->				Tuple:




   """simple docstring"""







   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:




   """simple docstring"""







   (
       (
       _A      
   )     , (
       _A      
   )     , (
       _A      
   )     , (
       _A      
   )     ,       
   )							     =     self.prepare_config_and_inputs()

   _A							     =     True
   _A							     =     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
   _A							     =     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						,     a						,     a						,     a						,     a						,     **a						,     )						->				Optional[Any]:




   """simple docstring"""







   _A							     =     BertGenerationEncoder(config=a       )
   model.to(a       )
   model.eval()
   _A							     =     model(a						,     attention_mask=a       )
   _A							     =     model(a       )
   self.parent.assertEqual(result.last_hidden_state.shape						,     (self.batch_size, self.seq_length, self.hidden_size)       )
 def lowercase_		(    self						,     a						,     a						,     a						,     a						,     a						,     a						,     **a						,     )						->				List[Any]:




   """simple docstring"""







   _A							     =     True
   _A							     =     BertGenerationEncoder(config=a       )
   model.to(a       )
   model.eval()
   _A							     =     model(
       a						,     attention_mask=a						,     encoder_hidden_states=a						,     encoder_attention_mask=a						,     )
   _A							     =     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						,     a						,     a						,     a						,     a						,     a						,     a						,     **a						,     )						->				Union[str, Any]:




   """simple docstring"""







   _A							     =     True
   _A							     =     True
   _A							     =     BertGenerationDecoder(config=a       ).to(a       ).eval()

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

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

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

   _A							     =     model(
       a						,     attention_mask=a						,     encoder_hidden_states=a						,     encoder_attention_mask=a						,     output_hidden_states=a						,     )['''hidden_states'''][0]
   _A							     =     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
   _A							     =     ids_tensor((1,)						,     output_from_past.shape[-1]       ).item()
   _A							     =     output_from_no_past[:, -3:, random_slice_idx].detach()
   _A							     =     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						,     a						,     a						,     a						,     a						,     *a						,     )						->				str:




   """simple docstring"""







   _A							     =     BertGenerationDecoder(a       )
   model.to(a       )
   model.eval()
   _A							     =     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       )						->				int:




   """simple docstring"""







   _A     , _A     , _A     , _A							     =     self.prepare_config_and_inputs()
   _A							     =     {'''input_ids''': input_ids, '''attention_mask''': input_mask}
   return config, inputs_dict







@require_torch
class 		_snake_case   (      _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 lowercase_		(    self       )						->				List[Any]:




   """simple docstring"""







   _A							     =     BertGenerationEncoderTester(self       )
   _A							     =     ConfigTester(self						,     config_class=a						,     hidden_size=3_7       )
 def lowercase_		(    self       )						->				Tuple:




   """simple docstring"""







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




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*a       )
 def lowercase_		(    self       )						->				Any:




   """simple docstring"""







   _A     , _A     , _A     , _A							     =     self.model_tester.prepare_config_and_inputs()
   _A							     =     '''bert'''
   self.model_tester.create_and_check_model(a						,     a						,     a						,     a       )
 def lowercase_		(    self       )						->				int:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_model_as_decoder(*a       )
 def lowercase_		(    self       )						->				List[Any]:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_decoder_model_past_large_inputs(*a       )
 def lowercase_		(    self       )						->				Tuple:




   """simple docstring"""







   (
       (
       _A      
   )     , (
       _A      
   )     , (
       _A      
   )     , (
       _A      
   )     , (
       _A      
   )     , (
       _A      
   )     ,       
   )							     =     self.model_tester.prepare_config_and_inputs_for_decoder()

   _A							     =     None

   self.model_tester.create_and_check_model_as_decoder(
       a						,     a						,     a						,     a						,     a						,     a						,     )
 def lowercase_		(    self       )						->				Optional[int]:




   """simple docstring"""







   _A							     =     self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_for_causal_lm(*a       )



 @slow
 def lowercase_		(    self       )						->				Optional[Any]:




   """simple docstring"""







   _A							     =     BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''       )
   self.assertIsNotNone(a       )







@require_torch
class 		_snake_case   (      unittest.TestCase					):





 """simple docstring"""



 @slow
 def lowercase_		(    self       )						->				Optional[int]:




   """simple docstring"""







   _A							     =     BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''       )
   _A							     =     torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]       )
   with torch.no_grad():
     _A							     =     model(a       )[0]
   _A							     =     torch.Size([1, 8, 1_0_2_4]       )
   self.assertEqual(output.shape						,     a       )
   _A							     =     torch.tensor(
       [[[0.1_775, 0.0_083, -0.0_321], [1.6_002, 0.1_287, 0.3_912], [2.1_473, 0.5_791, 0.6_066]]]       )
   self.assertTrue(torch.allclose(output[:, :3, :3]						,     a						,     atol=1e-4       )       )







@require_torch
class 		_snake_case   (      unittest.TestCase					):





 """simple docstring"""



 @slow
 def lowercase_		(    self       )						->				str:




   """simple docstring"""







   _A							     =     BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''       )
   _A							     =     torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]       )
   with torch.no_grad():
     _A							     =     model(a       )[0]
   _A							     =     torch.Size([1, 8, 5_0_3_5_8]       )
   self.assertEqual(output.shape						,     a       )
   _A							     =     torch.tensor(
       [[[-0.5_788, -2.5_994, -3.7_054], [0.0_438, 4.7_997, 1.8_795], [1.5_862, 6.6_409, 4.4_638]]]       )
   self.assertTrue(torch.allclose(output[:, :3, :3]						,     a						,     atol=1e-4       )       )







'''simple docstring'''



import copy
import os
from typing import TYPE_CHECKING, List, Union


if TYPE_CHECKING:
 pass

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase   :							int         =   logging.get_logger(__name__)

__UpperCamelCase   :							int         =   {
    """kakaobrain/align-base""": """https://huggingface.co/kakaobrain/align-base/resolve/main/config.json""",
}
class 				__SCREAMING_SNAKE_CASE							(				_lowerCAmelCase  ):
      __a	       ="align_text_model"




      def __init__(	self			, lowerCamelCase=3_0522			, lowerCamelCase=768			, lowerCamelCase=12			, lowerCamelCase=12			, lowerCamelCase=3072			, lowerCamelCase="gelu"			, lowerCamelCase=0.1			, lowerCamelCase=0.1			, lowerCamelCase=512			, lowerCamelCase=2			, lowerCamelCase=0.02			, lowerCamelCase=1e-12			, lowerCamelCase=0			, lowerCamelCase="absolute"			, lowerCamelCase=True			, **lowerCamelCase			, )				->List[Any]:



         '''simple docstring'''
         super().__init__(**lowerCamelCase   )

         __a					  =     vocab_size
         __a					  =     hidden_size
         __a					  =     num_hidden_layers
         __a					  =     num_attention_heads
         __a					  =     hidden_act
         __a					  =     intermediate_size
         __a					  =     hidden_dropout_prob
         __a					  =     attention_probs_dropout_prob
         __a					  =     max_position_embeddings
         __a					  =     type_vocab_size
         __a					  =     initializer_range
         __a					  =     layer_norm_eps
         __a					  =     position_embedding_type
         __a					  =     use_cache
         __a					  =     pad_token_id






      @classmethod
      def 			__UpperCamelCase				(	cls			, lowerCamelCase			, **lowerCamelCase   )				->"PretrainedConfig":



         '''simple docstring'''
         cls._set_token_in_kwargs(lowerCamelCase   )

         __a  ,      __a					  =     cls.get_config_dict(lowerCamelCase			, **lowerCamelCase   )

         # get the text config dict if we are loading from AlignConfig
         if config_dict.get('model_type'   ) == "align":
            __a					  =     config_dict['text_config']

         if "model_type" in config_dict and hasattr(cls			, 'model_type'   ) and config_dict["model_type"] != cls.model_type:
            logger.warning(
                F"""You are using a model of type {config_dict["model_type"]} to instantiate a model of type """
                F"""{cls.model_type}. This is not supported for all configurations of models and can yield errors."""   )

         return cls.from_dict(lowerCamelCase			, **lowerCamelCase   )







class 				__SCREAMING_SNAKE_CASE							(				_lowerCAmelCase  ):
      __a	       ="align_vision_model"




      def __init__(	self			, lowerCamelCase = 3			, lowerCamelCase = 600			, lowerCamelCase = 2.0			, lowerCamelCase = 3.1			, lowerCamelCase = 8			, lowerCamelCase = [3, 3, 5, 3, 5, 5, 3]			, lowerCamelCase = [32, 16, 24, 40, 80, 112, 192]			, lowerCamelCase = [16, 24, 40, 80, 112, 192, 320]			, lowerCamelCase = []			, lowerCamelCase = [1, 2, 2, 2, 1, 2, 1]			, lowerCamelCase = [1, 2, 2, 3, 3, 4, 1]			, lowerCamelCase = [1, 6, 6, 6, 6, 6, 6]			, lowerCamelCase = 0.25			, lowerCamelCase = "swish"			, lowerCamelCase = 2560			, lowerCamelCase = "mean"			, lowerCamelCase = 0.02			, lowerCamelCase = 0.001			, lowerCamelCase = 0.99			, lowerCamelCase = 0.2			, **lowerCamelCase			, )				->Dict:



         '''simple docstring'''
         super().__init__(**lowerCamelCase   )

         __a					  =     num_channels
         __a					  =     image_size
         __a					  =     width_coefficient
         __a					  =     depth_coefficient
         __a					  =     depth_divisor
         __a					  =     kernel_sizes
         __a					  =     in_channels
         __a					  =     out_channels
         __a					  =     depthwise_padding
         __a					  =     strides
         __a					  =     num_block_repeats
         __a					  =     expand_ratios
         __a					  =     squeeze_expansion_ratio
         __a					  =     hidden_act
         __a					  =     hidden_dim
         __a					  =     pooling_type
         __a					  =     initializer_range
         __a					  =     batch_norm_eps
         __a					  =     batch_norm_momentum
         __a					  =     drop_connect_rate
         __a					  =     sum(lowerCamelCase   ) * 4






      @classmethod
      def 			__UpperCamelCase				(	cls			, lowerCamelCase			, **lowerCamelCase   )				->"PretrainedConfig":



         '''simple docstring'''
         cls._set_token_in_kwargs(lowerCamelCase   )

         __a  ,      __a					  =     cls.get_config_dict(lowerCamelCase			, **lowerCamelCase   )

         # get the vision config dict if we are loading from AlignConfig
         if config_dict.get('model_type'   ) == "align":
            __a					  =     config_dict['vision_config']

         if "model_type" in config_dict and hasattr(cls			, 'model_type'   ) and config_dict["model_type"] != cls.model_type:
            logger.warning(
                F"""You are using a model of type {config_dict["model_type"]} to instantiate a model of type """
                F"""{cls.model_type}. This is not supported for all configurations of models and can yield errors."""   )

         return cls.from_dict(lowerCamelCase			, **lowerCamelCase   )







class 				__SCREAMING_SNAKE_CASE							(				_lowerCAmelCase  ):
      __a	       ="align"
      __a	       =True




      def __init__(	self			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=640			, lowerCamelCase=1.0			, lowerCamelCase=0.02			, **lowerCamelCase			, )				->Union[str, Any]:



         '''simple docstring'''
         super().__init__(**lowerCamelCase   )

         if text_config is None:
            __a					  =     {}
            logger.info('text_config is None. Initializing the AlignTextConfig with default values.'   )

         if vision_config is None:
            __a					  =     {}
            logger.info('vision_config is None. Initializing the AlignVisionConfig with default values.'   )

         __a					  =     AlignTextConfig(**lowerCamelCase   )
         __a					  =     AlignVisionConfig(**lowerCamelCase   )

         __a					  =     projection_dim
         __a					  =     temperature_init_value
         __a					  =     initializer_range




      @classmethod
      def 			__UpperCamelCase				(	cls			, lowerCamelCase			, lowerCamelCase			, **lowerCamelCase   )				->int:



         '''simple docstring'''
         return cls(text_config=text_config.to_dict()			, vision_config=vision_config.to_dict()			, **lowerCamelCase   )






      def 			__UpperCamelCase				(	self   )				->Tuple:



         '''simple docstring'''
         __a					  =     copy.deepcopy(self.__dict__   )
         __a					  =     self.text_config.to_dict()
         __a					  =     self.vision_config.to_dict()
         __a					  =     self.__class__.model_type
         return output