






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        SCREAMING_SNAKE_CASE	:
     '''simple docstring'''


     def __init__( self				:		Union[str, Any]    ,       UpperCAmelCase_				:		Optional[Any]    ,       UpperCAmelCase_				:		Optional[int]=13    ,       UpperCAmelCase_				:		str=7    ,       UpperCAmelCase_				:		List[Any]=True    ,       UpperCAmelCase_				:		Any=True    ,       UpperCAmelCase_				:		List[str]=True    ,       UpperCAmelCase_				:		int=True    ,       UpperCAmelCase_				:		List[str]=99    ,       UpperCAmelCase_				:		Optional[Any]=[1, 1, 2]    ,       UpperCAmelCase_				:		int=1    ,       UpperCAmelCase_				:		Optional[Any]=32    ,       UpperCAmelCase_				:		List[str]=4    ,       UpperCAmelCase_				:		List[str]=8    ,       UpperCAmelCase_				:		Optional[Any]=37    ,       UpperCAmelCase_				:		List[str]="gelu_new"    ,       UpperCAmelCase_				:		Optional[int]=0.1    ,       UpperCAmelCase_				:		str=0.1    ,       UpperCAmelCase_				:		Optional[Any]=0.0    ,       UpperCAmelCase_				:		Any=512    ,       UpperCAmelCase_				:		List[str]=3    ,       UpperCAmelCase_				:		Optional[Any]=0.02    ,       UpperCAmelCase_				:		Dict=3    ,       UpperCAmelCase_				:		Optional[int]=4    ,       UpperCAmelCase_				:		Union[str, Any]=None    ,       UpperCAmelCase_				:		Optional[Any]=False    ,       ):
       SCREAMING_SNAKE_CASE				:					int		  =			parent
       SCREAMING_SNAKE_CASE				:					Tuple		  =			batch_size
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			seq_length
       SCREAMING_SNAKE_CASE				:					List[str]		  =			is_training
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			use_input_mask
       SCREAMING_SNAKE_CASE				:					Dict		  =			use_token_type_ids
       SCREAMING_SNAKE_CASE				:					Dict		  =			use_labels
       SCREAMING_SNAKE_CASE				:					List[str]		  =			vocab_size
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			block_sizes
       SCREAMING_SNAKE_CASE				:					Dict		  =			num_decoder_layers
       SCREAMING_SNAKE_CASE				:					Dict		  =			d_model
       SCREAMING_SNAKE_CASE				:					Any		  =			n_head
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			d_head
       SCREAMING_SNAKE_CASE				:					List[str]		  =			d_inner
       SCREAMING_SNAKE_CASE				:					List[str]		  =			hidden_act
       SCREAMING_SNAKE_CASE				:					int		  =			hidden_dropout
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			attention_dropout
       SCREAMING_SNAKE_CASE				:					str		  =			activation_dropout
       SCREAMING_SNAKE_CASE				:					Tuple		  =			max_position_embeddings
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			type_vocab_size
       SCREAMING_SNAKE_CASE				:					str		  =			2
       SCREAMING_SNAKE_CASE				:					List[str]		  =			num_labels
       SCREAMING_SNAKE_CASE				:					str		  =			num_choices
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			scope
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			initializer_std

       # Used in the tests to check the size of the first attention layer
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			n_head
       # Used in the tests to check the size of the first hidden state
       SCREAMING_SNAKE_CASE				:					List[str]		  =			self.d_model
       # Used in the tests to check the number of output hidden states/attentions
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			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:
         SCREAMING_SNAKE_CASE				:					int		  =			self.num_hidden_layers + 2


     def  _A				( self				:		Optional[Any]					):
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			ids_tensor([self.batch_size, self.seq_length]    ,       self.vocab_size					)

       SCREAMING_SNAKE_CASE				:					Any		  =			None
       if self.use_input_mask:
         SCREAMING_SNAKE_CASE				:					Optional[int]		  =			ids_tensor([self.batch_size, self.seq_length]    ,       vocab_size=2					)

       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			None
       if self.use_token_type_ids:
         SCREAMING_SNAKE_CASE				:					str		  =			ids_tensor([self.batch_size, self.seq_length]    ,       self.type_vocab_size					)

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			None
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			None
       SCREAMING_SNAKE_CASE				:					int		  =			None
       if self.use_labels:
         SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			ids_tensor([self.batch_size]    ,       self.type_sequence_label_size					)
         SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			ids_tensor([self.batch_size, self.seq_length]    ,       self.num_labels					)
         SCREAMING_SNAKE_CASE				:					List[Any]		  =			ids_tensor([self.batch_size]    ,       self.num_choices					)
         SCREAMING_SNAKE_CASE				:					List[Any]		  =			ids_tensor([self.batch_size, self.seq_length]    ,       1					)

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

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


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

       SCREAMING_SNAKE_CASE				:					List[Any]		  =			False
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			model(UpperCAmelCase_					)
       self.parent.assertEqual(result.last_hidden_state.shape    ,       (self.batch_size, self.seq_length, self.d_model)					)

       SCREAMING_SNAKE_CASE				:					str		  =			False
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			model(UpperCAmelCase_					)
       self.parent.assertEqual(result.last_hidden_state.shape    ,       (self.batch_size, self.seq_length, self.d_model)					)


     def  _A				( self				:		Tuple    ,       UpperCAmelCase_				:		Tuple    ,       UpperCAmelCase_				:		Dict    ,       UpperCAmelCase_				:		Any    ,       UpperCAmelCase_				:		Tuple    ,       UpperCAmelCase_				:		Dict    ,       UpperCAmelCase_				:		int    ,       UpperCAmelCase_				:		Optional[Any]    ,       UpperCAmelCase_				:		List[Any]    ,       ):
       SCREAMING_SNAKE_CASE				:					int		  =			FunnelBaseModel(config=UpperCAmelCase_					)
       model.to(UpperCAmelCase_					)
       model.eval()
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			model(UpperCAmelCase_    ,       attention_mask=UpperCAmelCase_    ,       token_type_ids=UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					Tuple		  =			model(UpperCAmelCase_    ,       token_type_ids=UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					Any		  =			model(UpperCAmelCase_					)
       self.parent.assertEqual(result.last_hidden_state.shape    ,       (self.batch_size, 2, self.d_model)					)

       SCREAMING_SNAKE_CASE				:					List[str]		  =			False
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			model(UpperCAmelCase_					)
       self.parent.assertEqual(result.last_hidden_state.shape    ,       (self.batch_size, 3, self.d_model)					)

       SCREAMING_SNAKE_CASE				:					List[str]		  =			False
       SCREAMING_SNAKE_CASE				:					Any		  =			model(UpperCAmelCase_					)
       self.parent.assertEqual(result.last_hidden_state.shape    ,       (self.batch_size, 2, self.d_model)					)


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


     def  _A				( self				:		str    ,       UpperCAmelCase_				:		Dict    ,       UpperCAmelCase_				:		Union[str, Any]    ,       UpperCAmelCase_				:		Optional[Any]    ,       UpperCAmelCase_				:		Dict    ,       UpperCAmelCase_				:		Optional[Any]    ,       UpperCAmelCase_				:		Any    ,       UpperCAmelCase_				:		Optional[int]    ,       UpperCAmelCase_				:		Union[str, Any]    ,       ):
       SCREAMING_SNAKE_CASE				:					Any		  =			FunnelForMaskedLM(config=UpperCAmelCase_					)
       model.to(UpperCAmelCase_					)
       model.eval()
       SCREAMING_SNAKE_CASE				:					Dict		  =			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  _A				( self				:		List[Any]    ,       UpperCAmelCase_				:		Any    ,       UpperCAmelCase_				:		str    ,       UpperCAmelCase_				:		Optional[int]    ,       UpperCAmelCase_				:		int    ,       UpperCAmelCase_				:		Dict    ,       UpperCAmelCase_				:		List[Any]    ,       UpperCAmelCase_				:		Dict    ,       UpperCAmelCase_				:		List[Any]    ,       ):
       SCREAMING_SNAKE_CASE				:					Any		  =			self.num_labels
       SCREAMING_SNAKE_CASE				:					Any		  =			FunnelForSequenceClassification(UpperCAmelCase_					)
       model.to(UpperCAmelCase_					)
       model.eval()
       SCREAMING_SNAKE_CASE				:					int		  =			model(UpperCAmelCase_    ,       attention_mask=UpperCAmelCase_    ,       token_type_ids=UpperCAmelCase_    ,       labels=UpperCAmelCase_					)
       self.parent.assertEqual(result.logits.shape    ,       (self.batch_size, self.num_labels)					)


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


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





@require_torch
class        SCREAMING_SNAKE_CASE	(	lowerCAmelCase					,						lowerCAmelCase					,						unittest.TestCase    ):
     '''simple docstring'''


     UpperCamelCase_							:						Tuple								=		False
     UpperCamelCase_							:						Any								=		False
     UpperCamelCase_							:						int								=		(
         (
             FunnelModel,
             FunnelForMaskedLM,
             FunnelForPreTraining,
             FunnelForQuestionAnswering,
             FunnelForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     UpperCamelCase_							:						str								=		(
         {
             '''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  _A				( self				:		int    ,       UpperCAmelCase_				:		str    ,       UpperCAmelCase_				:		List[Any]    ,       UpperCAmelCase_				:		List[Any]=False					):
       SCREAMING_SNAKE_CASE				:					str		  =			super()._prepare_for_class(UpperCAmelCase_    ,       UpperCAmelCase_    ,       return_labels=UpperCAmelCase_					)

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


     def  _A				( self				:		int					):
       SCREAMING_SNAKE_CASE				:					str		  =			FunnelModelTester(self					)
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			ConfigTester(self    ,       config_class=UpperCAmelCase_					)


     def  _A				( self				:		Dict					):
       self.config_tester.run_common_tests()


     def  _A				( self				:		Union[str, Any]					):
       SCREAMING_SNAKE_CASE				:					Dict		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*UpperCAmelCase_					)


     def  _A				( self				:		List[str]					):
       SCREAMING_SNAKE_CASE				:					str		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_pretraining(*UpperCAmelCase_					)


     def  _A				( self				:		Tuple					):
       SCREAMING_SNAKE_CASE				:					Tuple		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase_					)


     def  _A				( self				:		List[str]					):
       SCREAMING_SNAKE_CASE				:					str		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase_					)


     def  _A				( self				:		Any					):
       SCREAMING_SNAKE_CASE				:					str		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*UpperCAmelCase_					)






     def  _A				( self				:		Optional[Any]    ,       UpperCAmelCase_				:		Optional[int]					):
       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:
           SCREAMING_SNAKE_CASE				:					List[str]		  =			getattr(UpperCAmelCase_    ,       UpperCAmelCase_					)
           weight.data.fill_(3					)





@require_torch
class        SCREAMING_SNAKE_CASE	(	lowerCAmelCase					,						unittest.TestCase    ):
     '''simple docstring'''


     UpperCamelCase_							:						List[Any]								=		False
     UpperCamelCase_							:						Tuple								=		False
     UpperCamelCase_							:						List[Any]								=		(
         (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
     )


     def  _A				( self				:		List[str]					):
       SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			FunnelModelTester(self    ,       base=UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					int		  =			ConfigTester(self    ,       config_class=UpperCAmelCase_					)


     def  _A				( self				:		Optional[int]					):
       self.config_tester.run_common_tests()


     def  _A				( self				:		str					):
       SCREAMING_SNAKE_CASE				:					List[Any]		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_base_model(*UpperCAmelCase_					)


     def  _A				( self				:		List[str]					):
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*UpperCAmelCase_					)


     def  _A				( self				:		Tuple					):
       SCREAMING_SNAKE_CASE				:					int		  =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*UpperCAmelCase_					)


     def  _A				( self				:		Optional[Any]					):
       SCREAMING_SNAKE_CASE  ,     SCREAMING_SNAKE_CASE				:					List[Any]		  =			self.model_tester.prepare_config_and_inputs_for_common()
       SCREAMING_SNAKE_CASE				:					Tuple		  =			True

       for model_class in self.all_model_classes:
         if model_class.__name__ == "FunnelBaseModel":
           continue
         SCREAMING_SNAKE_CASE				:					Dict		  =			model_class(UpperCAmelCase_					)
         model.to(UpperCAmelCase_					)
         model.train()
         SCREAMING_SNAKE_CASE				:					List[Any]		  =			self._prepare_for_class(UpperCAmelCase_    ,       UpperCAmelCase_    ,       return_labels=UpperCAmelCase_					)
         SCREAMING_SNAKE_CASE				:					Dict		  =			model(**UpperCAmelCase_					).loss
         loss.backward()






     def  _A				( self				:		Optional[int]    ,       UpperCAmelCase_				:		int					):
       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:
           SCREAMING_SNAKE_CASE				:					Tuple		  =			getattr(UpperCAmelCase_    ,       UpperCAmelCase_					)
           weight.data.fill_(3					)





@require_torch
@require_sentencepiece
@require_tokenizers
class        SCREAMING_SNAKE_CASE	(	unittest.TestCase    ):
     '''simple docstring'''


     def  _A				( self				:		List[Any]					):
       SCREAMING_SNAKE_CASE				:					int		  =			13
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			7
       SCREAMING_SNAKE_CASE				:					List[str]		  =			torch.arange(0    ,       batch_size * sequence_length					).long().reshape(UpperCAmelCase_    ,       UpperCAmelCase_					)
       SCREAMING_SNAKE_CASE				:					Dict		  =			[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
       SCREAMING_SNAKE_CASE				:					Dict		  =			torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]					)

       SCREAMING_SNAKE_CASE				:					Tuple		  =			FunnelModel.from_pretrained("sgugger/funnel-random-tiny"					)
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			model(UpperCAmelCase_    ,       token_type_ids=UpperCAmelCase_					)[0].abs()

       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			torch.tensor(2_344.8_352					)
       SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			torch.tensor(0.8_052					)
       self.assertTrue(torch.allclose(output.sum()    ,       UpperCAmelCase_    ,       atol=1E-4					)					)
       self.assertTrue(torch.allclose(output.mean()    ,       UpperCAmelCase_    ,       atol=1E-4					)					)

       SCREAMING_SNAKE_CASE				:					Dict		  =			torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]					)
       SCREAMING_SNAKE_CASE				:					int		  =			model(UpperCAmelCase_    ,       attention_mask=UpperCAmelCase_    ,       token_type_ids=UpperCAmelCase_					)[0].abs()

       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			torch.tensor(2_343.8_425					)
       SCREAMING_SNAKE_CASE				:					List[str]		  =			torch.tensor(0.8_049					)
       self.assertTrue(torch.allclose(output.sum()    ,       UpperCAmelCase_    ,       atol=1E-4					)					)
       self.assertTrue(torch.allclose(output.mean()    ,       UpperCAmelCase_    ,       atol=1E-4					)					)






     @slow
     def  _A				( self				:		Union[str, Any]					):
       SCREAMING_SNAKE_CASE				:					Optional[int]		  =			FunnelTokenizer.from_pretrained("huggingface/funnel-small"					)
       SCREAMING_SNAKE_CASE				:					str		  =			FunnelModel.from_pretrained("huggingface/funnel-small"					)
       SCREAMING_SNAKE_CASE				:					Any		  =			tokenizer("Hello! I am the Funnel Transformer model."    ,       return_tensors="pt"					)
       SCREAMING_SNAKE_CASE				:					Tuple		  =			model(**UpperCAmelCase_					)[0]

       SCREAMING_SNAKE_CASE				:					Tuple		  =			torch.tensor(235.7_246					)
       SCREAMING_SNAKE_CASE				:					Dict		  =			torch.tensor(0.0_256					)
       self.assertTrue(torch.allclose(output.sum()    ,       UpperCAmelCase_    ,       atol=1E-4					)					)
       self.assertTrue(torch.allclose(output.mean()    ,       UpperCAmelCase_    ,       atol=1E-4					)					)














import argparse
import logging

import numpy as np
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer


logging.basicConfig(
    format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s""",
    datefmt="""%m/%d/%Y %H:%M:%S""",
    level=logging.INFO,
)
snake_case			 =						logging.getLogger(__name__)


def      lowerCamelCase__	(   lowercase ):

  """simple docstring"""






  np.random.seed(args.seed )
  torch.manual_seed(args.seed )
  if args.n_gpu > 0:
    torch.cuda.manual_seed_all(args.seed )


def      lowerCamelCase__	(   ):

  """simple docstring"""






  SCREAMING_SNAKE_CASE				:					List[str]		  =			argparse.ArgumentParser()
  parser.add_argument(
      "--model_name_or_path"  ,   default=lowercase  ,   type=lowercase  ,   required=lowercase  ,   )
  parser.add_argument("--prompt"  ,   type=lowercase  ,   default="" )
  parser.add_argument("--length"  ,   type=lowercase  ,   default=20 )
  parser.add_argument("--stop_token"  ,   type=lowercase  ,   default=lowercase  ,   help="Token at which text generation is stopped" )
  parser.add_argument(
      "--temperature"  ,   type=lowercase  ,   default=1.0  ,   help="temperature of 1.0 has no effect, lower tend toward greedy sampling"  ,   )
  parser.add_argument(
      "--repetition_penalty"  ,   type=lowercase  ,   default=1.0  ,   help="primarily useful for CTRL model; in that case, use 1.2" )
  parser.add_argument("--k"  ,   type=lowercase  ,   default=0 )
  parser.add_argument("--penalty_alpha"  ,   type=lowercase  ,   default=0.0 )
  parser.add_argument("--p"  ,   type=lowercase  ,   default=0.9 )

  parser.add_argument("--prefix"  ,   type=lowercase  ,   default=""  ,   help="Text added prior to input." )
  parser.add_argument("--padding_text"  ,   type=lowercase  ,   default=""  ,   help="Deprecated, the use of `--prefix` is preferred." )
  parser.add_argument("--xlm_language"  ,   type=lowercase  ,   default=""  ,   help="Optional language when used with the XLM model." )

  parser.add_argument("--seed"  ,   type=lowercase  ,   default=42  ,   help="random seed for initialization" )
  parser.add_argument("--no_cuda"  ,   action="store_true"  ,   help="Avoid using CUDA when available" )
  parser.add_argument(
      "--fp16"  ,   action="store_true"  ,   help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"  ,   )
  SCREAMING_SNAKE_CASE				:					str		  =			parser.parse_args()

  SCREAMING_SNAKE_CASE				:					Dict		  =			torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu" )
  SCREAMING_SNAKE_CASE				:					List[str]		  =			0 if args.no_cuda else torch.cuda.device_count()

  logger.warning(F'''device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fpaa}''' )

  set_seed(lowercase )

  # Initialize the model and tokenizer
  SCREAMING_SNAKE_CASE				:					Tuple		  =			AutoTokenizer.from_pretrained(args.model_name_or_path )
  SCREAMING_SNAKE_CASE				:					int		  =			AutoModelForCausalLM.from_pretrained(args.model_name_or_path )

  # tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path)
  # model = OPTForCausalLM.from_pretrained(args.model_name_or_path)
  model.to(args.device )

  if args.fpaa:
    model.half()

  logger.info(lowercase )
  SCREAMING_SNAKE_CASE				:					Any		  =			args.prompt if args.prompt else input("Model prompt >>> " )

  SCREAMING_SNAKE_CASE				:					str		  =			tokenizer(lowercase  ,   return_tensors="pt"  ,   add_special_tokens=lowercase )
  SCREAMING_SNAKE_CASE				:					Dict		  =			{key: value.to(args.device ) for key, value in inputs.items()}

  SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			model.generate(
      **lowercase  ,   max_length=args.length + len(inputs["input_ids"][0] )  ,   penalty_alpha=args.penalty_alpha  ,   top_k=args.k  ,   )

  SCREAMING_SNAKE_CASE				:					str		  =			[]
  for generated_sequence_idx, generated_sequence in enumerate(lowercase ):
    print(F'''=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===''' )
    SCREAMING_SNAKE_CASE				:					Union[str, Any]		  =			generated_sequence.tolist()

    # Decode text
    SCREAMING_SNAKE_CASE				:					Optional[Any]		  =			tokenizer.decode(lowercase  ,   clean_up_tokenization_spaces=lowercase  ,   add_special_tokens=lowercase )

    # Remove all text after the stop token
    SCREAMING_SNAKE_CASE				:					List[str]		  =			text[: text.find(args.stop_token ) if args.stop_token else None]

    # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
    SCREAMING_SNAKE_CASE				:					str		  =			(
        prompt_text + text[len(tokenizer.decode(inputs["input_ids"][0]  ,   clean_up_tokenization_spaces=lowercase ) ) :]
    )

    generated_sequences.append(lowercase )
    print(lowercase )

  return generated_sequences


if __name__ == "__main__":
   main()






