"""simple docstring"""






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



    """simple docstring"""

    def __init__(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase=13    ,			_UpperCAmelCase=7    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase=False    ,			_UpperCAmelCase=True    ,			_UpperCAmelCase=99    ,			_UpperCAmelCase=32    ,			_UpperCAmelCase=5    ,			_UpperCAmelCase=4    ,			_UpperCAmelCase=64    ,			_UpperCAmelCase="gelu"    ,			_UpperCAmelCase=0.1    ,			_UpperCAmelCase=0.1    ,			_UpperCAmelCase=512    ,			_UpperCAmelCase=16    ,			_UpperCAmelCase=2    ,			_UpperCAmelCase=0.02    ,			_UpperCAmelCase=3    ,			_UpperCAmelCase=4    ,			_UpperCAmelCase=None    ,			_UpperCAmelCase=2    ,			_UpperCAmelCase=2    ,			_UpperCAmelCase=2    ,			_UpperCAmelCase=2    ,			_UpperCAmelCase=4    ,			_UpperCAmelCase=1    ,			):
       lowercase__:   int					=      parent
       lowercase__:   Dict					=      batch_size
       lowercase__:   Tuple					=      seq_length
       lowercase__:   int					=      is_training
       lowercase__:   Union[str, Any]					=      use_input_mask
       lowercase__:   Any					=      use_token_type_ids
       lowercase__:   Union[str, Any]					=      use_labels
       lowercase__:   Any					=      vocab_size
       lowercase__:   int					=      hidden_size
       lowercase__:   List[Any]					=      num_hidden_layers
       lowercase__:   Optional[int]					=      num_attention_heads
       lowercase__:   int					=      intermediate_size
       lowercase__:   str					=      hidden_act
       lowercase__:   Dict					=      hidden_dropout_prob
       lowercase__:   Optional[int]					=      attention_probs_dropout_prob
       lowercase__:   str					=      max_position_embeddings
       lowercase__:   List[str]					=      type_vocab_size
       lowercase__:   Optional[int]					=      type_sequence_label_size
       lowercase__:   Dict					=      initializer_range
       lowercase__:   Union[str, Any]					=      num_labels
       lowercase__:   Optional[int]					=      num_choices
       lowercase__:   List[str]					=      scope
       lowercase__:   Any					=      q_groups
       lowercase__:   int					=      k_groups
       lowercase__:   Dict					=      v_groups
       lowercase__:   str					=      post_attention_groups
       lowercase__:   Tuple					=      intermediate_groups
       lowercase__:   List[str]					=      output_groups






    def _snake_case			(   self							):
       lowercase__:   str					=      ids_tensor([self.batch_size, self.seq_length]    ,			self.vocab_size							)

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

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

       lowercase__:   Optional[int]					=      self.get_config()

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






    def _snake_case			(   self							):
       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 _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Optional[Any]					=      SqueezeBertModel(config=_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()
       lowercase__:   Optional[Any]					=      model(_UpperCAmelCase    ,			_UpperCAmelCase							)
       lowercase__:   Union[str, Any]					=      model(_UpperCAmelCase							)
       self.parent.assertEqual(result.last_hidden_state.shape    ,			(self.batch_size, self.seq_length, self.hidden_size)							)






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Union[str, Any]					=      SqueezeBertForMaskedLM(config=_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()
       lowercase__:   List[str]					=      model(_UpperCAmelCase    ,			attention_mask=_UpperCAmelCase    ,			labels=_UpperCAmelCase							)
       self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.seq_length, self.vocab_size)							)






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Dict					=      SqueezeBertForQuestionAnswering(config=_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()
       lowercase__:   Dict					=      model(
           _UpperCAmelCase    ,			attention_mask=_UpperCAmelCase    ,			start_positions=_UpperCAmelCase    ,			end_positions=_UpperCAmelCase							)
       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 _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   Optional[int]					=      self.num_labels
       lowercase__:   Tuple					=      SqueezeBertForSequenceClassification(_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()
       lowercase__:   Dict					=      model(_UpperCAmelCase    ,			attention_mask=_UpperCAmelCase    ,			labels=_UpperCAmelCase							)
       self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.num_labels)							)






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   int					=      self.num_labels
       lowercase__:   Tuple					=      SqueezeBertForTokenClassification(config=_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()

       lowercase__:   Dict					=      model(_UpperCAmelCase    ,			attention_mask=_UpperCAmelCase    ,			labels=_UpperCAmelCase							)
       self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.seq_length, self.num_labels)							)






    def _snake_case			(   self    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase    ,			_UpperCAmelCase							):
       lowercase__:   int					=      self.num_choices
       lowercase__:   Dict					=      SqueezeBertForMultipleChoice(config=_UpperCAmelCase							)
       model.to(_UpperCAmelCase							)
       model.eval()
       lowercase__:   Optional[int]					=      input_ids.unsqueeze(1							).expand(-1    ,			self.num_choices    ,			-1							).contiguous()
       lowercase__:   List[Any]					=      input_mask.unsqueeze(1							).expand(-1    ,			self.num_choices    ,			-1							).contiguous()
       lowercase__:   List[str]					=      model(
           _UpperCAmelCase    ,			attention_mask=_UpperCAmelCase    ,			labels=_UpperCAmelCase    ,			)
       self.parent.assertEqual(result.logits.shape    ,			(self.batch_size, self.num_choices)							)






    def _snake_case			(   self							):
       lowercase__:   Tuple					=      self.prepare_config_and_inputs()
       ((lowercase__), (lowercase__), (lowercase__), (lowercase__), (lowercase__), (lowercase__)):   int					=      config_and_inputs
       lowercase__:   List[Any]					=      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict



@require_torch
class     UpperCAmelCase						(_UpperCAmelCase    ,_UpperCAmelCase    ,unittest.TestCase  ):



    """simple docstring"""

    _UpperCAmelCase	:Tuple						    = (
        (
            SqueezeBertModel,
            SqueezeBertForMaskedLM,
            SqueezeBertForMultipleChoice,
            SqueezeBertForQuestionAnswering,
            SqueezeBertForSequenceClassification,
            SqueezeBertForTokenClassification,
        )
        if is_torch_available()
        else None
    )
    _UpperCAmelCase	:int						    = (
        {
            "feature-extraction": SqueezeBertModel,
            "fill-mask": SqueezeBertForMaskedLM,
            "question-answering": SqueezeBertForQuestionAnswering,
            "text-classification": SqueezeBertForSequenceClassification,
            "token-classification": SqueezeBertForTokenClassification,
            "zero-shot": SqueezeBertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    _UpperCAmelCase	:Optional[int]						    = False
    _UpperCAmelCase	:Dict						    = True
    _UpperCAmelCase	:Union[str, Any]						    = False






    def _snake_case			(   self							):
       lowercase__:   List[str]					=      SqueezeBertModelTester(self							)
       lowercase__:   Optional[int]					=      ConfigTester(self    ,			config_class=_UpperCAmelCase    ,			dim=37							)






    def _snake_case			(   self							):
       self.config_tester.run_common_tests()






    def _snake_case			(   self							):
       lowercase__:   Tuple					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_squeezebert_model(*_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   Union[str, Any]					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_squeezebert_for_masked_lm(*_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   Any					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_squeezebert_for_question_answering(*_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   List[str]					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_squeezebert_for_sequence_classification(*_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   Optional[Any]					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_squeezebert_for_token_classification(*_UpperCAmelCase							)






    def _snake_case			(   self							):
       lowercase__:   List[str]					=      self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_squeezebert_for_multiple_choice(*_UpperCAmelCase							)






    @slow
    def _snake_case			(   self							):
       for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          lowercase__:   Optional[int]					=      SqueezeBertModel.from_pretrained(_UpperCAmelCase							)
          self.assertIsNotNone(_UpperCAmelCase							)



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



    """simple docstring"""

    @slow
    def _snake_case			(   self							):
       lowercase__:   Tuple					=      SqueezeBertForSequenceClassification.from_pretrained('''squeezebert/squeezebert-mnli'''							)

       lowercase__:   Union[str, Any]					=      torch.tensor([[1, 29414, 232, 328, 740, 1140, 12695, 69, 13, 1588, 2]]							)
       lowercase__:   Dict					=      model(_UpperCAmelCase							)[0]
       lowercase__:   Optional[int]					=      torch.Size((1, 3)							)
       self.assertEqual(output.shape    ,			_UpperCAmelCase							)
       lowercase__:   Tuple					=      torch.tensor([[0.6_401, -0.0_349, -0.6_041]]							)
       self.assertTrue(torch.allclose(_UpperCAmelCase    ,			_UpperCAmelCase    ,			atol=1e-4							)							)

"""simple docstring"""






import argparse

import torch

from transformers import FunnelBaseModel, FunnelConfig, FunnelModel, load_tf_weights_in_funnel
from transformers.utils import logging


logging.set_verbosity_info()



def SCREAMING_SNAKE_CASE__			(   __UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase )  ->  List[Any]:
   # Initialise PyTorch model
   lowercase__:   Optional[Any]					=      FunnelConfig.from_json_file(__UpperCAmelCase )
   print(F"""Building PyTorch model from configuration: {config}""" )
   lowercase__:   List[Any]					=      FunnelBaseModel(__UpperCAmelCase ) if base_model else FunnelModel(__UpperCAmelCase )

   # Load weights from tf checkpoint
   load_tf_weights_in_funnel(__UpperCAmelCase       ,    __UpperCAmelCase       ,    __UpperCAmelCase )

   # Save pytorch-model
   print(F"""Save PyTorch model to {pytorch_dump_path}""" )
   torch.save(model.state_dict()       ,    __UpperCAmelCase )


if __name__ == "__main__":
       __A			   = argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
       )
       parser.add_argument(
           "--config_file",
           default=None,
           type=str,
           required=True,
           help="The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.",
       )
       parser.add_argument(
           "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
       )
       parser.add_argument(
           "--base_model", action="store_true", help="Whether you want just the base model (no decoder) or not."
       )
       __A			   = parser.parse_args()
       convert_tf_checkpoint_to_pytorch(
           args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.base_model
       )
