





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     SCREAMING_SNAKE_CASE       (UpperCAmelCase				):


   def __init__(   self      :    Union[str, Any]      ,					a      :    Optional[Any]      ,					a      :    str=13      ,					a      :    List[str]=7      ,					a      :    Optional[Any]=True      ,					a      :    List[str]=True      ,					a      :    Union[str, Any]=False      ,					a      :    Optional[Any]=True      ,					a      :    Dict=99      ,					a      :    Tuple=32      ,					a      :    Any=5      ,					a      :    Any=4      ,					a      :    Any=64      ,					a      :    Any="gelu"      ,					a      :    List[str]=0.1      ,					a      :    Optional[Any]=0.1      ,					a      :    str=512      ,					a      :    List[str]=16      ,					a      :    int=2      ,					a      :    Union[str, Any]=0.02      ,					a      :    List[Any]=3      ,					a      :    str=4      ,					a      :    int=None      ,					a      :    str=2      ,					a      :    Tuple=2      ,					a      :    str=2      ,					a      :    Dict=2      ,					a      :    str=4      ,					a      :    List[Any]=1      ,					)->       List[str]:





     """simple docstring"""







     lowercase__           =  parent
     lowercase__           =  batch_size
     lowercase__           =  seq_length
     lowercase__           =  is_training
     lowercase__           =  use_input_mask
     lowercase__           =  use_token_type_ids
     lowercase__           =  use_labels
     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__           =  type_vocab_size
     lowercase__           =  type_sequence_label_size
     lowercase__           =  initializer_range
     lowercase__           =  num_labels
     lowercase__           =  num_choices
     lowercase__           =  scope
     lowercase__           =  q_groups
     lowercase__           =  k_groups
     lowercase__           =  v_groups
     lowercase__           =  post_attention_groups
     lowercase__           =  intermediate_groups
     lowercase__           =  output_groups


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





     """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]	)

     lowercase__           =  None
     lowercase__           =  None
     lowercase__           =  None
     if self.use_labels:
       lowercase__           =  ids_tensor([self.batch_size]      ,					self.type_sequence_label_size	)
       lowercase__           =  ids_tensor([self.batch_size, self.seq_length]      ,					self.num_labels	)
       lowercase__           =  ids_tensor([self.batch_size]      ,					self.num_choices	)

     lowercase__           =  self.get_config()

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


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





     """simple docstring"""







     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 	SCREAMING_SNAKE_CASE_							(   self      :    str      ,					a      :    Optional[Any]      ,					a      :    Any      ,					a      :    Optional[int]      ,					a      :    Dict      ,					a      :    Optional[int]      ,					a      :    Any	)->       str:





     """simple docstring"""







     lowercase__           =  SqueezeBertModel(config=a	)
     model.to(a	)
     model.eval()
     lowercase__           =  model(a      ,					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      :    Optional[Any]      ,					a      :    Optional[int]      ,					a      :    Optional[int]      ,					a      :    Any      ,					a      :    Optional[int]      ,					a      :    Tuple      ,					a      :    Dict	)->       Tuple:





     """simple docstring"""







     lowercase__           =  SqueezeBertForMaskedLM(config=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      :    List[str]      ,					a      :    str      ,					a      :    str      ,					a      :    Optional[Any]      ,					a      :    int      ,					a      :    Tuple      ,					a      :    Any	)->       Optional[int]:





     """simple docstring"""







     lowercase__           =  SqueezeBertForQuestionAnswering(config=a	)
     model.to(a	)
     model.eval()
     lowercase__           =  model(
         a      ,					attention_mask=a      ,					start_positions=a      ,					end_positions=a	)
     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 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[int]      ,					a      :    Dict      ,					a      :    Optional[Any]      ,					a      :    str      ,					a      :    Tuple      ,					a      :    Union[str, Any]      ,					a      :    Optional[Any]	)->       Tuple:





     """simple docstring"""







     lowercase__           =  self.num_labels
     lowercase__           =  SqueezeBertForSequenceClassification(a	)
     model.to(a	)
     model.eval()
     lowercase__           =  model(a      ,					attention_mask=a      ,					labels=a	)
     self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.num_labels)	)


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





     """simple docstring"""







     lowercase__           =  self.num_labels
     lowercase__           =  SqueezeBertForTokenClassification(config=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.num_labels)	)


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





     """simple docstring"""







     lowercase__           =  self.num_choices
     lowercase__           =  SqueezeBertForMultipleChoice(config=a	)
     model.to(a	)
     model.eval()
     lowercase__           =  input_ids.unsqueeze(1	).expand(-1      ,					self.num_choices      ,					-1	).contiguous()
     lowercase__           =  input_mask.unsqueeze(1	).expand(-1      ,					self.num_choices      ,					-1	).contiguous()
     lowercase__           =  model(
         a      ,					attention_mask=a      ,					labels=a      ,					)
     self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.num_choices)	)


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





     """simple docstring"""







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






@require_torch
class     SCREAMING_SNAKE_CASE       (UpperCAmelCase      ,     UpperCAmelCase      ,     unittest.TestCase				):
   _UpperCamelCase   :		str			 =	(
       (
           SqueezeBertModel,
           SqueezeBertForMaskedLM,
           SqueezeBertForMultipleChoice,
           SqueezeBertForQuestionAnswering,
           SqueezeBertForSequenceClassification,
           SqueezeBertForTokenClassification,
       )
       if is_torch_available()
       else None
   )
   _UpperCamelCase   :		str			 =	(
       {
           'feature-extraction': SqueezeBertModel,
           'fill-mask': SqueezeBertForMaskedLM,
           'question-answering': SqueezeBertForQuestionAnswering,
           'text-classification': SqueezeBertForSequenceClassification,
           'token-classification': SqueezeBertForTokenClassification,
           'zero-shot': SqueezeBertForSequenceClassification,
       }
       if is_torch_available()
       else {}
   )
   _UpperCamelCase   :		Dict			 =	False
   _UpperCamelCase   :		Union[str, Any]			 =	True
   _UpperCamelCase   :		List[Any]			 =	False


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





     """simple docstring"""







     lowercase__           =  SqueezeBertModelTester(self	)
     lowercase__           =  ConfigTester(self      ,					config_class=a      ,					dim=37	)


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





     """simple docstring"""







     self.config_tester.run_common_tests()


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





     """simple docstring"""







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


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





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict	)->       Tuple:





     """simple docstring"""







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


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





     """simple docstring"""







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


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





     """simple docstring"""







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


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





     """simple docstring"""







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


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





     """simple docstring"""







     for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       lowercase__           =  SqueezeBertModel.from_pretrained(a	)
       self.assertIsNotNone(a	)






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


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





     """simple docstring"""







     lowercase__           =  SqueezeBertForSequenceClassification.from_pretrained('squeezebert/squeezebert-mnli'	)

     lowercase__           =  torch.tensor([[1, 29_414, 232, 328, 740, 1_140, 12_695, 69, 13, 1_588, 2]]	)
     lowercase__           =  model(a	)[0]
     lowercase__           =  torch.Size((1, 3)	)
     self.assertEqual(output.shape      ,					a	)
     lowercase__           =  torch.tensor([[0.6401, -0.0349, -0.6041]]	)
     self.assertTrue(torch.allclose(a      ,					a      ,					atol=1E-4	)	)









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


lowercase_          =		logging.get_logger(__name__)

lowercase_          =		{
    """transfo-xl-wt103""": """https://huggingface.co/transfo-xl-wt103/resolve/main/config.json""",
}
class     SCREAMING_SNAKE_CASE       (UpperCAmelCase				):
   _UpperCamelCase   :		Optional[Any]			 =	'transfo-xl'
   _UpperCamelCase   :		Any			 =	['mems']
   _UpperCamelCase   :		Any			 =	{
       'n_token': 'vocab_size',
       'hidden_size': 'd_model',
       'num_attention_heads': 'n_head',
       'num_hidden_layers': 'n_layer',
   }


   def __init__(   self      :    Optional[Any]      ,					a      :    Optional[int]=267_735      ,					a      :    str=[20_000, 40_000, 200_000]      ,					a      :    str=1_024      ,					a      :    str=1_024      ,					a      :    int=16      ,					a      :    Optional[int]=64      ,					a      :    Optional[int]=4_096      ,					a      :    int=4      ,					a      :    Tuple=False      ,					a      :    Any=18      ,					a      :    Tuple=1_600      ,					a      :    Union[str, Any]=1_000      ,					a      :    str=True      ,					a      :    Dict=True      ,					a      :    Any=0      ,					a      :    List[Any]=-1      ,					a      :    List[Any]=True      ,					a      :    Tuple=0.1      ,					a      :    List[Any]=0.0      ,					a      :    Optional[Any]=True      ,					a      :    int="normal"      ,					a      :    Optional[Any]=0.01      ,					a      :    str=0.01      ,					a      :    List[Any]=0.02      ,					a      :    List[Any]=1E-5      ,					a      :    Optional[Any]=0      ,					**a      :    Optional[int]      ,					)->       Optional[int]:





     """simple docstring"""







     lowercase__           =  vocab_size
     lowercase__           =  []
     self.cutoffs.extend(a	)
     if proj_share_all_but_first:
       lowercase__           =  [False] + [True] * len(self.cutoffs	)
     else:
       lowercase__           =  [False] + [False] * len(self.cutoffs	)
     lowercase__           =  d_model
     lowercase__           =  d_embed
     lowercase__           =  d_head
     lowercase__           =  d_inner
     lowercase__           =  div_val
     lowercase__           =  pre_lnorm
     lowercase__           =  n_layer
     lowercase__           =  n_head
     lowercase__           =  mem_len
     lowercase__           =  same_length
     lowercase__           =  attn_type
     lowercase__           =  clamp_len
     lowercase__           =  sample_softmax
     lowercase__           =  adaptive
     lowercase__           =  dropout
     lowercase__           =  dropatt
     lowercase__           =  untie_r
     lowercase__           =  init
     lowercase__           =  init_range
     lowercase__           =  proj_init_std
     lowercase__           =  init_std
     lowercase__           =  layer_norm_epsilon
     super().__init__(eos_token_id=a      ,					**a	)


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]	)->       Union[str, Any]:





     """simple docstring"""







     logger.info(f"""The model {self.model_type} is one of the few models that has no sequence length limit."""	)
     return -1


   @max_position_embeddings.setter
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any      ,					a      :    Optional[int]	)->       Optional[int]:





     """simple docstring"""







     raise NotImplementedError(
         f"""The model {self.model_type} is one of the few models that has no sequence length limit."""	)


