



import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_PRETRAINING_MAPPING,
        FunnelBaseModel,
        FunnelForMaskedLM,
        FunnelForMultipleChoice,
        FunnelForPreTraining,
        FunnelForQuestionAnswering,
        FunnelForSequenceClassification,
        FunnelForTokenClassification,
        FunnelModel,
    )



class 						_snake_case					:





 """simple docstring"""



 def __init__(  self		,			a		,			a=1_3		,			a=7		,			a=True		,			a=True		,			a=True		,			a=True		,			a=9_9		,			a=[1, 1, 2]		,			a=1		,			a=3_2		,			a=4		,			a=8		,			a=3_7		,			a="gelu_new"		,			a=0.1		,			a=0.1		,			a=0.0		,			a=5_1_2		,			a=3		,			a=0.02		,			a=3		,			a=4		,			a=None		,			a=False		,			)		->							Optional[int]:




   """simple docstring"""







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

   # Used in the tests to check the size of the first attention layer
   _A							=			n_head
   # Used in the tests to check the size of the first hidden state
   _A							=			self.d_model
   # Used in the tests to check the number of output hidden states/attentions
   _A							=			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:
     _A							=			self.num_hidden_layers + 2
 def lowercase_						(  self    )		->							Any:




   """simple docstring"""







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

   _A							=			None
   if self.use_input_mask:
     _A							=			ids_tensor([self.batch_size, self.seq_length]		,			vocab_size=2    )

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

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

   _A							=			self.get_config()

   return (
       config,
       input_ids,
       token_type_ids,
       input_mask,
       sequence_labels,
       token_labels,
       choice_labels,
       fake_token_labels,
   )
 def lowercase_						(  self    )		->							Optional[int]:




   """simple docstring"""







   return 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		,			)
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							Optional[int]:




   """simple docstring"""







   _A							=			FunnelModel(config=a    )
   model.to(a    )
   model.eval()
   _A							=			model(a		,			attention_mask=a		,			token_type_ids=a    )
   _A							=			model(a		,			token_type_ids=a    )
   _A							=			model(a    )
   self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.d_model)    )

   _A							=			False
   _A							=			model(a    )
   self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.d_model)    )

   _A							=			False
   _A							=			model(a    )
   self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, self.seq_length, self.d_model)    )
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							Any:




   """simple docstring"""







   _A							=			FunnelBaseModel(config=a    )
   model.to(a    )
   model.eval()
   _A							=			model(a		,			attention_mask=a		,			token_type_ids=a    )
   _A							=			model(a		,			token_type_ids=a    )
   _A							=			model(a    )
   self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, 2, self.d_model)    )

   _A							=			False
   _A							=			model(a    )
   self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, 3, self.d_model)    )

   _A							=			False
   _A							=			model(a    )
   self.parent.assertEqual(result.last_hidden_state.shape		,			(self.batch_size, 2, self.d_model)    )
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							Any:




   """simple docstring"""







   _A							=			self.num_labels
   _A							=			FunnelForPreTraining(config=a    )
   model.to(a    )
   model.eval()
   _A							=			model(a		,			attention_mask=a		,			token_type_ids=a		,			labels=a    )
   self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length)    )
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							Dict:




   """simple docstring"""







   _A							=			FunnelForMaskedLM(config=a    )
   model.to(a    )
   model.eval()
   _A							=			model(a		,			attention_mask=a		,			token_type_ids=a		,			labels=a    )
   self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.vocab_size)    )
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							List[Any]:




   """simple docstring"""







   _A							=			self.num_labels
   _A							=			FunnelForSequenceClassification(a    )
   model.to(a    )
   model.eval()
   _A							=			model(a		,			attention_mask=a		,			token_type_ids=a		,			labels=a    )
   self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_labels)    )
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							Optional[Any]:




   """simple docstring"""







   _A							=			self.num_choices
   _A							=			FunnelForMultipleChoice(config=a    )
   model.to(a    )
   model.eval()
   _A							=			input_ids.unsqueeze(1    ).expand(-1		,			self.num_choices		,			-1    ).contiguous()
   _A							=			token_type_ids.unsqueeze(1    ).expand(-1		,			self.num_choices		,			-1    ).contiguous()
   _A							=			input_mask.unsqueeze(1    ).expand(-1		,			self.num_choices		,			-1    ).contiguous()
   _A							=			model(
       a		,			attention_mask=a		,			token_type_ids=a		,			labels=a		,			)
   self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.num_choices)    )
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							Union[str, Any]:




   """simple docstring"""







   _A							=			self.num_labels
   _A							=			FunnelForTokenClassification(config=a    )
   model.to(a    )
   model.eval()
   _A							=			model(a		,			attention_mask=a		,			token_type_ids=a		,			labels=a    )
   self.parent.assertEqual(result.logits.shape		,			(self.batch_size, self.seq_length, self.num_labels)    )
 def lowercase_						(  self		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			a		,			)		->							str:




   """simple docstring"""







   _A							=			FunnelForQuestionAnswering(config=a    )
   model.to(a    )
   model.eval()
   _A							=			model(
       a		,			attention_mask=a		,			token_type_ids=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 lowercase_						(  self    )		->							Tuple:




   """simple docstring"""







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







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





 """simple docstring"""



 lowerCamelCase_								=	False
 lowerCamelCase_								=	False
 lowerCamelCase_								=	(
     (
         FunnelModel,
         FunnelForMaskedLM,
         FunnelForPreTraining,
         FunnelForQuestionAnswering,
         FunnelForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 lowerCamelCase_								=	(
     {
         'feature-extraction': (FunnelBaseModel, FunnelModel),
         'fill-mask': FunnelForMaskedLM,
         'question-answering': FunnelForQuestionAnswering,
         'text-classification': FunnelForSequenceClassification,
         'token-classification': FunnelForTokenClassification,
         'zero-shot': FunnelForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 def lowercase_						(  self		,			a		,			a		,			a=False    )		->							Optional[int]:




   """simple docstring"""







   _A							=			super()._prepare_for_class(a		,			a		,			return_labels=a    )

   if return_labels:
     if model_class in get_values(a    ):
       _A							=			torch.zeros(
           (self.model_tester.batch_size, self.model_tester.seq_length)		,			dtype=torch.long		,			device=a    )
   return inputs_dict
 def lowercase_						(  self    )		->							Tuple:




   """simple docstring"""







   _A							=			FunnelModelTester(self    )
   _A							=			ConfigTester(self		,			config_class=a    )
 def lowercase_						(  self    )		->							str:




   """simple docstring"""







   self.config_tester.run_common_tests()
 def lowercase_						(  self    )		->							List[str]:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*a    )
 def lowercase_						(  self    )		->							Any:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_pretraining(*a    )
 def lowercase_						(  self    )		->							Union[str, Any]:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*a    )
 def lowercase_						(  self    )		->							Optional[Any]:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*a    )
 def lowercase_						(  self    )		->							str:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*a    )



 def lowercase_						(  self		,			a    )		->							Tuple:




   """simple docstring"""







   if hasattr(a		,			'''weight'''    ) and module.weight is not None:
     module.weight.data.fill_(3    )
   if hasattr(a		,			'''bias'''    ) and module.bias is not None:
     module.bias.data.fill_(3    )

   for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
     if hasattr(a		,			a    ) and getattr(a		,			a    ) is not None:
       _A							=			getattr(a		,			a    )
       weight.data.fill_(3    )







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





 """simple docstring"""



 lowerCamelCase_								=	False
 lowerCamelCase_								=	False
 lowerCamelCase_								=	(
     (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
 )
 def lowercase_						(  self    )		->							Optional[Any]:




   """simple docstring"""







   _A							=			FunnelModelTester(self		,			base=a    )
   _A							=			ConfigTester(self		,			config_class=a    )
 def lowercase_						(  self    )		->							Union[str, Any]:




   """simple docstring"""







   self.config_tester.run_common_tests()
 def lowercase_						(  self    )		->							Optional[int]:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_base_model(*a    )
 def lowercase_						(  self    )		->							Tuple:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*a    )
 def lowercase_						(  self    )		->							Optional[int]:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_multiple_choice(*a    )
 def lowercase_						(  self    )		->							List[str]:




   """simple docstring"""







   _A							=			self.model_tester.prepare_config_and_inputs_for_common()
   _A							=			True

   for model_class in self.all_model_classes:
     if model_class.__name__ == "FunnelBaseModel":
       continue
     _A							=			model_class(a    )
     model.to(a    )
     model.train()
     _A							=			self._prepare_for_class(a		,			a		,			return_labels=a    )
     _A							=			model(**a    ).loss
     loss.backward()



 def lowercase_						(  self		,			a    )		->							Tuple:




   """simple docstring"""







   if hasattr(a		,			'''weight'''    ) and module.weight is not None:
     module.weight.data.fill_(3    )
   if hasattr(a		,			'''bias'''    ) and module.bias is not None:
     module.bias.data.fill_(3    )

   for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
     if hasattr(a		,			a    ) and getattr(a		,			a    ) is not None:
       _A							=			getattr(a		,			a    )
       weight.data.fill_(3    )







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





 """simple docstring"""



 def lowercase_						(  self    )		->							Any:




   """simple docstring"""







   _A							=			1_3
   _A							=			7
   _A							=			torch.arange(0		,			batch_size * sequence_length    ).long().reshape(a		,			a    )
   _A							=			[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
   _A							=			torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]    )

   _A							=			FunnelModel.from_pretrained('''sgugger/funnel-random-tiny'''    )
   _A							=			model(a		,			token_type_ids=a    )[0].abs()

   _A							=			torch.tensor(2_3_4_4.8_3_5_2    )
   _A							=			torch.tensor(0.8_052    )
   self.assertTrue(torch.allclose(output.sum()		,			a		,			atol=1e-4    )    )
   self.assertTrue(torch.allclose(output.mean()		,			a		,			atol=1e-4    )    )

   _A							=			torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]    )
   _A							=			model(a		,			attention_mask=a		,			token_type_ids=a    )[0].abs()

   _A							=			torch.tensor(2_3_4_3.8_4_2_5    )
   _A							=			torch.tensor(0.8_049    )
   self.assertTrue(torch.allclose(output.sum()		,			a		,			atol=1e-4    )    )
   self.assertTrue(torch.allclose(output.mean()		,			a		,			atol=1e-4    )    )



 @slow
 def lowercase_						(  self    )		->							int:




   """simple docstring"""







   _A							=			FunnelTokenizer.from_pretrained('''huggingface/funnel-small'''    )
   _A							=			FunnelModel.from_pretrained('''huggingface/funnel-small'''    )
   _A							=			tokenizer('''Hello! I am the Funnel Transformer model.'''		,			return_tensors='''pt'''    )
   _A							=			model(**a    )[0]

   _A							=			torch.tensor(2_3_5.7_2_4_6    )
   _A							=			torch.tensor(0.0_256    )
   self.assertTrue(torch.allclose(output.sum()		,			a		,			atol=1e-4    )    )
   self.assertTrue(torch.allclose(output.mean()		,			a		,			atol=1e-4    )    )
import os
from pathlib import Path
from typing import Dict, List

import fire
import torch

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from transformers.utils.logging import get_logger


UpperCAmelCase							= get_logger(__name__)


def SCREAMING_SNAKE_CASE (   lowercase_:		str						,  lowercase_:		str ):
       if text.startswith(lowercase_ ):
              return text[len(lowercase_ ) :]
       return text  # or whatever


def SCREAMING_SNAKE_CASE (   lowercase_:		Union[str, Any] ):
       return {remove_prefix(lowercase_						,  'model.' ): v for k, v in sd.items()}


def SCREAMING_SNAKE_CASE (   lowercase_:		List[Dict[str, torch.Tensor]] ):
       _UpperCamelCase   :				Union[str, Any]       =     {}
       for k in state_dicts[0].keys():
              _UpperCamelCase   :				Tuple       =     [sd[k] for sd in state_dicts]
              _UpperCamelCase   :				str       =     sum(lowercase_ ) / len(lowercase_ )
              assert isinstance(lowercase_						,  torch.Tensor )
              _UpperCamelCase   :				Dict       =     new_t
       return new_sd


def SCREAMING_SNAKE_CASE (   lowercase_:		str						,  lowercase_:		str						,  lowercase_:		str ):
       _UpperCamelCase   :				Optional[Any]       =     AutoModelForSeqaSeqLM.from_pretrained(lowercase_ )
       if os.path.isfile(lowercase_ ):
              _UpperCamelCase   :				Any       =     [pl_ckpt_path]
       else:
              assert os.path.isdir(lowercase_ )
              _UpperCamelCase   :				str       =     list(Path(lowercase_ ).glob('*.ckpt' ) )
              assert ckpt_files, F"could not find any ckpt files inside the {pl_ckpt_path} directory"

       if len(lowercase_ ) > 1:
              logger.info(F"averaging the weights of {ckpt_files}" )

       _UpperCamelCase   :				Tuple       =     [sanitize(torch.load(lowercase_						,  map_location='cpu' )['state_dict'] ) for x in ckpt_files]
       _UpperCamelCase   :				Optional[Any]       =     average_state_dicts(lowercase_ )

       _UpperCamelCase  ,					_UpperCamelCase   :				Tuple       =     hf_model.load_state_dict(lowercase_						,  strict=lowercase_ )
       assert not missing, F"missing keys: {missing}"
       hf_model.save_pretrained(lowercase_ )
       try:
              _UpperCamelCase   :				str       =     AutoTokenizer.from_pretrained(lowercase_ )
              tok.save_pretrained(lowercase_ )
       except Exception:
              pass
              # dont copy tokenizer if cant


if __name__ == "__main__":
 fire.Fire(convert_pl_to_hf)



