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							:


   def __init__(self  :		List[str] ,						__UpperCAmelCase  :		int ,						__UpperCAmelCase  :		str=1_3 ,						__UpperCAmelCase  :		Optional[Any]=7 ,						__UpperCAmelCase  :		Union[str, Any]=True ,						__UpperCAmelCase  :		Optional[Any]=True ,						__UpperCAmelCase  :		Any=True ,						__UpperCAmelCase  :		str=True ,						__UpperCAmelCase  :		Union[str, Any]=9_9 ,						__UpperCAmelCase  :		Any=[1, 1, 2] ,						__UpperCAmelCase  :		str=1 ,						__UpperCAmelCase  :		List[Any]=3_2 ,						__UpperCAmelCase  :		int=4 ,						__UpperCAmelCase  :		int=8 ,						__UpperCAmelCase  :		Tuple=3_7 ,						__UpperCAmelCase  :		int="gelu_new" ,						__UpperCAmelCase  :		Dict=0.1 ,						__UpperCAmelCase  :		Optional[Any]=0.1 ,						__UpperCAmelCase  :		List[str]=0.0 ,						__UpperCAmelCase  :		Union[str, Any]=5_1_2 ,						__UpperCAmelCase  :		Union[str, Any]=3 ,						__UpperCAmelCase  :		List[str]=0.02 ,						__UpperCAmelCase  :		Union[str, Any]=3 ,						__UpperCAmelCase  :		str=4 ,						__UpperCAmelCase  :		Optional[Any]=None ,						__UpperCAmelCase  :		Any=False ,						)		->       Optional[int]:
         """simple docstring"""



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

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


   def     lowercase_   (self  :		int	)		->       List[str]:
         """simple docstring"""



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

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

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

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

         UpperCAmelCase__			=						self.get_config()

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


   def     lowercase_   (self  :		List[str]	)		->       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  :		Optional[Any] ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		Optional[Any] ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		List[str] ,						__UpperCAmelCase  :		Optional[int] ,						__UpperCAmelCase  :		List[str] ,						__UpperCAmelCase  :		List[str] ,						)		->       int:
         """simple docstring"""



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

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

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


   def     lowercase_   (self  :		List[Any] ,						__UpperCAmelCase  :		int ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		Optional[Any] ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		int ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		Tuple ,						)		->       List[Any]:
         """simple docstring"""



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

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

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


   def     lowercase_   (self  :		Dict ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		List[Any] ,						__UpperCAmelCase  :		List[Any] ,						__UpperCAmelCase  :		Optional[Any] ,						__UpperCAmelCase  :		Any ,						__UpperCAmelCase  :		str ,						__UpperCAmelCase  :		int ,						)		->       int:
         """simple docstring"""



         UpperCAmelCase__			=						self.num_labels
         UpperCAmelCase__			=						FunnelForPreTraining(config=__UpperCAmelCase	)
         model.to(__UpperCAmelCase	)
         model.eval()
         UpperCAmelCase__			=						model(__UpperCAmelCase ,						attention_mask=__UpperCAmelCase ,						token_type_ids=__UpperCAmelCase ,						labels=__UpperCAmelCase	)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length)	)


   def     lowercase_   (self  :		Tuple ,						__UpperCAmelCase  :		str ,						__UpperCAmelCase  :		Optional[Any] ,						__UpperCAmelCase  :		List[str] ,						__UpperCAmelCase  :		int ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		str ,						__UpperCAmelCase  :		List[str] ,						__UpperCAmelCase  :		List[Any] ,						)		->       List[Any]:
         """simple docstring"""



         UpperCAmelCase__			=						FunnelForMaskedLM(config=__UpperCAmelCase	)
         model.to(__UpperCAmelCase	)
         model.eval()
         UpperCAmelCase__			=						model(__UpperCAmelCase ,						attention_mask=__UpperCAmelCase ,						token_type_ids=__UpperCAmelCase ,						labels=__UpperCAmelCase	)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size)	)


   def     lowercase_   (self  :		int ,						__UpperCAmelCase  :		str ,						__UpperCAmelCase  :		Any ,						__UpperCAmelCase  :		Any ,						__UpperCAmelCase  :		Union[str, Any] ,						__UpperCAmelCase  :		Union[str, Any] ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		List[Any] ,						__UpperCAmelCase  :		str ,						)		->       List[str]:
         """simple docstring"""



         UpperCAmelCase__			=						self.num_labels
         UpperCAmelCase__			=						FunnelForSequenceClassification(__UpperCAmelCase	)
         model.to(__UpperCAmelCase	)
         model.eval()
         UpperCAmelCase__			=						model(__UpperCAmelCase ,						attention_mask=__UpperCAmelCase ,						token_type_ids=__UpperCAmelCase ,						labels=__UpperCAmelCase	)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.num_labels)	)


   def     lowercase_   (self  :		Union[str, Any] ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		List[str] ,						__UpperCAmelCase  :		List[str] ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		List[Any] ,						)		->       Optional[int]:
         """simple docstring"""



         UpperCAmelCase__			=						self.num_choices
         UpperCAmelCase__			=						FunnelForMultipleChoice(config=__UpperCAmelCase	)
         model.to(__UpperCAmelCase	)
         model.eval()
         UpperCAmelCase__			=						input_ids.unsqueeze(1	).expand(-1 ,						self.num_choices ,						-1	).contiguous()
         UpperCAmelCase__			=						token_type_ids.unsqueeze(1	).expand(-1 ,						self.num_choices ,						-1	).contiguous()
         UpperCAmelCase__			=						input_mask.unsqueeze(1	).expand(-1 ,						self.num_choices ,						-1	).contiguous()
         UpperCAmelCase__			=						model(
             __UpperCAmelCase ,						attention_mask=__UpperCAmelCase ,						token_type_ids=__UpperCAmelCase ,						labels=__UpperCAmelCase ,						)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.num_choices)	)


   def     lowercase_   (self  :		Any ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		List[Any] ,						__UpperCAmelCase  :		Union[str, Any] ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		List[Any] ,						__UpperCAmelCase  :		Union[str, Any] ,						__UpperCAmelCase  :		int ,						__UpperCAmelCase  :		List[str] ,						)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						self.num_labels
         UpperCAmelCase__			=						FunnelForTokenClassification(config=__UpperCAmelCase	)
         model.to(__UpperCAmelCase	)
         model.eval()
         UpperCAmelCase__			=						model(__UpperCAmelCase ,						attention_mask=__UpperCAmelCase ,						token_type_ids=__UpperCAmelCase ,						labels=__UpperCAmelCase	)
         self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels)	)


   def     lowercase_   (self  :		Any ,						__UpperCAmelCase  :		List[str] ,						__UpperCAmelCase  :		str ,						__UpperCAmelCase  :		Dict ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		Optional[Any] ,						__UpperCAmelCase  :		List[Any] ,						__UpperCAmelCase  :		str ,						__UpperCAmelCase  :		Any ,						)		->       Tuple:
         """simple docstring"""



         UpperCAmelCase__			=						FunnelForQuestionAnswering(config=__UpperCAmelCase	)
         model.to(__UpperCAmelCase	)
         model.eval()
         UpperCAmelCase__			=						model(
             __UpperCAmelCase ,						attention_mask=__UpperCAmelCase ,						token_type_ids=__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     lowercase_   (self  :		Union[str, Any]	)		->       Any:
         """simple docstring"""



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





@require_torch
class   A							(       UpperCAmelCase_	,       UpperCAmelCase_	,       unittest.TestCase  ):
   __UpperCAmelCase				:				str          =    False
   __UpperCAmelCase				:				Optional[int]          =    False
   __UpperCAmelCase				:				Any          =    (
       (
           FunnelModel,
           FunnelForMaskedLM,
           FunnelForPreTraining,
           FunnelForQuestionAnswering,
           FunnelForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   __UpperCAmelCase				:				int          =    (
       {
           '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  :		Optional[Any] ,						__UpperCAmelCase  :		int ,						__UpperCAmelCase  :		Tuple ,						__UpperCAmelCase  :		List[str]=False	)		->       Any:
         """simple docstring"""



         UpperCAmelCase__			=						super()._prepare_for_class(__UpperCAmelCase ,						__UpperCAmelCase ,						return_labels=__UpperCAmelCase	)

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


   def     lowercase_   (self  :		Optional[Any]	)		->       int:
         """simple docstring"""



         UpperCAmelCase__			=						FunnelModelTester(self	)
         UpperCAmelCase__			=						ConfigTester(self ,						config_class=__UpperCAmelCase	)


   def     lowercase_   (self  :		Union[str, Any]	)		->       Optional[Any]:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def     lowercase_   (self  :		int	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__UpperCAmelCase	)


   def     lowercase_   (self  :		int	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_pretraining(*__UpperCAmelCase	)


   def     lowercase_   (self  :		Any	)		->       List[Any]:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*__UpperCAmelCase	)


   def     lowercase_   (self  :		Optional[Any]	)		->       Optional[int]:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*__UpperCAmelCase	)


   def     lowercase_   (self  :		Any	)		->       str:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*__UpperCAmelCase	)





   def     lowercase_   (self  :		int ,						__UpperCAmelCase  :		Dict	)		->       Dict:
         """simple docstring"""



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





@require_torch
class   A							(       UpperCAmelCase_	,       unittest.TestCase  ):
   __UpperCAmelCase				:				Optional[int]          =    False
   __UpperCAmelCase				:				str          =    False
   __UpperCAmelCase				:				Optional[Any]          =    (
       (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
   )


   def     lowercase_   (self  :		List[Any]	)		->       Optional[int]:
         """simple docstring"""



         UpperCAmelCase__			=						FunnelModelTester(self ,						base=__UpperCAmelCase	)
         UpperCAmelCase__			=						ConfigTester(self ,						config_class=__UpperCAmelCase	)


   def     lowercase_   (self  :		Tuple	)		->       List[Any]:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def     lowercase_   (self  :		List[Any]	)		->       List[Any]:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_base_model(*__UpperCAmelCase	)


   def     lowercase_   (self  :		Any	)		->       Optional[int]:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*__UpperCAmelCase	)


   def     lowercase_   (self  :		Union[str, Any]	)		->       Dict:
         """simple docstring"""



         UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_multiple_choice(*__UpperCAmelCase	)


   def     lowercase_   (self  :		str	)		->       int:
         """simple docstring"""



         UpperCAmelCase__	,   UpperCAmelCase__			=						self.model_tester.prepare_config_and_inputs_for_common()
         UpperCAmelCase__			=						True

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





   def     lowercase_   (self  :		List[Any] ,						__UpperCAmelCase  :		Optional[Any]	)		->       Union[str, Any]:
         """simple docstring"""



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





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


   def     lowercase_   (self  :		str	)		->       str:
         """simple docstring"""



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

         UpperCAmelCase__			=						FunnelModel.from_pretrained("sgugger/funnel-random-tiny"	)
         UpperCAmelCase__			=						model(__UpperCAmelCase ,						token_type_ids=__UpperCAmelCase	)[0].abs()

         UpperCAmelCase__			=						torch.tensor(2344.8352	)
         UpperCAmelCase__			=						torch.tensor(0.8052	)
         self.assertTrue(torch.allclose(output.sum() ,						__UpperCAmelCase ,						atol=1E-4	)	)
         self.assertTrue(torch.allclose(output.mean() ,						__UpperCAmelCase ,						atol=1E-4	)	)

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

         UpperCAmelCase__			=						torch.tensor(2343.8425	)
         UpperCAmelCase__			=						torch.tensor(0.8049	)
         self.assertTrue(torch.allclose(output.sum() ,						__UpperCAmelCase ,						atol=1E-4	)	)
         self.assertTrue(torch.allclose(output.mean() ,						__UpperCAmelCase ,						atol=1E-4	)	)





   @slow
   def     lowercase_   (self  :		int	)		->       Union[str, Any]:
         """simple docstring"""



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

         UpperCAmelCase__			=						torch.tensor(235.7246	)
         UpperCAmelCase__			=						torch.tensor(0.0256	)
         self.assertTrue(torch.allclose(output.sum() ,						__UpperCAmelCase ,						atol=1E-4	)	)
         self.assertTrue(torch.allclose(output.mean() ,						__UpperCAmelCase ,						atol=1E-4	)	)




import os
import tempfile
from pathlib import Path
from unittest import TestCase

import pyarrow as pa
import pytest

from datasets.arrow_dataset import Dataset
from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
from datasets.info import DatasetInfo
from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo

from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases





class   A							(       UpperCAmelCase_  ):


   def     lowercase_   (self  :		Dict ,						__UpperCAmelCase  :		Optional[int] ,						__UpperCAmelCase  :		List[str]=False	)		->       Union[str, Any]:
         """simple docstring"""



         UpperCAmelCase__	,   UpperCAmelCase__	,   UpperCAmelCase__			=						(
             filename_skip_take["filename"],
             filename_skip_take["skip"] if "skip" in filename_skip_take else None,
             filename_skip_take["take"] if "take" in filename_skip_take else None,
         )
         open(os.path.join(__UpperCAmelCase	) ,						"wb"	).close()
         UpperCAmelCase__			=						pa.Table.from_pydict({"filename": [Path(__UpperCAmelCase	).name] * 1_0_0}	)
         if take == -1:
               UpperCAmelCase__			=						len(__UpperCAmelCase	) - skip
         if skip is not None and take is not None:
               UpperCAmelCase__			=						pa_table.slice(__UpperCAmelCase ,						__UpperCAmelCase	)
         return pa_table





class   A							(       UpperCAmelCase_  ):


   def     lowercase_   (self  :		Tuple	)		->       Tuple:
         """simple docstring"""



         UpperCAmelCase__			=						"my_name"
         UpperCAmelCase__			=						SplitInfo(name="train" ,						num_examples=1_0_0	)
         UpperCAmelCase__			=						SplitInfo(name="test" ,						num_examples=1_0_0	)
         UpperCAmelCase__			=						[train_info, test_info]
         UpperCAmelCase__			=						SplitDict()
         split_dict.add(__UpperCAmelCase	)
         split_dict.add(__UpperCAmelCase	)
         UpperCAmelCase__			=						DatasetInfo(splits=__UpperCAmelCase	)

         with tempfile.TemporaryDirectory() as tmp_dir:
               UpperCAmelCase__			=						ReaderTest(__UpperCAmelCase ,						__UpperCAmelCase	)

               UpperCAmelCase__			=						"test[:33%]"
               UpperCAmelCase__			=						Dataset(**reader.read(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase	)	)
               self.assertEqual(dset["filename"][0] ,						f"""{name}-test"""	)
               self.assertEqual(dset.num_rows ,						3_3	)
               self.assertEqual(dset.num_columns ,						1	)

               UpperCAmelCase__			=						["train", "test[:33%]"]
               UpperCAmelCase__			=						[Split.TRAIN, ReadInstruction.from_spec("test[:33%]"	)]
               for instructions in [instructionsa, instructionsa]:
                     UpperCAmelCase__			=						[reader.read(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase	) for instr in instructions]
                     UpperCAmelCase__	,   UpperCAmelCase__			=						(Dataset(**__UpperCAmelCase	) for dataset_kwargs in datasets_kwargs)
                     self.assertEqual(train_dset["filename"][0] ,						f"""{name}-train"""	)
                     self.assertEqual(train_dset.num_rows ,						1_0_0	)
                     self.assertEqual(train_dset.num_columns ,						1	)
                     self.assertIsInstance(train_dset.split ,						__UpperCAmelCase	)
                     self.assertEqual(str(train_dset.split	) ,						"train"	)
                     self.assertEqual(test_dset["filename"][0] ,						f"""{name}-test"""	)
                     self.assertEqual(test_dset.num_rows ,						3_3	)
                     self.assertEqual(test_dset.num_columns ,						1	)
                     self.assertIsInstance(test_dset.split ,						__UpperCAmelCase	)
                     self.assertEqual(str(test_dset.split	) ,						"test[:33%]"	)
                     del train_dset, test_dset


   def     lowercase_   (self  :		Any	)		->       Any:
         """simple docstring"""



         UpperCAmelCase__			=						"my_name"
         UpperCAmelCase__			=						SplitInfo(name="train" ,						num_examples=1_0_0_0 ,						shard_lengths=[1_0_0] * 1_0	)
         UpperCAmelCase__			=						[train_info]
         UpperCAmelCase__			=						SplitDict()
         split_dict.add(__UpperCAmelCase	)
         UpperCAmelCase__			=						DatasetInfo(splits=__UpperCAmelCase	)

         with tempfile.TemporaryDirectory() as tmp_dir:
               UpperCAmelCase__			=						ReaderTest(__UpperCAmelCase ,						__UpperCAmelCase	)

               UpperCAmelCase__			=						"train[:33%]"
               UpperCAmelCase__			=						Dataset(**reader.read(__UpperCAmelCase ,						__UpperCAmelCase ,						__UpperCAmelCase	)	)
               self.assertEqual(dset["filename"][0] ,						f"""{name}-train-00000-of-00010"""	)
               self.assertEqual(dset["filename"][-1] ,						f"""{name}-train-00003-of-00010"""	)
               self.assertEqual(dset.num_rows ,						3_3_0	)
               self.assertEqual(dset.num_columns ,						1	)





   def     lowercase_   (self  :		Dict	)		->       Dict:
         """simple docstring"""



         UpperCAmelCase__			=						SplitInfo(name="train" ,						num_examples=1_0_0	)
         UpperCAmelCase__			=						SplitInfo(name="test" ,						num_examples=1_0_0	)
         UpperCAmelCase__			=						SplitDict()
         split_dict.add(__UpperCAmelCase	)
         split_dict.add(__UpperCAmelCase	)
         UpperCAmelCase__			=						DatasetInfo(splits=__UpperCAmelCase	)

         with tempfile.TemporaryDirectory() as tmp_dir:
               UpperCAmelCase__			=						ReaderTest(__UpperCAmelCase ,						__UpperCAmelCase	)

               UpperCAmelCase__			=						[
                   {"filename": os.path.join(__UpperCAmelCase ,						"train"	)},
                   {"filename": os.path.join(__UpperCAmelCase ,						"test"	), "skip": 1_0, "take": 1_0},
               ]
               UpperCAmelCase__			=						Dataset(**reader.read_files(__UpperCAmelCase ,						original_instructions="train+test[10:20]"	)	)
               self.assertEqual(dset.num_rows ,						1_1_0	)
               self.assertEqual(dset.num_columns ,						1	)
               del dset




@pytest.mark.parametrize("in_memory",						[False, True]							)
def      lowerCAmelCase_				(  __A,						__A,						__A							)  ->       List[Any]:







      '''simple docstring'''


      UpperCAmelCase__			=						arrow_file
      with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
            UpperCAmelCase__			=						ArrowReader.read_table(__A,						in_memory=__A							)
      assert table.shape == dataset.data.shape
      assert set(table.column_names							) == set(dataset.data.column_names							)
      assert dict(table.to_pydict()							) == dict(dataset.data.to_pydict()							)  # to_pydict returns OrderedDict




@pytest.mark.parametrize("in_memory",						[False, True]							)
def      lowerCAmelCase_				(  __A,						__A,						__A							)  ->       Tuple:







      '''simple docstring'''


      UpperCAmelCase__			=						arrow_file
      UpperCAmelCase__			=						ArrowReader("",						__A							)
      with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
            UpperCAmelCase__			=						reader.read_files([{"filename": filename}],						in_memory=__A							)
      assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
      UpperCAmelCase__			=						dataset_kwargs["arrow_table"]
      assert table.shape == dataset.data.shape
      assert set(table.column_names							) == set(dataset.data.column_names							)
      assert dict(table.to_pydict()							) == dict(dataset.data.to_pydict()							)  # to_pydict returns OrderedDict




def      lowerCAmelCase_				(  )  ->       Any:







      '''simple docstring'''


      assert ReadInstruction("train",						to=10,						unit="abs"							).to_spec() == "train[:10]"
      assert ReadInstruction("train",						from_=-80,						to=10,						unit="%"							).to_spec() == "train[-80%:10%]"

      UpperCAmelCase__			=						"train+test"
      assert ReadInstruction.from_spec(__A							).to_spec() == spec_train_test

      UpperCAmelCase__			=						"train[2:10]"
      assert ReadInstruction.from_spec(__A							).to_spec() == spec_train_abs

      UpperCAmelCase__			=						"train[15%:-20%]"
      assert ReadInstruction.from_spec(__A							).to_spec() == spec_train_pct

      UpperCAmelCase__			=						"train[:10%](closest)"
      assert ReadInstruction.from_spec(__A							).to_spec() == "train[:10%]"

      UpperCAmelCase__			=						"train[:10%](pct1_dropremainder)"
      assert ReadInstruction.from_spec(__A							).to_spec() == spec_train_pct_rounding

      UpperCAmelCase__			=						"train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)"
      assert ReadInstruction.from_spec(__A							).to_spec() == spec_train_test_pct_rounding




def      lowerCAmelCase_				(  )  ->       int:







      '''simple docstring'''


      UpperCAmelCase__			=						"dummy"
      UpperCAmelCase__			=						[SplitInfo(name="train",						num_examples=100							)]
      UpperCAmelCase__			=						"train[:33%]"
      UpperCAmelCase__			=						"arrow"
      UpperCAmelCase__			=						"prefix"

      UpperCAmelCase__			=						make_file_instructions(__A,						__A,						__A,						__A,						__A							)
      assert isinstance(__A,						__A							)
      assert file_instructions.num_examples == 33
      assert file_instructions.file_instructions == [
          {"filename": os.path.join(__A,						f"""{name}-train.arrow"""							), "skip": 0, "take": 33}
      ]

      UpperCAmelCase__			=						[SplitInfo(name="train",						num_examples=100,						shard_lengths=[10] * 10							)]
      UpperCAmelCase__			=						make_file_instructions(__A,						__A,						__A,						__A,						__A							)
      assert isinstance(__A,						__A							)
      assert file_instructions.num_examples == 33
      assert file_instructions.file_instructions == [
          {"filename": os.path.join(__A,						f"""{name}-train-00000-of-00010.arrow"""							), "skip": 0, "take": -1},
          {"filename": os.path.join(__A,						f"""{name}-train-00001-of-00010.arrow"""							), "skip": 0, "take": -1},
          {"filename": os.path.join(__A,						f"""{name}-train-00002-of-00010.arrow"""							), "skip": 0, "take": -1},
          {"filename": os.path.join(__A,						f"""{name}-train-00003-of-00010.arrow"""							), "skip": 0, "take": 3},
      ]







@pytest.mark.parametrize("name, expected_exception",						[(None, TypeError), ("", ValueError)]							)
def      lowerCAmelCase_				(  __A,						__A							)  ->       Optional[Any]:







      '''simple docstring'''


      UpperCAmelCase__			=						[SplitInfo(name="train",						num_examples=100							)]
      UpperCAmelCase__			=						"train"
      UpperCAmelCase__			=						"arrow"
      UpperCAmelCase__			=						"prefix_path"
      with pytest.raises(__A							):
            UpperCAmelCase__			=						make_file_instructions(__A,						__A,						__A,						__A,						__A							)



