





import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
           SqueezeBertForMaskedLM,
           SqueezeBertForMultipleChoice,
           SqueezeBertForQuestionAnswering,
           SqueezeBertForSequenceClassification,
           SqueezeBertForTokenClassification,
           SqueezeBertModel,
       )





class 			lowerCAmelCase__( _snake_case    ):






       '''simple docstring'''


       def __init__(      self    ,						__lowerCamelCase    ,						__lowerCamelCase=1_3    ,						__lowerCamelCase=7    ,						__lowerCamelCase=True    ,						__lowerCamelCase=True    ,						__lowerCamelCase=False    ,						__lowerCamelCase=True    ,						__lowerCamelCase=9_9    ,						__lowerCamelCase=3_2    ,						__lowerCamelCase=5    ,						__lowerCamelCase=4    ,						__lowerCamelCase=6_4    ,						__lowerCamelCase="gelu"    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=5_1_2    ,						__lowerCamelCase=1_6    ,						__lowerCamelCase=2    ,						__lowerCamelCase=0.02    ,						__lowerCamelCase=3    ,						__lowerCamelCase=4    ,						__lowerCamelCase=None    ,						__lowerCamelCase=2    ,						__lowerCamelCase=2    ,						__lowerCamelCase=2    ,						__lowerCamelCase=2    ,						__lowerCamelCase=4    ,						__lowerCamelCase=1    ,						)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							parent
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							batch_size
            _SCREAMING_SNAKE_CASE       :	str						    =							seq_length
            _SCREAMING_SNAKE_CASE       :	str						    =							is_training
            _SCREAMING_SNAKE_CASE       :	int						    =							use_input_mask
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							use_token_type_ids
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							use_labels
            _SCREAMING_SNAKE_CASE       :	Dict						    =							vocab_size
            _SCREAMING_SNAKE_CASE       :	Any						    =							hidden_size
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							num_hidden_layers
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							num_attention_heads
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							intermediate_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							hidden_act
            _SCREAMING_SNAKE_CASE       :	int						    =							hidden_dropout_prob
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							attention_probs_dropout_prob
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							max_position_embeddings
            _SCREAMING_SNAKE_CASE       :	int						    =							type_vocab_size
            _SCREAMING_SNAKE_CASE       :	str						    =							type_sequence_label_size
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							initializer_range
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							num_labels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							num_choices
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							scope
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							q_groups
            _SCREAMING_SNAKE_CASE       :	int						    =							k_groups
            _SCREAMING_SNAKE_CASE       :	str						    =							v_groups
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							post_attention_groups
            _SCREAMING_SNAKE_CASE       :	int						    =							intermediate_groups
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							output_groups




       def     UpperCamelCase_			(      self							)			->     str:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size							)

            _SCREAMING_SNAKE_CASE       :	str						    =							None
            if self.use_input_mask:
                 _SCREAMING_SNAKE_CASE       :	Any						    =							random_attention_mask([self.batch_size, self.seq_length]							)

            _SCREAMING_SNAKE_CASE       :	Any						    =							None
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							None
            _SCREAMING_SNAKE_CASE       :	int						    =							None
            if self.use_labels:
                 _SCREAMING_SNAKE_CASE       :	int						    =							ids_tensor([self.batch_size]    ,						self.type_sequence_label_size							)
                 _SCREAMING_SNAKE_CASE       :	Any						    =							ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels							)
                 _SCREAMING_SNAKE_CASE       :	str						    =							ids_tensor([self.batch_size]    ,						self.num_choices							)

            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.get_config()

            return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels




       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            return SqueezeBertConfig(
                embedding_size=self.hidden_size    ,						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    ,						attention_probs_dropout_prob=self.hidden_dropout_prob    ,						attention_dropout=self.attention_probs_dropout_prob    ,						max_position_embeddings=self.max_position_embeddings    ,						initializer_range=self.initializer_range    ,						q_groups=self.q_groups    ,						k_groups=self.k_groups    ,						v_groups=self.v_groups    ,						post_attention_groups=self.post_attention_groups    ,						intermediate_groups=self.intermediate_groups    ,						output_groups=self.output_groups    ,						)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							SqueezeBertModel(config=lowerCAmelCase__							)
            model.to(lowerCAmelCase__							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	Dict						    =							model(lowerCAmelCase__    ,						lowerCAmelCase__							)
            _SCREAMING_SNAKE_CASE       :	int						    =							model(lowerCAmelCase__							)
            self.parent.assertEqual(result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     Optional[int]:
            _SCREAMING_SNAKE_CASE       :	Dict						    =							SqueezeBertForMaskedLM(config=lowerCAmelCase__							)
            model.to(lowerCAmelCase__							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							model(lowerCAmelCase__    ,						attention_mask=lowerCAmelCase__    ,						labels=lowerCAmelCase__							)
            self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     int:
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							SqueezeBertForQuestionAnswering(config=lowerCAmelCase__							)
            model.to(lowerCAmelCase__							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							model(
                lowerCAmelCase__    ,						attention_mask=lowerCAmelCase__    ,						start_positions=lowerCAmelCase__    ,						end_positions=lowerCAmelCase__							)
            self.parent.assertEqual(result.start_logits.shape    ,						(self.batch_size, self.seq_length)							)
            self.parent.assertEqual(result.end_logits.shape    ,						(self.batch_size, self.seq_length)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     List[Any]:
            _SCREAMING_SNAKE_CASE       :	Dict						    =							self.num_labels
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							SqueezeBertForSequenceClassification(lowerCAmelCase__							)
            model.to(lowerCAmelCase__							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	int						    =							model(lowerCAmelCase__    ,						attention_mask=lowerCAmelCase__    ,						labels=lowerCAmelCase__							)
            self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_labels)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     Union[str, Any]:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.num_labels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							SqueezeBertForTokenClassification(config=lowerCAmelCase__							)
            model.to(lowerCAmelCase__							)
            model.eval()

            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							model(lowerCAmelCase__    ,						attention_mask=lowerCAmelCase__    ,						labels=lowerCAmelCase__							)
            self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)							)




       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     Any:
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.num_choices
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							SqueezeBertForMultipleChoice(config=lowerCAmelCase__							)
            model.to(lowerCAmelCase__							)
            model.eval()
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							input_ids.unsqueeze(1							).expand(-1    ,						self.num_choices    ,						-1							).contiguous()
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							input_mask.unsqueeze(1							).expand(-1    ,						self.num_choices    ,						-1							).contiguous()
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							model(
                lowerCAmelCase__    ,						attention_mask=lowerCAmelCase__    ,						labels=lowerCAmelCase__    ,						)
            self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)							)







       def     UpperCamelCase_			(      self							)			->     Dict:
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							self.prepare_config_and_inputs()
            ((_SCREAMING_SNAKE_CASE)    ,			(_SCREAMING_SNAKE_CASE)    ,			(_SCREAMING_SNAKE_CASE)    ,			(_SCREAMING_SNAKE_CASE)    ,			(_SCREAMING_SNAKE_CASE)    ,			(_SCREAMING_SNAKE_CASE))       :	Optional[int]						    =							config_and_inputs
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							{"input_ids": input_ids, "attention_mask": input_mask}
            return config, inputs_dict





@require_torch
class 			lowerCAmelCase__( _snake_case  ,     _snake_case  ,     unittest.TestCase    ):






       '''simple docstring'''


       __snake_case 			=    (
           (
               SqueezeBertModel,
               SqueezeBertForMaskedLM,
               SqueezeBertForMultipleChoice,
               SqueezeBertForQuestionAnswering,
               SqueezeBertForSequenceClassification,
               SqueezeBertForTokenClassification,
           )
           if is_torch_available()
           else None
       )
       __snake_case 			=    (
           {
               'feature-extraction': SqueezeBertModel,
               'fill-mask': SqueezeBertForMaskedLM,
               'question-answering': SqueezeBertForQuestionAnswering,
               'text-classification': SqueezeBertForSequenceClassification,
               'token-classification': SqueezeBertForTokenClassification,
               'zero-shot': SqueezeBertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       __snake_case 			=    False
       __snake_case 			=    True
       __snake_case 			=    False




       def     UpperCamelCase_			(      self							)			->     Optional[Any]:
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							SqueezeBertModelTester(self							)
            _SCREAMING_SNAKE_CASE       :	str						    =							ConfigTester(self    ,						config_class=lowerCAmelCase__    ,						dim=3_7							)




       def     UpperCamelCase_			(      self							)			->     Any:
            self.config_tester.run_common_tests()




       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_model(*lowerCAmelCase__							)




       def     UpperCamelCase_			(      self							)			->     int:
            _SCREAMING_SNAKE_CASE       :	Dict						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_masked_lm(*lowerCAmelCase__							)




       def     UpperCamelCase_			(      self							)			->     Dict:
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_question_answering(*lowerCAmelCase__							)




       def     UpperCamelCase_			(      self							)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_sequence_classification(*lowerCAmelCase__							)




       def     UpperCamelCase_			(      self							)			->     str:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_token_classification(*lowerCAmelCase__							)




       def     UpperCamelCase_			(      self							)			->     List[str]:
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_squeezebert_for_multiple_choice(*lowerCAmelCase__							)







       @slow
       def     UpperCamelCase_			(      self							)			->     Optional[int]:
            for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							SqueezeBertModel.from_pretrained(lowerCAmelCase__							)
                 self.assertIsNotNone(lowerCAmelCase__							)





@require_sentencepiece
@require_tokenizers
@require_torch
class 			lowerCAmelCase__( unittest.TestCase    ):






       '''simple docstring'''


       @slow
       def     UpperCamelCase_			(      self							)			->     Any:
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							SqueezeBertForSequenceClassification.from_pretrained("squeezebert/squeezebert-mnli"							)

            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							torch.tensor([[1, 2_9_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 1_3, 1_5_8_8, 2]]							)
            _SCREAMING_SNAKE_CASE       :	str						    =							model(lowerCAmelCase__							)[0]
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							torch.Size((1, 3)							)
            self.assertEqual(output.shape    ,						lowerCAmelCase__							)
            _SCREAMING_SNAKE_CASE       :	str						    =							torch.tensor([[0.6401, -0.0349, -0.6041]]							)
            self.assertTrue(torch.allclose(lowerCAmelCase__    ,						lowerCAmelCase__    ,						atol=1E-4							)							)







from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer





@dataclass
class 			lowerCAmelCase__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    42





class 			lowerCAmelCase__( __lowercase  ,     __lowercase    ):






       '''simple docstring'''


       @register_to_config
       def __init__(      self    ,						__lowerCamelCase = 3    ,						__lowerCamelCase = 3    ,						__lowerCamelCase = ("DownEncoderBlock2D",)    ,						__lowerCamelCase = ("UpDecoderBlock2D",)    ,						__lowerCamelCase = (6_4,)    ,						__lowerCamelCase = 1    ,						__lowerCamelCase = "silu"    ,						__lowerCamelCase = 3    ,						__lowerCamelCase = 3_2    ,						__lowerCamelCase = 2_5_6    ,						__lowerCamelCase = 3_2    ,						__lowerCamelCase = None    ,						__lowerCamelCase = 0.1_8215    ,						__lowerCamelCase = "group"    ,						)			->     List[str]:
            super().__init__()

            # pass init params to Encoder
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							Encoder(
                in_channels=__lowerCamelCase    ,						out_channels=__lowerCamelCase    ,						down_block_types=__lowerCamelCase    ,						block_out_channels=__lowerCamelCase    ,						layers_per_block=__lowerCamelCase    ,						act_fn=__lowerCamelCase    ,						norm_num_groups=__lowerCamelCase    ,						double_z=__lowerCamelCase    ,						)

            _SCREAMING_SNAKE_CASE       :	Dict						    =							vq_embed_dim if vq_embed_dim is not None else latent_channels

            _SCREAMING_SNAKE_CASE       :	str						    =							nn.Convad(__lowerCamelCase    ,						__lowerCamelCase    ,						1							)
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							VectorQuantizer(__lowerCamelCase    ,						__lowerCamelCase    ,						beta=0.25    ,						remap=__lowerCamelCase    ,						sane_index_shape=__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	int						    =							nn.Convad(__lowerCamelCase    ,						__lowerCamelCase    ,						1							)

            # pass init params to Decoder
            _SCREAMING_SNAKE_CASE       :	Dict						    =							Decoder(
                in_channels=__lowerCamelCase    ,						out_channels=__lowerCamelCase    ,						up_block_types=__lowerCamelCase    ,						block_out_channels=__lowerCamelCase    ,						layers_per_block=__lowerCamelCase    ,						act_fn=__lowerCamelCase    ,						norm_num_groups=__lowerCamelCase    ,						norm_type=__lowerCamelCase    ,						)




       @apply_forward_hook
       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase = True							)			->     VQEncoderOutput:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.encoder(__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							self.quant_conv(__lowerCamelCase							)

            if not return_dict:
                 return (h,)

            return VQEncoderOutput(latents=__lowerCamelCase							)




       @apply_forward_hook
       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase = False    ,						__lowerCamelCase = True							)			->     Union[DecoderOutput, torch.FloatTensor]:
            # also go through quantization layer
            if not force_not_quantize:
                 _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.quantize(__lowerCamelCase							)
            else:
                 _SCREAMING_SNAKE_CASE       :	Dict						    =							h
            _SCREAMING_SNAKE_CASE       :	str						    =							self.post_quant_conv(__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Dict						    =							self.decoder(__lowerCamelCase    ,						quant if self.config.norm_type == "spatial" else None							)

            if not return_dict:
                 return (dec,)

            return DecoderOutput(sample=__lowerCamelCase							)







       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase = True							)			->     Union[DecoderOutput, torch.FloatTensor]:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							sample
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.encode(__lowerCamelCase							).latents
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.decode(__lowerCamelCase							).sample

            if not return_dict:
                 return (dec,)

            return DecoderOutput(sample=__lowerCamelCase							)