




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       A      :
     """simple docstring"""

     def __init__(     self	: str						,	A__	: Any						,	A__	: List[str]=1_3						,	A__	: Any=7						,	A__	: Tuple=True						,	A__	: int=True						,	A__	: int=True						,	A__	: Optional[int]=True						,	A__	: Any=9_9						,	A__	: Dict=[1, 1, 2]						,	A__	: Optional[int]=1						,	A__	: Union[str, Any]=3_2						,	A__	: Dict=4						,	A__	: List[Any]=8						,	A__	: Dict=3_7						,	A__	: Tuple="gelu_new"						,	A__	: Any=0.1						,	A__	: List[Any]=0.1						,	A__	: str=0.0						,	A__	: Any=5_1_2						,	A__	: int=3						,	A__	: Dict=0.02						,	A__	: List[str]=3						,	A__	: List[str]=4						,	A__	: List[Any]=None						,	A__	: List[str]=False						,	)					->							Optional[Any]:
           lowerCamelCase__   	=					parent
           lowerCamelCase__   	=					batch_size
           lowerCamelCase__   	=					seq_length
           lowerCamelCase__   	=					is_training
           lowerCamelCase__   	=					use_input_mask
           lowerCamelCase__   	=					use_token_type_ids
           lowerCamelCase__   	=					use_labels
           lowerCamelCase__   	=					vocab_size
           lowerCamelCase__   	=					block_sizes
           lowerCamelCase__   	=					num_decoder_layers
           lowerCamelCase__   	=					d_model
           lowerCamelCase__   	=					n_head
           lowerCamelCase__   	=					d_head
           lowerCamelCase__   	=					d_inner
           lowerCamelCase__   	=					hidden_act
           lowerCamelCase__   	=					hidden_dropout
           lowerCamelCase__   	=					attention_dropout
           lowerCamelCase__   	=					activation_dropout
           lowerCamelCase__   	=					max_position_embeddings
           lowerCamelCase__   	=					type_vocab_size
           lowerCamelCase__   	=					2
           lowerCamelCase__   	=					num_labels
           lowerCamelCase__   	=					num_choices
           lowerCamelCase__   	=					scope
           lowerCamelCase__   	=					initializer_std

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




     def 			__snake_case (     self	: str      )					->							Optional[Any]:
           lowerCamelCase__   	=					ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size      )

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

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

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

           lowerCamelCase__   	=					self.get_config()

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




     def 			__snake_case (     self	: Union[str, Any]      )					->							Tuple:
           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 			__snake_case (     self	: Tuple						,	A__	: Tuple						,	A__	: Tuple						,	A__	: Optional[int]						,	A__	: Optional[int]						,	A__	: List[Any]						,	A__	: List[str]						,	A__	: Optional[int]						,	A__	: List[str]						,	)					->							str:
           lowerCamelCase__   	=					FunnelModel(config=__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					model(__snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case      )
           lowerCamelCase__   	=					model(__snake_case						,	token_type_ids=__snake_case      )
           lowerCamelCase__   	=					model(__snake_case      )
           self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.d_model)      )

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

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




     def 			__snake_case (     self	: Tuple						,	A__	: int						,	A__	: int						,	A__	: Optional[Any]						,	A__	: int						,	A__	: Tuple						,	A__	: Union[str, Any]						,	A__	: Tuple						,	A__	: Optional[int]						,	)					->							List[str]:
           lowerCamelCase__   	=					FunnelBaseModel(config=__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					model(__snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case      )
           lowerCamelCase__   	=					model(__snake_case						,	token_type_ids=__snake_case      )
           lowerCamelCase__   	=					model(__snake_case      )
           self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, 2, self.d_model)      )

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

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




     def 			__snake_case (     self	: int						,	A__	: Dict						,	A__	: List[str]						,	A__	: str						,	A__	: Any						,	A__	: Dict						,	A__	: Any						,	A__	: Union[str, Any]						,	A__	: Tuple						,	)					->							Optional[Any]:
           lowerCamelCase__   	=					self.num_labels
           lowerCamelCase__   	=					FunnelForPreTraining(config=__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					model(__snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case						,	labels=__snake_case      )
           self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length)      )




     def 			__snake_case (     self	: List[str]						,	A__	: Union[str, Any]						,	A__	: int						,	A__	: Dict						,	A__	: Optional[Any]						,	A__	: str						,	A__	: int						,	A__	: str						,	A__	: Optional[int]						,	)					->							Dict:
           lowerCamelCase__   	=					FunnelForMaskedLM(config=__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					model(__snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case						,	labels=__snake_case      )
           self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.vocab_size)      )




     def 			__snake_case (     self	: Union[str, Any]						,	A__	: List[str]						,	A__	: Union[str, Any]						,	A__	: int						,	A__	: Optional[int]						,	A__	: str						,	A__	: List[str]						,	A__	: Dict						,	A__	: Any						,	)					->							Tuple:
           lowerCamelCase__   	=					self.num_labels
           lowerCamelCase__   	=					FunnelForSequenceClassification(__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					model(__snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case						,	labels=__snake_case      )
           self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_labels)      )




     def 			__snake_case (     self	: Union[str, Any]						,	A__	: Tuple						,	A__	: Tuple						,	A__	: List[Any]						,	A__	: Any						,	A__	: Dict						,	A__	: Optional[int]						,	A__	: Any						,	A__	: Dict						,	)					->							Tuple:
           lowerCamelCase__   	=					self.num_choices
           lowerCamelCase__   	=					FunnelForMultipleChoice(config=__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					input_ids.unsqueeze(1      ).expand(-1						,	self.num_choices						,	-1      ).contiguous()
           lowerCamelCase__   	=					token_type_ids.unsqueeze(1      ).expand(-1						,	self.num_choices						,	-1      ).contiguous()
           lowerCamelCase__   	=					input_mask.unsqueeze(1      ).expand(-1						,	self.num_choices						,	-1      ).contiguous()
           lowerCamelCase__   	=					model(
               __snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case						,	labels=__snake_case						,	)
           self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_choices)      )




     def 			__snake_case (     self	: Union[str, Any]						,	A__	: Optional[int]						,	A__	: str						,	A__	: Tuple						,	A__	: Optional[Any]						,	A__	: Dict						,	A__	: Tuple						,	A__	: Any						,	A__	: Any						,	)					->							List[str]:
           lowerCamelCase__   	=					self.num_labels
           lowerCamelCase__   	=					FunnelForTokenClassification(config=__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					model(__snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case						,	labels=__snake_case      )
           self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.num_labels)      )




     def 			__snake_case (     self	: Any						,	A__	: Any						,	A__	: List[Any]						,	A__	: Optional[int]						,	A__	: Union[str, Any]						,	A__	: Optional[Any]						,	A__	: Optional[int]						,	A__	: Optional[int]						,	A__	: Dict						,	)					->							Optional[int]:
           lowerCamelCase__   	=					FunnelForQuestionAnswering(config=__snake_case      )
           model.to(__snake_case      )
           model.eval()
           lowerCamelCase__   	=					model(
               __snake_case						,	attention_mask=__snake_case						,	token_type_ids=__snake_case						,	start_positions=__snake_case						,	end_positions=__snake_case						,	)
           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	: int      )					->							Tuple:
           lowerCamelCase__   	=					self.prepare_config_and_inputs()
           (
               lowerCamelCase__			
           )   	=					config_and_inputs
           lowerCamelCase__   	=					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
           return config, inputs_dict






@require_torch
class       A      (		__A      ,    __A      ,    unittest.TestCase		):
     """simple docstring"""

     _snake_case:						Dict  =False
     _snake_case:						str  =False
     _snake_case:						Any  =(
         (
             FunnelModel,
             FunnelForMaskedLM,
             FunnelForPreTraining,
             FunnelForQuestionAnswering,
             FunnelForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     _snake_case:						Tuple  =(
         {
             '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 			__snake_case (     self	: int						,	A__	: Dict						,	A__	: List[Any]						,	A__	: List[Any]=False      )					->							Union[str, Any]:
           lowerCamelCase__   	=					super()._prepare_for_class(__snake_case						,	__snake_case						,	return_labels=__snake_case      )

           if return_labels:
                 if model_class in get_values(__snake_case      ):
                       lowerCamelCase__   	=					torch.zeros(
                           (self.model_tester.batch_size, self.model_tester.seq_length)						,	dtype=torch.long						,	device=__snake_case      )
           return inputs_dict




     def 			__snake_case (     self	: List[Any]      )					->							Tuple:
           lowerCamelCase__   	=					FunnelModelTester(self      )
           lowerCamelCase__   	=					ConfigTester(self						,	config_class=__snake_case      )




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




     def 			__snake_case (     self	: str      )					->							Union[str, Any]:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__snake_case      )




     def 			__snake_case (     self	: List[Any]      )					->							Tuple:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_pretraining(*__snake_case      )




     def 			__snake_case (     self	: int      )					->							Any:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*__snake_case      )




     def 			__snake_case (     self	: List[str]      )					->							Optional[Any]:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*__snake_case      )




     def 			__snake_case (     self	: List[str]      )					->							Optional[int]:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*__snake_case      )




     def 			__snake_case (     self	: Tuple						,	A__	: List[str]      )					->							str:

           if hasattr(__snake_case						,	'weight'      ) and module.weight is not None:
                 module.weight.data.fill_(3      )
           if hasattr(__snake_case						,	'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(__snake_case						,	__snake_case      ) and getattr(__snake_case						,	__snake_case      ) is not None:
                       lowerCamelCase__   	=					getattr(__snake_case						,	__snake_case      )
                       weight.data.fill_(3      )






@require_torch
class       A      (		__A      ,    unittest.TestCase		):
     """simple docstring"""

     _snake_case:						Optional[Any]  =False
     _snake_case:						List[str]  =False
     _snake_case:						Optional[int]  =(
         (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
     )




     def 			__snake_case (     self	: Any      )					->							Optional[int]:
           lowerCamelCase__   	=					FunnelModelTester(self						,	base=__snake_case      )
           lowerCamelCase__   	=					ConfigTester(self						,	config_class=__snake_case      )




     def 			__snake_case (     self	: List[Any]      )					->							int:
           self.config_tester.run_common_tests()




     def 			__snake_case (     self	: Optional[int]      )					->							int:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_base_model(*__snake_case      )




     def 			__snake_case (     self	: Dict      )					->							Dict:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*__snake_case      )




     def 			__snake_case (     self	: Union[str, Any]      )					->							int:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*__snake_case      )




     def 			__snake_case (     self	: Optional[int]      )					->							List[Any]:
           lowerCamelCase__   	=					self.model_tester.prepare_config_and_inputs_for_common()
           lowerCamelCase__   	=					True

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




     def 			__snake_case (     self	: Union[str, Any]						,	A__	: Optional[Any]      )					->							Optional[Any]:

           if hasattr(__snake_case						,	'weight'      ) and module.weight is not None:
                 module.weight.data.fill_(3      )
           if hasattr(__snake_case						,	'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(__snake_case						,	__snake_case      ) and getattr(__snake_case						,	__snake_case      ) is not None:
                       lowerCamelCase__   	=					getattr(__snake_case						,	__snake_case      )
                       weight.data.fill_(3      )






@require_torch
@require_sentencepiece
@require_tokenizers
class       A      (		unittest.TestCase		):
     """simple docstring"""

     def 			__snake_case (     self	: Optional[Any]      )					->							int:
           lowerCamelCase__   	=					1_3
           lowerCamelCase__   	=					7
           lowerCamelCase__   	=					torch.arange(0						,	batch_size * sequence_length      ).long().reshape(__snake_case						,	__snake_case      )
           lowerCamelCase__   	=					[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
           lowerCamelCase__   	=					torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]      )

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

           lowerCamelCase__   	=					torch.tensor(2_3_4_4.8_3_5_2      )
           lowerCamelCase__   	=					torch.tensor(0.8052      )
           self.assertTrue(torch.allclose(output.sum()						,	__snake_case						,	atol=1E-4      )      )
           self.assertTrue(torch.allclose(output.mean()						,	__snake_case						,	atol=1E-4      )      )

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

           lowerCamelCase__   	=					torch.tensor(2_3_4_3.8_4_2_5      )
           lowerCamelCase__   	=					torch.tensor(0.8049      )
           self.assertTrue(torch.allclose(output.sum()						,	__snake_case						,	atol=1E-4      )      )
           self.assertTrue(torch.allclose(output.mean()						,	__snake_case						,	atol=1E-4      )      )




     @slow
     def 			__snake_case (     self	: Tuple      )					->							int:
           lowerCamelCase__   	=					FunnelTokenizer.from_pretrained('huggingface/funnel-small'      )
           lowerCamelCase__   	=					FunnelModel.from_pretrained('huggingface/funnel-small'      )
           lowerCamelCase__   	=					tokenizer('Hello! I am the Funnel Transformer model.'						,	return_tensors='pt'      )
           lowerCamelCase__   	=					model(**__snake_case      )[0]

           lowerCamelCase__   	=					torch.tensor(2_3_5.7_2_4_6      )
           lowerCamelCase__   	=					torch.tensor(0.0256      )
           self.assertTrue(torch.allclose(output.sum()						,	__snake_case						,	atol=1E-4      )      )
           self.assertTrue(torch.allclose(output.mean()						,	__snake_case						,	atol=1E-4      )      )






import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp







def  A					(				):
					_snake_case							:       Optional[Any]												= ArgumentParser(
					    description=(
					        "PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes"
					    )					)

					# Optional arguments for the launch helper
					parser.add_argument("--num_cores"	,	type=UpperCAmelCase	,	default=1	,	help="Number of TPU cores to use (1 or 8)."					)

					# positional
					parser.add_argument(
					    "training_script"	,	type=UpperCAmelCase	,	help=(
					        "The full path to the single TPU training "
					        "program/script to be launched in parallel, "
					        "followed by all the arguments for the "
					        "training script"
					    )	,	)

					# rest from the training program
					parser.add_argument("training_script_args"	,	nargs=UpperCAmelCase					)

					return parser.parse_args()







def  A					(				):
					_snake_case							:       Any												= parse_args()

					# Import training_script as a module.
					_snake_case							:       Any												= Path(args.training_script					)
					sys.path.append(str(script_fpath.parent.resolve()					)					)
					_snake_case							:       Tuple												= script_fpath.stem
					_snake_case							:       str												= importlib.import_module(UpperCAmelCase					)

					# Patch sys.argv
					_snake_case							:       str												= [args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores					)]

					xmp.spawn(mod._mp_fn	,	args=()	,	nprocs=args.num_cores					)


if __name__ == "__main__":
							main()