code: 
  
  
  
  '''simple docstring'''
  
  
  
  import unittest
  
  import numpy as np
  import timeout_decorator  # noqa
  
  from transformers import MBartConfig, is_flax_available
  from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
  from transformers.utils import cached_property
  
  from ...generation.test_flax_utils import FlaxGenerationTesterMixin
  from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor
  
  
  if is_flax_available():
   import os
  
   # The slow tests are often failing with OOM error on GPU
   # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
   # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
   a__						:			Dict										=			'platform'
  
   import jax
   import jax.numpy as jnp
  
   from transformers import AutoTokenizer
   from transformers.models.mbart.modeling_flax_mbart import (
       FlaxMBartForConditionalGeneration,
       FlaxMBartForQuestionAnswering,
       FlaxMBartForSequenceClassification,
       FlaxMBartModel,
       shift_tokens_right,
   )
  
  
  
  
  
  
  
  def _lowercase   (  __A							,__A							,__A=None							,__A=None							,__A=None							,__A=None							,__A=None							,__A=None							,):
  
  
  
   '''simple docstring'''
  
  
  
   if attention_mask is None:
    __UpperCamelCase				      =  np.where(input_ids != config.pad_token_id							,1							,0       )
   if decoder_attention_mask is None:
    __UpperCamelCase				      =  np.where(decoder_input_ids != config.pad_token_id							,1							,0       )
   if head_mask is None:
    __UpperCamelCase				      =  np.ones((config.encoder_layers, config.encoder_attention_heads)       )
   if decoder_head_mask is None:
    __UpperCamelCase				      =  np.ones((config.decoder_layers, config.decoder_attention_heads)       )
   if cross_attn_head_mask is None:
    __UpperCamelCase				      =  np.ones((config.decoder_layers, config.decoder_attention_heads)       )
   return {
       "input_ids": input_ids,
       "decoder_input_ids": decoder_input_ids,
       "attention_mask": attention_mask,
       "decoder_attention_mask": decoder_attention_mask,
   }
  class        UpperCAmelCase__  (    unittest.TestCase):
  
  
  
   def __init__(  self		,	lowercase		,	lowercase=1_3		,	lowercase=7		,	lowercase=True		,	lowercase=False		,	lowercase=9_9		,	lowercase=1_6		,	lowercase=2		,	lowercase=4		,	lowercase=4		,	lowercase="gelu"		,	lowercase=0.1		,	lowercase=0.1		,	lowercase=3_2		,	lowercase=2		,	lowercase=1		,	lowercase=0		,	lowercase=2		,	lowercase=0.02		,	)  ->							Optional[Any]:
    __UpperCamelCase				      =  parent
    __UpperCamelCase				      =  batch_size
    __UpperCamelCase				      =  seq_length
    __UpperCamelCase				      =  is_training
    __UpperCamelCase				      =  use_labels
    __UpperCamelCase				      =  vocab_size
    __UpperCamelCase				      =  hidden_size
    __UpperCamelCase				      =  num_hidden_layers
    __UpperCamelCase				      =  num_attention_heads
    __UpperCamelCase				      =  intermediate_size
    __UpperCamelCase				      =  hidden_act
    __UpperCamelCase				      =  hidden_dropout_prob
    __UpperCamelCase				      =  attention_probs_dropout_prob
    __UpperCamelCase				      =  max_position_embeddings
    __UpperCamelCase				      =  eos_token_id
    __UpperCamelCase				      =  pad_token_id
    __UpperCamelCase				      =  bos_token_id
    __UpperCamelCase				      =  decoder_start_token_id
    __UpperCamelCase				      =  initializer_range
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Optional[int]:
    __UpperCamelCase				      =  np.clip(ids_tensor([self.batch_size, self.seq_length - 1]		,	self.vocab_size			)		,	3		,	self.vocab_size			)
    __UpperCamelCase				      =  np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)		,	dtype=np.intaa			))		,	-1			)
  
    __UpperCamelCase				      =  shift_tokens_right(lowercase		,	1			)
  
    __UpperCamelCase				      =  MBartConfig(
        vocab_size=self.vocab_size		,	d_model=self.hidden_size		,	encoder_layers=self.num_hidden_layers		,	decoder_layers=self.num_hidden_layers		,	encoder_attention_heads=self.num_attention_heads		,	decoder_attention_heads=self.num_attention_heads		,	encoder_ffn_dim=self.intermediate_size		,	decoder_ffn_dim=self.intermediate_size		,	dropout=self.hidden_dropout_prob		,	attention_dropout=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	eos_token_id=self.eos_token_id		,	bos_token_id=self.bos_token_id		,	pad_token_id=self.pad_token_id		,	decoder_start_token_id=self.decoder_start_token_id		,	initializer_range=self.initializer_range		,	use_cache=lowercase		,	)
    __UpperCamelCase				      =  prepare_mbart_inputs_dict(lowercase		,	lowercase		,	lowercase			)
    return config, inputs_dict
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Any:
    __UpperCamelCase      ,		__UpperCamelCase				      =  self.prepare_config_and_inputs()
    return config, inputs_dict
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase			)  ->							List[str]:
    __UpperCamelCase				      =  2_0
    __UpperCamelCase				      =  model_class_name(lowercase			)
  
    __UpperCamelCase				      =  model.encode(inputs_dict["""input_ids"""]			)
  
    __UpperCamelCase      ,		__UpperCamelCase				      =  (
        inputs_dict["""decoder_input_ids"""],
        inputs_dict["""decoder_attention_mask"""],
    )
  
    __UpperCamelCase				      =  model.init_cache(decoder_input_ids.shape[0]		,	lowercase		,	lowercase			)
    __UpperCamelCase				      =  jnp.ones((decoder_input_ids.shape[0], max_decoder_length)		,	dtype="""i4"""			)
  
    __UpperCamelCase				      =  jnp.broadcast_to(
        jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :]		,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,	)
    __UpperCamelCase				      =  model.decode(
        decoder_input_ids[:, :-1]		,	lowercase		,	decoder_attention_mask=lowercase		,	past_key_values=lowercase		,	decoder_position_ids=lowercase		,	)
  
    __UpperCamelCase				      =  jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,	dtype="""i4"""			)
    __UpperCamelCase				      =  model.decode(
        decoder_input_ids[:, -1:]		,	lowercase		,	decoder_attention_mask=lowercase		,	past_key_values=outputs_cache.past_key_values		,	decoder_position_ids=lowercase		,	)
  
    __UpperCamelCase				      =  model.decode(lowercase		,	lowercase			)
  
    __UpperCamelCase				      =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])			)			)
    self.parent.assertTrue(diff < 1E-3		,	msg=f"Max diff is {diff}"			)
  
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase			)  ->							List[str]:
    __UpperCamelCase				      =  2_0
    __UpperCamelCase				      =  model_class_name(lowercase			)
  
    __UpperCamelCase				      =  model.encode(inputs_dict["""input_ids"""]			)
  
    __UpperCamelCase      ,		__UpperCamelCase				      =  (
        inputs_dict["""decoder_input_ids"""],
        inputs_dict["""decoder_attention_mask"""],
    )
  
    __UpperCamelCase				      =  jnp.concatenate(
        [
            decoder_attention_mask,
            jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])			),
        ]		,	axis=-1		,	)
  
    __UpperCamelCase				      =  model.init_cache(decoder_input_ids.shape[0]		,	lowercase		,	lowercase			)
    __UpperCamelCase				      =  jnp.broadcast_to(
        jnp.arange(decoder_input_ids.shape[-1] - 1			)[None, :]		,	(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)		,	)
  
    __UpperCamelCase				      =  model.decode(
        decoder_input_ids[:, :-1]		,	lowercase		,	decoder_attention_mask=lowercase		,	past_key_values=lowercase		,	decoder_position_ids=lowercase		,	)
    __UpperCamelCase				      =  jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]		,	dtype="""i4"""			)
    __UpperCamelCase				      =  model.decode(
        decoder_input_ids[:, -1:]		,	lowercase		,	past_key_values=outputs_cache.past_key_values		,	decoder_attention_mask=lowercase		,	decoder_position_ids=lowercase		,	)
  
    __UpperCamelCase				      =  model.decode(lowercase		,	lowercase		,	decoder_attention_mask=lowercase			)
  
    __UpperCamelCase				      =  np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])			)			)
    self.parent.assertTrue(diff < 1E-3		,	msg=f"Max diff is {diff}"			)
  
  
  @require_flax
  class        UpperCAmelCase__  (    unittest.TestCase):
   __SCREAMING_SNAKE_CASE					      =				9_9
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Tuple:
    __UpperCamelCase				      =  np.array(
        [
            [7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2],
            [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2],
            [5, 9_7, 1_7, 3_9, 9_4, 4_0, 2],
            [7_6, 8_3, 9_4, 2_5, 7_0, 7_8, 2],
            [8_7, 5_9, 4_1, 3_5, 4_8, 6_6, 2],
            [5_5, 1_3, 1_6, 5_8, 5, 2, 1],  # note padding
            [6_4, 2_7, 3_1, 5_1, 1_2, 7_5, 2],
            [5_2, 6_4, 8_6, 1_7, 8_3, 3_9, 2],
            [4_8, 6_1, 9, 2_4, 7_1, 8_2, 2],
            [2_6, 1, 6_0, 4_8, 2_2, 1_3, 2],
            [2_1, 5, 6_2, 2_8, 1_4, 7_6, 2],
            [4_5, 9_8, 3_7, 8_6, 5_9, 4_8, 2],
            [7_0, 7_0, 5_0, 9, 2_8, 0, 2],
        ]		,	dtype=np.intaa		,	)
  
    __UpperCamelCase				      =  input_ids.shape[0]
    __UpperCamelCase				      =  MBartConfig(
        vocab_size=self.vocab_size		,	d_model=2_4		,	encoder_layers=2		,	decoder_layers=2		,	encoder_attention_heads=2		,	decoder_attention_heads=2		,	encoder_ffn_dim=3_2		,	decoder_ffn_dim=3_2		,	max_position_embeddings=4_8		,	eos_token_id=2		,	pad_token_id=1		,	bos_token_id=0		,	)
    return config, input_ids, batch_size
  
  
  
   def    __lowerCamelCase				(  self			)  ->							List[str]:
    __UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase				      =  self._get_config_and_data()
    __UpperCamelCase				      =  FlaxMBartForSequenceClassification(lowercase			)
    __UpperCamelCase				      =  model(input_ids=lowercase		,	decoder_input_ids=lowercase			)
    __UpperCamelCase				      =  (batch_size, config.num_labels)
    self.assertEqual(outputs["""logits"""].shape		,	lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Optional[int]:
    __UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase				      =  self._get_config_and_data()
    __UpperCamelCase				      =  FlaxMBartForQuestionAnswering(lowercase			)
    __UpperCamelCase				      =  model(input_ids=lowercase			)
  
    self.assertEqual(outputs["""start_logits"""].shape		,	input_ids.shape			)
    self.assertEqual(outputs["""end_logits"""].shape		,	input_ids.shape			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Dict:
    __UpperCamelCase      ,		__UpperCamelCase      ,		__UpperCamelCase				      =  self._get_config_and_data()
    __UpperCamelCase				      =  FlaxMBartForConditionalGeneration(lowercase			)
    __UpperCamelCase				      =  lm_model(input_ids=lowercase			)
    __UpperCamelCase				      =  (batch_size, input_ids.shape[1], config.vocab_size)
    self.assertEqual(outputs["""logits"""].shape		,	lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							List[str]:
    __UpperCamelCase				      =  MBartConfig(
        vocab_size=self.vocab_size		,	d_model=1_4		,	encoder_layers=2		,	decoder_layers=2		,	encoder_attention_heads=2		,	decoder_attention_heads=2		,	encoder_ffn_dim=8		,	decoder_ffn_dim=8		,	max_position_embeddings=4_8		,	)
    __UpperCamelCase				      =  FlaxMBartForConditionalGeneration(lowercase			)
    __UpperCamelCase				      =  np.array([[7_1, 8_2, 1_8, 3_3, 4_6, 9_1, 2], [6_8, 3_4, 2_6, 5_8, 3_0, 2, 1]]		,	dtype=np.intaa			)
    __UpperCamelCase				      =  np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]]		,	dtype=np.intaa			)
    __UpperCamelCase				      =  lm_model(input_ids=lowercase		,	decoder_input_ids=lowercase			)
    __UpperCamelCase				      =  (*summary.shape, config.vocab_size)
    self.assertEqual(outputs["""logits"""].shape		,	lowercase			)
  
  
  
  
   def    __lowerCamelCase				(  self			)  ->							List[str]:
    __UpperCamelCase				      =  np.array([[7_1, 8_2, 1_8, 3_3, 2, 1, 1], [6_8, 3_4, 2_6, 5_8, 3_0, 8_2, 2]]		,	dtype=np.intaa			)
    __UpperCamelCase				      =  shift_tokens_right(lowercase		,	1			)
    __UpperCamelCase				      =  np.equal(lowercase		,	1			).astype(np.floataa			).sum()
    __UpperCamelCase				      =  np.equal(lowercase		,	1			).astype(np.floataa			).sum()
    self.assertEqual(shifted.shape		,	input_ids.shape			)
    self.assertEqual(lowercase		,	n_pad_before - 1			)
    self.assertTrue(np.equal(shifted[:, 0]		,	2			).all()			)
  
  
  @require_flax
  class        UpperCAmelCase__  (    UpperCAmelCase_							,			unittest.TestCase							,			UpperCAmelCase_):
   __SCREAMING_SNAKE_CASE					      =				True
   __SCREAMING_SNAKE_CASE					      =				(
       (
           FlaxMBartModel,
           FlaxMBartForConditionalGeneration,
           FlaxMBartForSequenceClassification,
           FlaxMBartForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )
   __SCREAMING_SNAKE_CASE					      =				(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Optional[int]:
    __UpperCamelCase				      =  FlaxMBartModelTester(self			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Dict:
    __UpperCamelCase      ,		__UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    for model_class in self.all_model_classes:
     self.model_tester.check_use_cache_forward(lowercase		,	lowercase		,	lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Optional[Any]:
    __UpperCamelCase      ,		__UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    for model_class in self.all_model_classes:
     self.model_tester.check_use_cache_forward_with_attn_mask(lowercase		,	lowercase		,	lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							List[Any]:
    __UpperCamelCase      ,		__UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs_for_common()
  
    for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__			):
      __UpperCamelCase				      =  self._prepare_for_class(lowercase		,	lowercase			)
      __UpperCamelCase				      =  model_class(lowercase			)
  
      @jax.jit
      def encode_jitted(lowercase		,	lowercase=None		,	**lowercase			):
       return model.encode(input_ids=lowercase		,	attention_mask=lowercase			)
  
      with self.subTest("""JIT Enabled"""			):
       __UpperCamelCase				      =  encode_jitted(**lowercase			).to_tuple()
  
      with self.subTest("""JIT Disabled"""			):
       with jax.disable_jit():
        __UpperCamelCase				      =  encode_jitted(**lowercase			).to_tuple()
  
      self.assertEqual(len(lowercase			)		,	len(lowercase			)			)
      for jitted_output, output in zip(lowercase		,	lowercase			):
       self.assertEqual(jitted_output.shape		,	output.shape			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							int:
    __UpperCamelCase      ,		__UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs_for_common()
  
    for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__			):
      __UpperCamelCase				      =  model_class(lowercase			)
      __UpperCamelCase				      =  model.encode(inputs_dict["""input_ids"""]		,	inputs_dict["""attention_mask"""]			)
  
      __UpperCamelCase				      =  {
          """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
          """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
          """encoder_outputs""": encoder_outputs,
      }
  
      @jax.jit
      def decode_jitted(lowercase		,	lowercase		,	lowercase			):
       return model.decode(
           decoder_input_ids=lowercase		,	decoder_attention_mask=lowercase		,	encoder_outputs=lowercase		,	)
  
      with self.subTest("""JIT Enabled"""			):
       __UpperCamelCase				      =  decode_jitted(**lowercase			).to_tuple()
  
      with self.subTest("""JIT Disabled"""			):
       with jax.disable_jit():
        __UpperCamelCase				      =  decode_jitted(**lowercase			).to_tuple()
  
      self.assertEqual(len(lowercase			)		,	len(lowercase			)			)
      for jitted_output, output in zip(lowercase		,	lowercase			):
       self.assertEqual(jitted_output.shape		,	output.shape			)
  
  
  
  
   @slow
   def    __lowerCamelCase				(  self			)  ->							str:
    for model_class_name in self.all_model_classes:
     __UpperCamelCase				      =  model_class_name.from_pretrained("""facebook/mbart-large-cc25"""		,	from_pt=lowercase			)
     # FlaxMBartForSequenceClassification expects eos token in input_ids
     __UpperCamelCase				      =  np.ones((1, 1)			) * model.config.eos_token_id
     __UpperCamelCase				      =  model(lowercase			)
     self.assertIsNotNone(lowercase			)
  
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class        UpperCAmelCase__  (    unittest.TestCase):
   __SCREAMING_SNAKE_CASE					      =				[
       ''' UN Chief Says There Is No Military Solution in Syria''',
   ]
   __SCREAMING_SNAKE_CASE					      =				[
       '''Şeful ONU declară că nu există o soluţie militară în Siria''',
   ]
   __SCREAMING_SNAKE_CASE					      =				'''facebook/mbart-large-en-ro'''
  
  
  
   @cached_property
   def    __lowerCamelCase				(  self			)  ->							int:
    return AutoTokenizer.from_pretrained(self.model_name			)
  
  
  
   @cached_property
   def    __lowerCamelCase				(  self			)  ->							Dict:
    __UpperCamelCase				      =  FlaxMBartForConditionalGeneration.from_pretrained(self.model_name		,	from_pt=lowercase			)
    return model
  
  
  
   def    __lowerCamelCase				(  self		,	**lowercase			)  ->							str:
    __UpperCamelCase				      =  self.translate_src_text(**lowercase			)
    self.assertListEqual(self.expected_text		,	lowercase			)
  
  
  
   def    __lowerCamelCase				(  self		,	**lowercase			)  ->							Any:
    __UpperCamelCase				      =  self.tokenizer(self.src_text		,	**lowercase		,	return_tensors="""np"""			)
    __UpperCamelCase				      =  self.model.generate(
        model_inputs.input_ids		,	attention_mask=model_inputs.attention_mask		,	decoder_start_token_id=self.tokenizer.lang_code_to_id["""ro_RO"""]		,	early_stopping=lowercase		,	num_beams=2		,	).sequences
    __UpperCamelCase				      =  self.tokenizer.batch_decode(lowercase		,	skip_special_tokens=lowercase			)
    return generated_words
  
  
  
  
   @slow
   def    __lowerCamelCase				(  self			)  ->							List[str]:
    self._assert_generated_batch_equal_expected()
  
  
  
  
  
  
style_context: 
  
  
  
  '''simple docstring'''
  
  
  
  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        UpperCAmelCase__  :
  
  
  
   def __init__(  self		,	lowercase		,	lowercase=1_3		,	lowercase=7		,	lowercase=True		,	lowercase=True		,	lowercase=True		,	lowercase=True		,	lowercase=9_9		,	lowercase=[1, 1, 2]		,	lowercase=1		,	lowercase=3_2		,	lowercase=4		,	lowercase=8		,	lowercase=3_7		,	lowercase="gelu_new"		,	lowercase=0.1		,	lowercase=0.1		,	lowercase=0.0		,	lowercase=5_1_2		,	lowercase=3		,	lowercase=0.02		,	lowercase=3		,	lowercase=4		,	lowercase=None		,	lowercase=False		,	)  ->							Any:
    __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    __lowerCamelCase				(  self			)  ->							Tuple:
    __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    __lowerCamelCase				(  self			)  ->							Any:
    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    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							Any:
    __UpperCamelCase				      =  FunnelModel(config=lowercase			)
    model.to(lowercase			)
    model.eval()
    __UpperCamelCase				      =  model(lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase			)
    __UpperCamelCase				      =  model(lowercase		,	token_type_ids=lowercase			)
    __UpperCamelCase				      =  model(lowercase			)
    self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, self.seq_length, self.d_model)			)
  
    __UpperCamelCase				      =  False
    __UpperCamelCase				      =  model(lowercase			)
    self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, self.seq_length, self.d_model)			)
  
    __UpperCamelCase				      =  False
    __UpperCamelCase				      =  model(lowercase			)
    self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, self.seq_length, self.d_model)			)
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							str:
    __UpperCamelCase				      =  FunnelBaseModel(config=lowercase			)
    model.to(lowercase			)
    model.eval()
    __UpperCamelCase				      =  model(lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase			)
    __UpperCamelCase				      =  model(lowercase		,	token_type_ids=lowercase			)
    __UpperCamelCase				      =  model(lowercase			)
    self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, 2, self.d_model)			)
  
    __UpperCamelCase				      =  False
    __UpperCamelCase				      =  model(lowercase			)
    self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, 3, self.d_model)			)
  
    __UpperCamelCase				      =  False
    __UpperCamelCase				      =  model(lowercase			)
    self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, 2, self.d_model)			)
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							Optional[Any]:
    __UpperCamelCase				      =  self.num_labels
    __UpperCamelCase				      =  FunnelForPreTraining(config=lowercase			)
    model.to(lowercase			)
    model.eval()
    __UpperCamelCase				      =  model(lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase		,	labels=lowercase			)
    self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length)			)
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							Optional[Any]:
    __UpperCamelCase				      =  FunnelForMaskedLM(config=lowercase			)
    model.to(lowercase			)
    model.eval()
    __UpperCamelCase				      =  model(lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase		,	labels=lowercase			)
    self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.vocab_size)			)
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							Any:
    __UpperCamelCase				      =  self.num_labels
    __UpperCamelCase				      =  FunnelForSequenceClassification(lowercase			)
    model.to(lowercase			)
    model.eval()
    __UpperCamelCase				      =  model(lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase		,	labels=lowercase			)
    self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.num_labels)			)
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							int:
    __UpperCamelCase				      =  self.num_choices
    __UpperCamelCase				      =  FunnelForMultipleChoice(config=lowercase			)
    model.to(lowercase			)
    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(
        lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase		,	labels=lowercase		,	)
    self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.num_choices)			)
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							Optional[Any]:
    __UpperCamelCase				      =  self.num_labels
    __UpperCamelCase				      =  FunnelForTokenClassification(config=lowercase			)
    model.to(lowercase			)
    model.eval()
    __UpperCamelCase				      =  model(lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase		,	labels=lowercase			)
    self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.seq_length, self.num_labels)			)
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	lowercase		,	)  ->							Any:
    __UpperCamelCase				      =  FunnelForQuestionAnswering(config=lowercase			)
    model.to(lowercase			)
    model.eval()
    __UpperCamelCase				      =  model(
        lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase		,	start_positions=lowercase		,	end_positions=lowercase		,	)
    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    __lowerCamelCase				(  self			)  ->							Any:
    __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        UpperCAmelCase__  (    UpperCAmelCase_							,			UpperCAmelCase_							,			unittest.TestCase):
   __SCREAMING_SNAKE_CASE					      =				False
   __SCREAMING_SNAKE_CASE					      =				False
   __SCREAMING_SNAKE_CASE					      =				(
       (
           FunnelModel,
           FunnelForMaskedLM,
           FunnelForPreTraining,
           FunnelForQuestionAnswering,
           FunnelForTokenClassification,
       )
       if is_torch_available()
       else ()
   )
   __SCREAMING_SNAKE_CASE					      =				(
       {
           '''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    __lowerCamelCase				(  self		,	lowercase		,	lowercase		,	lowercase=False			)  ->							Tuple:
    __UpperCamelCase				      =  super()._prepare_for_class(lowercase		,	lowercase		,	return_labels=lowercase			)
  
    if return_labels:
     if model_class in get_values(lowercase			):
      __UpperCamelCase				      =  torch.zeros(
          (self.model_tester.batch_size, self.model_tester.seq_length)		,	dtype=torch.long		,	device=lowercase			)
    return inputs_dict
  
  
  
   def    __lowerCamelCase				(  self			)  ->							List[str]:
    __UpperCamelCase				      =  FunnelModelTester(self			)
    __UpperCamelCase				      =  ConfigTester(self		,	config_class=lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Tuple:
    self.config_tester.run_common_tests()
  
  
  
   def    __lowerCamelCase				(  self			)  ->							List[str]:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Optional[int]:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_pretraining(*lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							str:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_masked_lm(*lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Dict:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_token_classification(*lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Tuple:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_question_answering(*lowercase			)
  
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase			)  ->							int:
    if hasattr(lowercase		,	"""weight"""			) and module.weight is not None:
     module.weight.data.fill_(3			)
    if hasattr(lowercase		,	"""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(lowercase		,	lowercase			) and getattr(lowercase		,	lowercase			) is not None:
      __UpperCamelCase				      =  getattr(lowercase		,	lowercase			)
      weight.data.fill_(3			)
  
  
  @require_torch
  class        UpperCAmelCase__  (    UpperCAmelCase_							,			unittest.TestCase):
   __SCREAMING_SNAKE_CASE					      =				False
   __SCREAMING_SNAKE_CASE					      =				False
   __SCREAMING_SNAKE_CASE					      =				(
       (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
   )
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Optional[int]:
    __UpperCamelCase				      =  FunnelModelTester(self		,	base=lowercase			)
    __UpperCamelCase				      =  ConfigTester(self		,	config_class=lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Any:
    self.config_tester.run_common_tests()
  
  
  
   def    __lowerCamelCase				(  self			)  ->							str:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_base_model(*lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Any:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_sequence_classification(*lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Any:
    __UpperCamelCase				      =  self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_multiple_choice(*lowercase			)
  
  
  
   def    __lowerCamelCase				(  self			)  ->							str:
    __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(lowercase			)
     model.to(lowercase			)
     model.train()
     __UpperCamelCase				      =  self._prepare_for_class(lowercase		,	lowercase		,	return_labels=lowercase			)
     __UpperCamelCase				      =  model(**lowercase			).loss
     loss.backward()
  
  
  
  
   def    __lowerCamelCase				(  self		,	lowercase			)  ->							List[Any]:
    if hasattr(lowercase		,	"""weight"""			) and module.weight is not None:
     module.weight.data.fill_(3			)
    if hasattr(lowercase		,	"""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(lowercase		,	lowercase			) and getattr(lowercase		,	lowercase			) is not None:
      __UpperCamelCase				      =  getattr(lowercase		,	lowercase			)
      weight.data.fill_(3			)
  
  
  @require_torch
  @require_sentencepiece
  @require_tokenizers
  class        UpperCAmelCase__  (    unittest.TestCase):
  
  
  
   def    __lowerCamelCase				(  self			)  ->							Tuple:
    __UpperCamelCase				      =  1_3
    __UpperCamelCase				      =  7
    __UpperCamelCase				      =  torch.arange(0		,	batch_size * sequence_length			).long().reshape(lowercase		,	lowercase			)
    __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(lowercase		,	token_type_ids=lowercase			)[0].abs()
  
    __UpperCamelCase				      =  torch.tensor(2_344.8_352			)
    __UpperCamelCase				      =  torch.tensor(0.8_052			)
    self.assertTrue(torch.allclose(output.sum()		,	lowercase		,	atol=1E-4			)			)
    self.assertTrue(torch.allclose(output.mean()		,	lowercase		,	atol=1E-4			)			)
  
    __UpperCamelCase				      =  torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]			)
    __UpperCamelCase				      =  model(lowercase		,	attention_mask=lowercase		,	token_type_ids=lowercase			)[0].abs()
  
    __UpperCamelCase				      =  torch.tensor(2_343.8_425			)
    __UpperCamelCase				      =  torch.tensor(0.8_049			)
    self.assertTrue(torch.allclose(output.sum()		,	lowercase		,	atol=1E-4			)			)
    self.assertTrue(torch.allclose(output.mean()		,	lowercase		,	atol=1E-4			)			)
  
  
  
  
   @slow
   def    __lowerCamelCase				(  self			)  ->							Optional[int]:
    __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(**lowercase			)[0]
  
    __UpperCamelCase				      =  torch.tensor(235.7_246			)
    __UpperCamelCase				      =  torch.tensor(0.0_256			)
    self.assertTrue(torch.allclose(output.sum()		,	lowercase		,	atol=1E-4			)			)
    self.assertTrue(torch.allclose(output.mean()		,	lowercase		,	atol=1E-4			)			)
  
  
  
  
  
  