





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


   def __init__(					self    :       Optional[int]    ,  a    :       List[str]    ,  a    :       int=13    ,  a    :       List[Any]=7    ,  a    :       Optional[Any]=True    ,  a    :       Dict=True    ,  a    :       str=99    ,  a    :       str=32    ,  a    :       int=5    ,  a    :       int=4    ,  a    :       Dict=37    ,  a    :       Optional[int]="gelu"    ,  a    :       Any=0.1    ,  a    :       Any=0.1    ,  a    :       int=50    ,  a    :       Optional[Any]=0.02    ,  a    :       Any=True    ,  a    :       Optional[int]=None    ,  )->     Any:





       """simple docstring"""







       lowercase__			    =						parent
       lowercase__			    =						batch_size
       lowercase__			    =						seq_length
       lowercase__			    =						is_training
       lowercase__			    =						use_input_mask
       lowercase__			    =						vocab_size
       lowercase__			    =						hidden_size
       lowercase__			    =						num_hidden_layers
       lowercase__			    =						num_attention_heads
       lowercase__			    =						intermediate_size
       lowercase__			    =						hidden_act
       lowercase__			    =						hidden_dropout_prob
       lowercase__			    =						attention_probs_dropout_prob
       lowercase__			    =						max_position_embeddings
       lowercase__			    =						initializer_range
       lowercase__			    =						use_labels
       lowercase__			    =						scope


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     str:





       """simple docstring"""







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

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

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

       lowercase__			    =						self.get_config()

       return config, input_ids, input_mask, token_labels


   def  SCREAMING_SNAKE_CASE_							(					self    :       Dict )->     Optional[Any]:





       """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  SCREAMING_SNAKE_CASE_							(					self    :       Union[str, Any] )->     Any:





       """simple docstring"""







       (
           (
           lowercase__					
       )    ,					(
           lowercase__					
       )    ,					(
           lowercase__					
       )    ,					(
           lowercase__					
       )    ,										
       )			    =						self.prepare_config_and_inputs()

       lowercase__			    =						True
       lowercase__			    =						floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
       lowercase__			    =						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  SCREAMING_SNAKE_CASE_							(					self    :       Union[str, Any]    ,  a    :       int    ,  a    :       List[Any]    ,  a    :       List[Any]    ,  a    :       Dict    ,  **a    :       List[Any]    ,  )->     Dict:





       """simple docstring"""







       lowercase__			    =						BertGenerationEncoder(config=a )
       model.to(a )
       model.eval()
       lowercase__			    =						model(a    ,  attention_mask=a )
       lowercase__			    =						model(a )
       self.parent.assertEqual(result.last_hidden_state.shape    ,  (self.batch_size, self.seq_length, self.hidden_size) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       int    ,  a    :       List[str]    ,  a    :       Any    ,  a    :       Dict    ,  a    :       int    ,  a    :       Any    ,  a    :       int    ,  **a    :       Optional[int]    ,  )->     Dict:





       """simple docstring"""







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





       """simple docstring"""







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

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

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

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

       lowercase__			    =						model(
           a    ,  attention_mask=a    ,  encoder_hidden_states=a    ,  encoder_attention_mask=a    ,  output_hidden_states=a    ,  )['hidden_states'][0]
       lowercase__			    =						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
       lowercase__			    =						ids_tensor((1,)    ,  output_from_past.shape[-1] ).item()
       lowercase__			    =						output_from_no_past[:, -3:, random_slice_idx].detach()
       lowercase__			    =						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  SCREAMING_SNAKE_CASE_							(					self    :       Dict    ,  a    :       List[str]    ,  a    :       Optional[int]    ,  a    :       int    ,  a    :       Any    ,  *a    :       Union[str, Any]    ,  )->     Tuple:





       """simple docstring"""







       lowercase__			    =						BertGenerationDecoder(a )
       model.to(a )
       model.eval()
       lowercase__			    =						model(a    ,  attention_mask=a    ,  labels=a )
       self.parent.assertEqual(result.logits.shape    ,  (self.batch_size, self.seq_length, self.vocab_size) )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Optional[int] )->     Any:





       """simple docstring"""







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






@require_torch
class        SCREAMING_SNAKE_CASE     (UpperCAmelCase    ,			UpperCAmelCase    ,			UpperCAmelCase    ,			unittest.TestCase   ):
   _UpperCamelCase					:	int				      = (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
   _UpperCamelCase					:	Union[str, Any]				      = (BertGenerationDecoder,) if is_torch_available() else ()
   _UpperCamelCase					:	str				      = (
       {'feature-extraction': BertGenerationEncoder, 'text-generation': BertGenerationDecoder}
       if is_torch_available()
       else {}
   )


   def  SCREAMING_SNAKE_CASE_							(					self    :       List[Any] )->     List[Any]:





       """simple docstring"""







       lowercase__			    =						BertGenerationEncoderTester(self )
       lowercase__			    =						ConfigTester(self    ,  config_class=a    ,  hidden_size=37 )


   def  SCREAMING_SNAKE_CASE_							(					self    :       str )->     Tuple:





       """simple docstring"""







       self.config_tester.run_common_tests()


   def  SCREAMING_SNAKE_CASE_							(					self    :       int )->     str:





       """simple docstring"""







       lowercase__			    =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*a )


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





       """simple docstring"""







       lowercase__    ,					lowercase__    ,					lowercase__    ,					lowercase__			    =						self.model_tester.prepare_config_and_inputs()
       lowercase__			    =						'bert'
       self.model_tester.create_and_check_model(a    ,  a    ,  a    ,  a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Any )->     Any:





       """simple docstring"""







       lowercase__			    =						self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Union[str, Any] )->     List[Any]:





       """simple docstring"""







       lowercase__			    =						self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_decoder_model_past_large_inputs(*a )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     Optional[Any]:





       """simple docstring"""







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

       lowercase__			    =						None

       self.model_tester.create_and_check_model_as_decoder(
           a    ,  a    ,  a    ,  a    ,  a    ,  a    ,  )


   def  SCREAMING_SNAKE_CASE_							(					self    :       Tuple )->     Optional[Any]:





       """simple docstring"""







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


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





       """simple docstring"""







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






@require_torch
class        SCREAMING_SNAKE_CASE     (unittest.TestCase   ):


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





       """simple docstring"""







       lowercase__			    =						BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder' )
       lowercase__			    =						torch.tensor([[101, 7_592, 1_010, 2_026, 3_899, 2_003, 10_140, 102]] )
       with torch.no_grad():
           lowercase__			    =						model(a )[0]
       lowercase__			    =						torch.Size([1, 8, 1_024] )
       self.assertEqual(output.shape    ,  a )
       lowercase__			    =						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        SCREAMING_SNAKE_CASE     (unittest.TestCase   ):


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





       """simple docstring"""







       lowercase__			    =						BertGenerationDecoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder' )
       lowercase__			    =						torch.tensor([[101, 7_592, 1_010, 2_026, 3_899, 2_003, 10_140, 102]] )
       with torch.no_grad():
           lowercase__			    =						model(a )[0]
       lowercase__			    =						torch.Size([1, 8, 50_358] )
       self.assertEqual(output.shape    ,  a )
       lowercase__			    =						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 OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowercase_					     =	{
    """configuration_squeezebert""": [
        """SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """SqueezeBertConfig""",
        """SqueezeBertOnnxConfig""",
    ],
    """tokenization_squeezebert""": ["""SqueezeBertTokenizer"""],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase_					     =	["""SqueezeBertTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase_					     =	[
          """SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """SqueezeBertForMaskedLM""",
          """SqueezeBertForMultipleChoice""",
          """SqueezeBertForQuestionAnswering""",
          """SqueezeBertForSequenceClassification""",
          """SqueezeBertForTokenClassification""",
          """SqueezeBertModel""",
          """SqueezeBertModule""",
          """SqueezeBertPreTrainedModel""",
      ]


if TYPE_CHECKING:
      from .configuration_squeezebert import (
          SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
          SqueezeBertConfig,
          SqueezeBertOnnxConfig,
      )
      from .tokenization_squeezebert import SqueezeBertTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_squeezebert_fast import SqueezeBertTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_squeezebert import (
                SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                SqueezeBertForMaskedLM,
                SqueezeBertForMultipleChoice,
                SqueezeBertForQuestionAnswering,
                SqueezeBertForSequenceClassification,
                SqueezeBertForTokenClassification,
                SqueezeBertModel,
                SqueezeBertModule,
                SqueezeBertPreTrainedModel,
            )

else:
      import sys

      lowercase_					     =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


