





from __future__ import annotations

import unittest

from transformers import FunnelConfig, is_tf_available
from transformers.testing_utils import require_tf

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers import (
          TFFunnelBaseModel,
          TFFunnelForMaskedLM,
          TFFunnelForMultipleChoice,
          TFFunnelForPreTraining,
          TFFunnelForQuestionAnswering,
          TFFunnelForSequenceClassification,
          TFFunnelForTokenClassification,
          TFFunnelModel,
      )

class 		SCREAMING_SNAKE_CASE__    :


      def __init__(self	:				str    ,      a__	:				Optional[Any]    ,      a__	:				Tuple=13    ,      a__	:				Union[str, Any]=7    ,      a__	:				List[str]=True    ,      a__	:				Dict=True    ,      a__	:				Tuple=True    ,      a__	:				int=True    ,      a__	:				Optional[int]=99    ,      a__	:				Any=[1, 1, 2]    ,      a__	:				List[Any]=1    ,      a__	:				Optional[Any]=32    ,      a__	:				List[str]=4    ,      a__	:				Union[str, Any]=8    ,      a__	:				int=37    ,      a__	:				Any="gelu_new"    ,      a__	:				Optional[Any]=0.1    ,      a__	:				List[str]=0.1    ,      a__	:				Dict=0.0    ,      a__	:				Tuple=512    ,      a__	:				Any=3    ,      a__	:				List[str]=0.0_2    ,      a__	:				List[str]=3    ,      a__	:				Optional[Any]=4    ,      a__	:				int=None    ,      a__	:				Dict=False    ,      ):





         """simple docstring"""



         __snake_case								=       parent
         __snake_case								=       batch_size
         __snake_case								=       seq_length
         __snake_case								=       is_training
         __snake_case								=       use_input_mask
         __snake_case								=       use_token_type_ids
         __snake_case								=       use_labels
         __snake_case								=       vocab_size
         __snake_case								=       block_sizes
         __snake_case								=       num_decoder_layers
         __snake_case								=       d_model
         __snake_case								=       n_head
         __snake_case								=       d_head
         __snake_case								=       d_inner
         __snake_case								=       hidden_act
         __snake_case								=       hidden_dropout
         __snake_case								=       attention_dropout
         __snake_case								=       activation_dropout
         __snake_case								=       max_position_embeddings
         __snake_case								=       type_vocab_size
         __snake_case								=       2
         __snake_case								=       num_labels
         __snake_case								=       num_choices
         __snake_case								=       scope
         __snake_case								=       initializer_std

         # Used in the tests to check the size of the first attention layer
         __snake_case								=       n_head
         # Used in the tests to check the size of the first hidden state
         __snake_case								=       self.d_model
         # Used in the tests to check the number of output hidden states/attentions
         __snake_case								=       sum(self.block_sizes				) + (0 if base else self.num_decoder_layers)
         # FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
         # the last hidden state of the first block (which is the first hidden state of the decoder).
         if not base:
            __snake_case								=       self.num_hidden_layers + 2


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



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

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

         __snake_case								=       None
         if self.use_token_type_ids:
            __snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      self.type_vocab_size				)

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

         __snake_case								=       FunnelConfig(
             vocab_size=self.vocab_size    ,      block_sizes=self.block_sizes    ,      num_decoder_layers=self.num_decoder_layers    ,      d_model=self.d_model    ,      n_head=self.n_head    ,      d_head=self.d_head    ,      d_inner=self.d_inner    ,      hidden_act=self.hidden_act    ,      hidden_dropout=self.hidden_dropout    ,      attention_dropout=self.attention_dropout    ,      activation_dropout=self.activation_dropout    ,      max_position_embeddings=self.max_position_embeddings    ,      type_vocab_size=self.type_vocab_size    ,      initializer_std=self.initializer_std    ,      )

         return (
             config,
             input_ids,
             token_type_ids,
             input_mask,
             sequence_labels,
             token_labels,
             choice_labels,
         )


      def        a			(self	:				Optional[Any]    ,      a__	:				Optional[Any]    ,      a__	:				Dict    ,      a__	:				str    ,      a__	:				List[str]    ,      a__	:				str    ,      a__	:				List[str]    ,      a__	:				List[str]    ,      ):





         """simple docstring"""



         __snake_case								=       TFFunnelModel(config=a__				)
         __snake_case								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __snake_case								=       model(a__				)

         __snake_case								=       [input_ids, input_mask]
         __snake_case								=       model(a__				)

         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.d_model)				)

         __snake_case								=       False
         __snake_case								=       TFFunnelModel(config=a__				)
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.d_model)				)

         __snake_case								=       False
         __snake_case								=       TFFunnelModel(config=a__				)
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.d_model)				)


      def        a			(self	:				str    ,      a__	:				Tuple    ,      a__	:				Union[str, Any]    ,      a__	:				List[Any]    ,      a__	:				Union[str, Any]    ,      a__	:				Optional[Any]    ,      a__	:				Tuple    ,      a__	:				Tuple    ,      ):





         """simple docstring"""



         __snake_case								=       TFFunnelBaseModel(config=a__				)
         __snake_case								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __snake_case								=       model(a__				)

         __snake_case								=       [input_ids, input_mask]
         __snake_case								=       model(a__				)

         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, 2, self.d_model)				)

         __snake_case								=       False
         __snake_case								=       TFFunnelBaseModel(config=a__				)
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, 3, self.d_model)				)

         __snake_case								=       False
         __snake_case								=       TFFunnelBaseModel(config=a__				)
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, 2, self.d_model)				)


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





         """simple docstring"""



         __snake_case								=       TFFunnelForPreTraining(config=a__				)
         __snake_case								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.seq_length)				)


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





         """simple docstring"""



         __snake_case								=       TFFunnelForMaskedLM(config=a__				)
         __snake_case								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.seq_length, self.vocab_size)				)


      def        a			(self	:				str    ,      a__	:				Optional[Any]    ,      a__	:				int    ,      a__	:				List[str]    ,      a__	:				Optional[int]    ,      a__	:				List[Any]    ,      a__	:				int    ,      a__	:				Union[str, Any]    ,      ):





         """simple docstring"""



         __snake_case								=       self.num_labels
         __snake_case								=       TFFunnelForSequenceClassification(config=a__				)
         __snake_case								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.num_labels)				)


      def        a			(self	:				Any    ,      a__	:				Tuple    ,      a__	:				List[str]    ,      a__	:				Optional[int]    ,      a__	:				Optional[int]    ,      a__	:				Optional[Any]    ,      a__	:				Optional[int]    ,      a__	:				Optional[Any]    ,      ):





         """simple docstring"""



         __snake_case								=       self.num_choices
         __snake_case								=       TFFunnelForMultipleChoice(config=a__				)
         __snake_case								=       tf.tile(tf.expand_dims(a__    ,      1				)    ,      (1, self.num_choices, 1)				)
         __snake_case								=       tf.tile(tf.expand_dims(a__    ,      1				)    ,      (1, self.num_choices, 1)				)
         __snake_case								=       tf.tile(tf.expand_dims(a__    ,      1				)    ,      (1, self.num_choices, 1)				)
         __snake_case								=       {
             '''input_ids''': multiple_choice_inputs_ids,
             '''attention_mask''': multiple_choice_input_mask,
             '''token_type_ids''': multiple_choice_token_type_ids,
         }
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.num_choices)				)


      def        a			(self	:				Optional[int]    ,      a__	:				Tuple    ,      a__	:				List[Any]    ,      a__	:				Any    ,      a__	:				Dict    ,      a__	:				Dict    ,      a__	:				int    ,      a__	:				Union[str, Any]    ,      ):





         """simple docstring"""



         __snake_case								=       self.num_labels
         __snake_case								=       TFFunnelForTokenClassification(config=a__				)
         __snake_case								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.seq_length, self.num_labels)				)


      def        a			(self	:				List[str]    ,      a__	:				Tuple    ,      a__	:				Tuple    ,      a__	:				Optional[int]    ,      a__	:				Union[str, Any]    ,      a__	:				List[Any]    ,      a__	:				Dict    ,      a__	:				List[str]    ,      ):





         """simple docstring"""



         __snake_case								=       TFFunnelForQuestionAnswering(config=a__				)
         __snake_case								=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
         __snake_case								=       model(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        a			(self	:				str				):





         """simple docstring"""



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






@require_tf
class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    Tuple							=   (
          (
              TFFunnelModel,
              TFFunnelForMaskedLM,
              TFFunnelForPreTraining,
              TFFunnelForQuestionAnswering,
              TFFunnelForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      A_						:    Union[str, Any]							=   (
          {
              'feature-extraction': (TFFunnelBaseModel, TFFunnelModel),
              'fill-mask': TFFunnelForMaskedLM,
              'question-answering': TFFunnelForQuestionAnswering,
              'text-classification': TFFunnelForSequenceClassification,
              'token-classification': TFFunnelForTokenClassification,
              'zero-shot': TFFunnelForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      A_						:    int							=   False
      A_						:    Any							=   False


      def        a			(self	:				Tuple				):





         """simple docstring"""



         __snake_case								=       TFFunnelModelTester(self				)
         __snake_case								=       ConfigTester(self    ,      config_class=a__				)


      def        a			(self	:				Tuple				):





         """simple docstring"""



         self.config_tester.run_common_tests()


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__				)


      def        a			(self	:				int				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*a__				)


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*a__				)


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*a__				)


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*a__				)






@require_tf
class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    List[str]							=   (
          (TFFunnelBaseModel, TFFunnelForMultipleChoice, TFFunnelForSequenceClassification) if is_tf_available() else ()
      )
      A_						:    Optional[Any]							=   False
      A_						:    Tuple							=   False


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       TFFunnelModelTester(self    ,      base=a__				)
         __snake_case								=       ConfigTester(self    ,      config_class=a__				)


      def        a			(self	:				Union[str, Any]				):





         """simple docstring"""



         self.config_tester.run_common_tests()


      def        a			(self	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_base_model(*a__				)


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*a__				)


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*a__				)








import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
snake_case_           =      data_utils.TransfoXLTokenizer
snake_case_           =      data_utils.TransfoXLCorpus
snake_case_           =      data_utils
snake_case_           =      data_utils






def lowerCamelCase__   (							snake_case_     :  Union[str, Any]	,				snake_case_     :  int	,				snake_case_     :  List[Any]	,				snake_case_     :  int	)				->    Dict:
   if transfo_xl_dataset_file:
      # Convert a pre-processed corpus (see original TensorFlow repo)
      with open(snake_case_	,				'''rb'''	) as fp:
         __snake_case								=       pickle.load(snake_case_	,				encoding='''latin1'''	)
      # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
      __snake_case								=       pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
      print(f"""Save vocabulary to {pytorch_vocab_dump_path}"""	)
      __snake_case								=       corpus.vocab.__dict__
      torch.save(snake_case_	,				snake_case_	)

      __snake_case								=       corpus.__dict__
      corpus_dict_no_vocab.pop('''vocab'''	,				snake_case_	)
      __snake_case								=       pytorch_dump_folder_path + '''/''' + CORPUS_NAME
      print(f"""Save dataset to {pytorch_dataset_dump_path}"""	)
      torch.save(snake_case_	,				snake_case_	)

   if tf_checkpoint_path:
      # Convert a pre-trained TensorFlow model
      __snake_case								=       os.path.abspath(snake_case_	)
      __snake_case								=       os.path.abspath(snake_case_	)

      print(f"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""	)
      # Initialise PyTorch model
      if transfo_xl_config_file == "":
         __snake_case								=       TransfoXLConfig()
      else:
         __snake_case								=       TransfoXLConfig.from_json_file(snake_case_	)
      print(f"""Building PyTorch model from configuration: {config}"""	)
      __snake_case								=       TransfoXLLMHeadModel(snake_case_	)

      __snake_case								=       load_tf_weights_in_transfo_xl(snake_case_	,				snake_case_	,				snake_case_	)
      # Save pytorch-model
      __snake_case								=       os.path.join(snake_case_	,				snake_case_	)
      __snake_case								=       os.path.join(snake_case_	,				snake_case_	)
      print(f"""Save PyTorch model to {os.path.abspath(snake_case_	)}"""	)
      torch.save(model.state_dict()	,				snake_case_	)
      print(f"""Save configuration file to {os.path.abspath(snake_case_	)}"""	)
      with open(snake_case_	,				'''w'''	,				encoding='''utf-8'''	) as f:
         f.write(config.to_json_string()	)


if __name__ == "__main__":
    snake_case_           =      argparse.ArgumentParser()
    parser.add_argument(
        '--pytorch_dump_folder_path',
        default=None,
        type=str,
        required=True,
        help='Path to the folder to store the PyTorch model or dataset/vocab.',
    )
    parser.add_argument(
        '--tf_checkpoint_path',
        default='',
        type=str,
        help='An optional path to a TensorFlow checkpoint path to be converted.',
    )
    parser.add_argument(
        '--transfo_xl_config_file',
        default='',
        type=str,
        help=(
            'An optional config json file corresponding to the pre-trained BERT model. \n'
            'This specifies the model architecture.'
        ),
    )
    parser.add_argument(
        '--transfo_xl_dataset_file',
        default='',
        type=str,
        help='An optional dataset file to be converted in a vocabulary.',
    )
    snake_case_           =      parser.parse_args()
    convert_transfo_xl_checkpoint_to_pytorch(
        args.tf_checkpoint_path,
        args.transfo_xl_config_file,
        args.pytorch_dump_folder_path,
        args.transfo_xl_dataset_file,
    )

