



"""simple docstring"""






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

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
   lowerCAmelCase__      =    '''platform'''

   import jax
   import jax.numpy as jnp

   from transformers import BlenderbotTokenizer
   from transformers.models.blenderbot.modeling_flax_blenderbot import (
       FlaxBlenderbotForConditionalGeneration,
       FlaxBlenderbotModel,
       shift_tokens_right,
   )






def 	snake_case_  (				A_ :		Tuple,    A_ :		List[Any],    A_ :		Dict=None,    A_ :		int=None,    A_ :		Any=None,    A_ :		List[Any]=None,    A_ :		Optional[int]=None,    A_ :		int=None,    ):


       '''simple docstring'''





       if attention_mask is None:
              _lowerCamelCase						:    Optional[int]         =			np.where(input_ids != config.pad_token_id,    1,    0						)
       if decoder_attention_mask is None:
              _lowerCamelCase						:    str         =			np.where(decoder_input_ids != config.pad_token_id,    1,    0						)
       if head_mask is None:
              _lowerCamelCase						:    str         =			np.ones((config.encoder_layers, config.encoder_attention_heads)						)
       if decoder_head_mask is None:
              _lowerCamelCase						:    str         =			np.ones((config.decoder_layers, config.decoder_attention_heads)						)
       if cross_attn_head_mask is None:
              _lowerCamelCase						:    List[Any]         =			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": attention_mask,
       }






class 	__snake_case	:




  def __init__(			self					:						Optional[int]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						List[Any]=1_3						, __lowerCAmelCase					:						Dict=7						, __lowerCAmelCase					:						Any=True						, __lowerCAmelCase					:						Optional[Any]=False						, __lowerCAmelCase					:						Optional[Any]=9_9						, __lowerCAmelCase					:						Dict=1_6						, __lowerCAmelCase					:						Optional[int]=2						, __lowerCAmelCase					:						Tuple=4						, __lowerCAmelCase					:						List[str]=4						, __lowerCAmelCase					:						str="gelu"						, __lowerCAmelCase					:						Tuple=0.1						, __lowerCAmelCase					:						List[Any]=0.1						, __lowerCAmelCase					:						str=3_2						, __lowerCAmelCase					:						Any=2						, __lowerCAmelCase					:						Optional[Any]=1						, __lowerCAmelCase					:						List[str]=0						, __lowerCAmelCase					:						List[str]=0.02						, ):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			parent
         _lowerCamelCase						:    Optional[int]         =			batch_size
         _lowerCamelCase						:    Any         =			seq_length
         _lowerCamelCase						:    List[Any]         =			is_training
         _lowerCamelCase						:    Optional[Any]         =			use_labels
         _lowerCamelCase						:    int         =			vocab_size
         _lowerCamelCase						:    Tuple         =			hidden_size
         _lowerCamelCase						:    List[str]         =			num_hidden_layers
         _lowerCamelCase						:    str         =			num_attention_heads
         _lowerCamelCase						:    List[str]         =			intermediate_size
         _lowerCamelCase						:    Optional[int]         =			hidden_act
         _lowerCamelCase						:    str         =			hidden_dropout_prob
         _lowerCamelCase						:    Tuple         =			attention_probs_dropout_prob
         _lowerCamelCase						:    List[str]         =			max_position_embeddings
         _lowerCamelCase						:    Tuple         =			eos_token_id
         _lowerCamelCase						:    Union[str, Any]         =			pad_token_id
         _lowerCamelCase						:    Union[str, Any]         =			bos_token_id
         _lowerCamelCase						:    Dict         =			initializer_range




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			np.clip(ids_tensor([self.batch_size, self.seq_length - 1]						, self.vocab_size		)						, 3						, self.vocab_size		)
         _lowerCamelCase						:    Tuple         =			np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)						, dtype=np.intaa		))						, -1		)

         _lowerCamelCase						:    int         =			shift_tokens_right(__lowerCAmelCase						, 1						, 2		)

         _lowerCamelCase						:    Union[str, Any]         =			BlenderbotConfig(
             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						, initializer_range=self.initializer_range						, use_cache=__lowerCAmelCase						, )
         _lowerCamelCase						:    Any         =			prepare_blenderbot_inputs_dict(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)
         return config, inputs_dict




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Union[str, Any]         =			self.prepare_config_and_inputs()
         return config, inputs_dict




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			2_0
         _lowerCamelCase						:    List[str]         =			model_class_name(__lowerCAmelCase		)

         _lowerCamelCase						:    List[str]         =			model.encode(inputs_dict['''input_ids''']		)

         _lowerCamelCase      ,   _lowerCamelCase						:    Optional[Any]         =			(
             inputs_dict['''decoder_input_ids'''],
             inputs_dict['''decoder_attention_mask'''],
         )

         _lowerCamelCase						:    Dict         =			model.init_cache(decoder_input_ids.shape[0]						, __lowerCAmelCase						, __lowerCAmelCase		)
         _lowerCamelCase						:    Optional[int]         =			jnp.ones((decoder_input_ids.shape[0], max_decoder_length)						, dtype='''i4'''		)

         _lowerCamelCase						:    int         =			jnp.broadcast_to(
             jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]						, (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)						, )
         _lowerCamelCase						:    Optional[Any]         =			model.decode(
             decoder_input_ids[:, :-1]						, __lowerCAmelCase						, decoder_attention_mask=__lowerCAmelCase						, past_key_values=__lowerCAmelCase						, decoder_position_ids=__lowerCAmelCase						, )

         _lowerCamelCase						:    Tuple         =			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]						, dtype='''i4'''		)
         _lowerCamelCase						:    Optional[Any]         =			model.decode(
             decoder_input_ids[:, -1:]						, __lowerCAmelCase						, decoder_attention_mask=__lowerCAmelCase						, past_key_values=outputs_cache.past_key_values						, decoder_position_ids=__lowerCAmelCase						, )

         _lowerCamelCase						:    Tuple         =			model.decode(__lowerCAmelCase						, __lowerCAmelCase		)

         _lowerCamelCase						:    Any         =			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 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						str		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			2_0
         _lowerCamelCase						:    Dict         =			model_class_name(__lowerCAmelCase		)

         _lowerCamelCase						:    Tuple         =			model.encode(inputs_dict['''input_ids''']		)

         _lowerCamelCase      ,   _lowerCamelCase						:    Dict         =			(
             inputs_dict['''decoder_input_ids'''],
             inputs_dict['''decoder_attention_mask'''],
         )

         _lowerCamelCase						:    Dict         =			jnp.concatenate(
             [
                 decoder_attention_mask,
                 jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])		),
             ]						, axis=-1						, )

         _lowerCamelCase						:    Optional[Any]         =			model.init_cache(decoder_input_ids.shape[0]						, __lowerCAmelCase						, __lowerCAmelCase		)
         _lowerCamelCase						:    Dict         =			jnp.broadcast_to(
             jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]						, (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)						, )

         _lowerCamelCase						:    List[str]         =			model.decode(
             decoder_input_ids[:, :-1]						, __lowerCAmelCase						, decoder_attention_mask=__lowerCAmelCase						, past_key_values=__lowerCAmelCase						, decoder_position_ids=__lowerCAmelCase						, )
         _lowerCamelCase						:    Optional[int]         =			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]						, dtype='''i4'''		)
         _lowerCamelCase						:    Any         =			model.decode(
             decoder_input_ids[:, -1:]						, __lowerCAmelCase						, past_key_values=outputs_cache.past_key_values						, decoder_attention_mask=__lowerCAmelCase						, decoder_position_ids=__lowerCAmelCase						, )

         _lowerCamelCase						:    List[str]         =			model.decode(__lowerCAmelCase						, __lowerCAmelCase						, decoder_attention_mask=__lowerCAmelCase		)

         _lowerCamelCase						:    Dict         =			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 	__snake_case	(     unittest.TestCase):
  snake_case__    :       Dict        =						9_9




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			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						, )

         _lowerCamelCase						:    Dict         =			input_ids.shape[0]
         _lowerCamelCase						:    str         =			BlenderbotConfig(
             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 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase      ,   _lowerCamelCase						:    str         =			self._get_config_and_data()
         _lowerCamelCase						:    int         =			FlaxBlenderbotForConditionalGeneration(__lowerCAmelCase		)
         _lowerCamelCase						:    Dict         =			lm_model(input_ids=__lowerCAmelCase		)
         _lowerCamelCase						:    Union[str, Any]         =			(batch_size, input_ids.shape[1], config.vocab_size)
         self.assertEqual(outputs['''logits'''].shape						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			BlenderbotConfig(
             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						, )
         _lowerCamelCase						:    Optional[Any]         =			FlaxBlenderbotForConditionalGeneration(__lowerCAmelCase		)
         _lowerCamelCase						:    Dict         =			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		)
         _lowerCamelCase						:    Optional[int]         =			np.array([[8_2, 7_1, 8_2, 1_8, 2], [5_8, 6_8, 2, 1, 1]]						, dtype=np.intaa		)
         _lowerCamelCase						:    Dict         =			lm_model(input_ids=__lowerCAmelCase						, decoder_input_ids=__lowerCAmelCase		)
         _lowerCamelCase						:    Union[str, Any]         =			(*summary.shape, config.vocab_size)
         self.assertEqual(outputs['''logits'''].shape						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			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		)
         _lowerCamelCase						:    List[str]         =			shift_tokens_right(__lowerCAmelCase						, 1						, 2		)
         _lowerCamelCase						:    Dict         =			np.equal(__lowerCAmelCase						, 1		).astype(np.floataa		).sum()
         _lowerCamelCase						:    Union[str, Any]         =			np.equal(__lowerCAmelCase						, 1		).astype(np.floataa		).sum()
         self.assertEqual(shifted.shape						, input_ids.shape		)
         self.assertEqual(__lowerCAmelCase						, n_pad_before - 1		)
         self.assertTrue(np.equal(shifted[:, 0]						, 2		).all()		)






@require_flax
class 	__snake_case	(     _lowercase  ,					unittest.TestCase  ,					_lowercase):
  snake_case__    :       int        =						True
  snake_case__    :       List[Any]        =						(
      (
          FlaxBlenderbotModel,
          FlaxBlenderbotForConditionalGeneration,
      )
      if is_flax_available()
      else ()
  )
  snake_case__    :       Any        =						(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			FlaxBlenderbotModelTester(self		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Tuple         =			self.model_tester.prepare_config_and_inputs()
         for model_class in self.all_model_classes:
                self.model_tester.check_use_cache_forward(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			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(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    List[str]         =			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
                with self.subTest(model_class.__name__		):
                       _lowerCamelCase						:    List[Any]         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase		)
                       _lowerCamelCase						:    Optional[int]         =			model_class(__lowerCAmelCase		)

                       @jax.jit
                       def encode_jitted(__lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Union[str, Any]=None						, **__lowerCAmelCase					:						Optional[Any]		):
                              return model.encode(input_ids=__lowerCAmelCase						, attention_mask=__lowerCAmelCase		)

                       with self.subTest('''JIT Enabled'''		):
                              _lowerCamelCase						:    Optional[int]         =			encode_jitted(**__lowerCAmelCase		).to_tuple()

                       with self.subTest('''JIT Disabled'''		):
                              with jax.disable_jit():
                                     _lowerCamelCase						:    Any         =			encode_jitted(**__lowerCAmelCase		).to_tuple()

                       self.assertEqual(len(__lowerCAmelCase		)						, len(__lowerCAmelCase		)		)
                       for jitted_output, output in zip(__lowerCAmelCase						, __lowerCAmelCase		):
                              self.assertEqual(jitted_output.shape						, output.shape		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    List[Any]         =			self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
                with self.subTest(model_class.__name__		):
                       _lowerCamelCase						:    List[Any]         =			model_class(__lowerCAmelCase		)
                       _lowerCamelCase						:    str         =			model.encode(inputs_dict['''input_ids''']						, inputs_dict['''attention_mask''']		)

                       _lowerCamelCase						:    Any         =			{
                           '''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(__lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Any		):
                              return model.decode(
                                  decoder_input_ids=__lowerCAmelCase						, decoder_attention_mask=__lowerCAmelCase						, encoder_outputs=__lowerCAmelCase						, )

                       with self.subTest('''JIT Enabled'''		):
                              _lowerCamelCase						:    Union[str, Any]         =			decode_jitted(**__lowerCAmelCase		).to_tuple()

                       with self.subTest('''JIT Disabled'''		):
                              with jax.disable_jit():
                                     _lowerCamelCase						:    Union[str, Any]         =			decode_jitted(**__lowerCAmelCase		).to_tuple()

                       self.assertEqual(len(__lowerCAmelCase		)						, len(__lowerCAmelCase		)		)
                       for jitted_output, output in zip(__lowerCAmelCase						, __lowerCAmelCase		):
                              self.assertEqual(jitted_output.shape						, output.shape		)




  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         for model_class_name in self.all_model_classes:
                _lowerCamelCase						:    Dict         =			model_class_name.from_pretrained('''facebook/blenderbot-400M-distill'''		)
                # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
                _lowerCamelCase						:    Tuple         =			np.ones((1, 1)		) * model.config.eos_token_id
                _lowerCamelCase						:    Tuple         =			model(__lowerCAmelCase		)
                self.assertIsNotNone(__lowerCAmelCase		)




  @unittest.skipUnless(jax_device != '''cpu'''						, '''3B test too slow on CPU.'''		)
  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			{'''num_beams''': 1, '''early_stopping''': True, '''min_length''': 1_5, '''max_length''': 2_5}
         _lowerCamelCase						:    List[Any]         =			{'''skip_special_tokens''': True, '''clean_up_tokenization_spaces''': True}

         _lowerCamelCase						:    Optional[Any]         =			FlaxBlenderbotForConditionalGeneration.from_pretrained('''facebook/blenderbot-3B'''						, from_pt=__lowerCAmelCase		)
         _lowerCamelCase						:    Optional[int]         =			BlenderbotTokenizer.from_pretrained('''facebook/blenderbot-3B'''		)

         _lowerCamelCase						:    List[str]         =			['''Sam''']
         _lowerCamelCase						:    Union[str, Any]         =			tokenizer(__lowerCAmelCase						, return_tensors='''jax'''		)

         _lowerCamelCase						:    List[Any]         =			model.generate(**__lowerCAmelCase						, **__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			'''Sam is a great name. It means "sun" in Gaelic.'''

         _lowerCamelCase						:    Dict         =			tokenizer.batch_decode(__lowerCAmelCase						, **__lowerCAmelCase		)
         assert generated_txt[0].strip() == tgt_text







"""simple docstring"""






import os
import tempfile
import unittest

from transformers import FlaubertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import (
       FlaubertForMultipleChoice,
       FlaubertForQuestionAnswering,
       FlaubertForQuestionAnsweringSimple,
       FlaubertForSequenceClassification,
       FlaubertForTokenClassification,
       FlaubertModel,
       FlaubertWithLMHeadModel,
   )
   from transformers.models.flaubert.modeling_flaubert import FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST






class 	__snake_case	(     _lowercase):




  def __init__(			self					:						List[Any]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						List[Any]=1_3						, __lowerCAmelCase					:						Any=7						, __lowerCAmelCase					:						Any=True						, __lowerCAmelCase					:						Any=True						, __lowerCAmelCase					:						Union[str, Any]=True						, __lowerCAmelCase					:						List[Any]=True						, __lowerCAmelCase					:						Any=True						, __lowerCAmelCase					:						Union[str, Any]=False						, __lowerCAmelCase					:						List[Any]=False						, __lowerCAmelCase					:						int=False						, __lowerCAmelCase					:						Tuple=2						, __lowerCAmelCase					:						Dict=9_9						, __lowerCAmelCase					:						str=0						, __lowerCAmelCase					:						Optional[Any]=3_2						, __lowerCAmelCase					:						Tuple=5						, __lowerCAmelCase					:						Tuple=4						, __lowerCAmelCase					:						List[str]=0.1						, __lowerCAmelCase					:						Union[str, Any]=0.1						, __lowerCAmelCase					:						Optional[int]=5_1_2						, __lowerCAmelCase					:						Any=1_2						, __lowerCAmelCase					:						Optional[int]=2						, __lowerCAmelCase					:						Any=0.02						, __lowerCAmelCase					:						str=3						, __lowerCAmelCase					:						Optional[int]=4						, __lowerCAmelCase					:						Optional[int]="last"						, __lowerCAmelCase					:						str=None						, __lowerCAmelCase					:						int=None						, ):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			parent
         _lowerCamelCase						:    List[str]         =			batch_size
         _lowerCamelCase						:    Dict         =			seq_length
         _lowerCamelCase						:    List[Any]         =			is_training
         _lowerCamelCase						:    Dict         =			use_input_lengths
         _lowerCamelCase						:    Tuple         =			use_token_type_ids
         _lowerCamelCase						:    Any         =			use_labels
         _lowerCamelCase						:    Optional[Any]         =			gelu_activation
         _lowerCamelCase						:    Optional[Any]         =			sinusoidal_embeddings
         _lowerCamelCase						:    Dict         =			causal
         _lowerCamelCase						:    Dict         =			asm
         _lowerCamelCase						:    str         =			n_langs
         _lowerCamelCase						:    str         =			vocab_size
         _lowerCamelCase						:    Optional[int]         =			n_special
         _lowerCamelCase						:    Dict         =			hidden_size
         _lowerCamelCase						:    int         =			num_hidden_layers
         _lowerCamelCase						:    str         =			num_attention_heads
         _lowerCamelCase						:    Dict         =			hidden_dropout_prob
         _lowerCamelCase						:    int         =			attention_probs_dropout_prob
         _lowerCamelCase						:    Any         =			max_position_embeddings
         _lowerCamelCase						:    Any         =			type_vocab_size
         _lowerCamelCase						:    Optional[int]         =			type_sequence_label_size
         _lowerCamelCase						:    List[str]         =			initializer_range
         _lowerCamelCase						:    List[Any]         =			num_labels
         _lowerCamelCase						:    Dict         =			num_choices
         _lowerCamelCase						:    str         =			summary_type
         _lowerCamelCase						:    List[str]         =			use_proj
         _lowerCamelCase						:    int         =			scope




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			ids_tensor([self.batch_size, self.seq_length]						, self.vocab_size		)
         _lowerCamelCase						:    Optional[Any]         =			random_attention_mask([self.batch_size, self.seq_length]		)

         _lowerCamelCase						:    Optional[int]         =			None
         if self.use_input_lengths:
                _lowerCamelCase						:    int         =			(
                    ids_tensor([self.batch_size]						, vocab_size=2		) + self.seq_length - 2
                )  # small variation of seq_length

         _lowerCamelCase						:    Union[str, Any]         =			None
         if self.use_token_type_ids:
                _lowerCamelCase						:    str         =			ids_tensor([self.batch_size, self.seq_length]						, self.n_langs		)

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

         _lowerCamelCase						:    Tuple         =			self.get_config()

         return (
             config,
             input_ids,
             token_type_ids,
             input_lengths,
             sequence_labels,
             token_labels,
             is_impossible_labels,
             choice_labels,
             input_mask,
         )




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         return FlaubertConfig(
             vocab_size=self.vocab_size						, n_special=self.n_special						, emb_dim=self.hidden_size						, n_layers=self.num_hidden_layers						, n_heads=self.num_attention_heads						, dropout=self.hidden_dropout_prob						, attention_dropout=self.attention_probs_dropout_prob						, gelu_activation=self.gelu_activation						, sinusoidal_embeddings=self.sinusoidal_embeddings						, asm=self.asm						, causal=self.causal						, n_langs=self.n_langs						, max_position_embeddings=self.max_position_embeddings						, initializer_range=self.initializer_range						, summary_type=self.summary_type						, use_proj=self.use_proj						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Tuple						, ):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			FlaubertModel(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         _lowerCamelCase						:    Union[str, Any]         =			model(__lowerCAmelCase						, lengths=__lowerCAmelCase						, langs=__lowerCAmelCase		)
         _lowerCamelCase						:    str         =			model(__lowerCAmelCase						, langs=__lowerCAmelCase		)
         _lowerCamelCase						:    List[str]         =			model(__lowerCAmelCase		)
         self.parent.assertEqual(result.last_hidden_state.shape						, (self.batch_size, self.seq_length, self.hidden_size)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						str						, ):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			FlaubertWithLMHeadModel(__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()

         _lowerCamelCase						:    str         =			model(__lowerCAmelCase						, token_type_ids=__lowerCAmelCase						, labels=__lowerCAmelCase		)
         self.parent.assertEqual(result.loss.shape						, ()		)
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.vocab_size)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Union[str, Any]						, ):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			FlaubertForQuestionAnsweringSimple(__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()

         _lowerCamelCase						:    List[str]         =			model(__lowerCAmelCase		)

         _lowerCamelCase						:    Union[str, Any]         =			model(__lowerCAmelCase						, start_positions=__lowerCAmelCase						, end_positions=__lowerCAmelCase		)
         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 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Optional[Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						int						, ):






         """simple docstring"""






         _lowerCamelCase						:    str         =			FlaubertForQuestionAnswering(__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()

         _lowerCamelCase						:    Optional[Any]         =			model(__lowerCAmelCase		)

         _lowerCamelCase						:    Optional[Any]         =			model(
             __lowerCAmelCase						, start_positions=__lowerCAmelCase						, end_positions=__lowerCAmelCase						, cls_index=__lowerCAmelCase						, is_impossible=__lowerCAmelCase						, p_mask=__lowerCAmelCase						, )

         _lowerCamelCase						:    List[str]         =			model(
             __lowerCAmelCase						, start_positions=__lowerCAmelCase						, end_positions=__lowerCAmelCase						, cls_index=__lowerCAmelCase						, is_impossible=__lowerCAmelCase						, )

         ((_lowerCamelCase)      ,   )						:    str         =			result_with_labels.to_tuple()

         _lowerCamelCase						:    Optional[Any]         =			model(__lowerCAmelCase						, start_positions=__lowerCAmelCase						, end_positions=__lowerCAmelCase		)

         ((_lowerCamelCase)      ,   )						:    Union[str, Any]         =			result_with_labels.to_tuple()

         self.parent.assertEqual(result_with_labels.loss.shape						, ()		)
         self.parent.assertEqual(result.start_top_log_probs.shape						, (self.batch_size, model.config.start_n_top)		)
         self.parent.assertEqual(result.start_top_index.shape						, (self.batch_size, model.config.start_n_top)		)
         self.parent.assertEqual(
             result.end_top_log_probs.shape						, (self.batch_size, model.config.start_n_top * model.config.end_n_top)		)
         self.parent.assertEqual(
             result.end_top_index.shape						, (self.batch_size, model.config.start_n_top * model.config.end_n_top)		)
         self.parent.assertEqual(result.cls_logits.shape						, (self.batch_size,)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						str						, ):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			FlaubertForSequenceClassification(__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()

         _lowerCamelCase						:    str         =			model(__lowerCAmelCase		)
         _lowerCamelCase						:    Tuple         =			model(__lowerCAmelCase						, labels=__lowerCAmelCase		)

         self.parent.assertEqual(result.loss.shape						, ()		)
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.type_sequence_label_size)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Tuple						, __lowerCAmelCase					:						Union[str, Any]						, ):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			self.num_labels
         _lowerCamelCase						:    List[str]         =			FlaubertForTokenClassification(__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()

         _lowerCamelCase						:    Union[str, Any]         =			model(__lowerCAmelCase						, attention_mask=__lowerCAmelCase						, labels=__lowerCAmelCase		)
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.seq_length, self.num_labels)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						Any						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Union[str, Any]						, __lowerCAmelCase					:						List[Any]						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Optional[Any]						, ):






         """simple docstring"""






         _lowerCamelCase						:    List[str]         =			self.num_choices
         _lowerCamelCase						:    Any         =			FlaubertForMultipleChoice(config=__lowerCAmelCase		)
         model.to(__lowerCAmelCase		)
         model.eval()
         _lowerCamelCase						:    List[str]         =			input_ids.unsqueeze(1		).expand(-1						, self.num_choices						, -1		).contiguous()
         _lowerCamelCase						:    List[str]         =			token_type_ids.unsqueeze(1		).expand(-1						, self.num_choices						, -1		).contiguous()
         _lowerCamelCase						:    List[Any]         =			input_mask.unsqueeze(1		).expand(-1						, self.num_choices						, -1		).contiguous()
         _lowerCamelCase						:    int         =			model(
             __lowerCAmelCase						, attention_mask=__lowerCAmelCase						, token_type_ids=__lowerCAmelCase						, labels=__lowerCAmelCase						, )
         self.parent.assertEqual(result.logits.shape						, (self.batch_size, self.num_choices)		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase						:    Any         =			self.prepare_config_and_inputs()
         (
             (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   (
             _lowerCamelCase
         )      ,   
         )						:    Optional[Any]         =			config_and_inputs
         _lowerCamelCase						:    int         =			{
             '''input_ids''': input_ids,
             '''token_type_ids''': token_type_ids,
             '''lengths''': input_lengths,
             '''attention_mask''': input_mask,
         }
         return config, inputs_dict






@require_torch
class 	__snake_case	(     _lowercase  ,					_lowercase  ,					unittest.TestCase):
  snake_case__    :       List[str]        =						(
      (
          FlaubertModel,
          FlaubertWithLMHeadModel,
          FlaubertForQuestionAnswering,
          FlaubertForQuestionAnsweringSimple,
          FlaubertForSequenceClassification,
          FlaubertForTokenClassification,
          FlaubertForMultipleChoice,
      )
      if is_torch_available()
      else ()
  )
  snake_case__    :       List[Any]        =						(
      {
          "feature-extraction": FlaubertModel,
          "fill-mask": FlaubertWithLMHeadModel,
          "question-answering": FlaubertForQuestionAnsweringSimple,
          "text-classification": FlaubertForSequenceClassification,
          "token-classification": FlaubertForTokenClassification,
          "zero-shot": FlaubertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )




  def 						SCREAMING_SNAKE_CASE    (			self					:						str						, __lowerCAmelCase					:						List[str]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Dict						, __lowerCAmelCase					:						Optional[int]						, __lowerCAmelCase					:						Tuple		):






         """simple docstring"""






         if (
             pipeline_test_casse_name == "QAPipelineTests"
             and tokenizer_name is not None
             and not tokenizer_name.endswith('''Fast'''		)
         ):
                # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
                # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
                # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
                return True

         return False




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						Optional[int]=False		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			super()._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase						, return_labels=__lowerCAmelCase		)

         if return_labels:
                if model_class.__name__ == "FlaubertForQuestionAnswering":
                       _lowerCamelCase						:    Dict         =			torch.zeros(
                           self.model_tester.batch_size						, dtype=torch.long						, device=__lowerCAmelCase		)
                       _lowerCamelCase						:    Any         =			torch.zeros(
                           self.model_tester.batch_size						, dtype=torch.long						, device=__lowerCAmelCase		)

         return inputs_dict




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			FlaubertModelTester(self		)
         _lowerCamelCase						:    str         =			ConfigTester(self						, config_class=__lowerCAmelCase						, emb_dim=3_7		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         self.config_tester.run_common_tests()




  def 						SCREAMING_SNAKE_CASE    (			self					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_model(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[str]		):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_lm_head(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_simple_qa(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple		):






         """simple docstring"""






         _lowerCamelCase						:    Tuple         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_qa(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[int]         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_sequence_classif(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Dict		):






         """simple docstring"""






         _lowerCamelCase						:    int         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_token_classif(*__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_flaubert_multiple_choice(*__lowerCAmelCase		)




  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         for model_name in FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                _lowerCamelCase						:    Dict         =			FlaubertModel.from_pretrained(__lowerCAmelCase		)
                self.assertIsNotNone(__lowerCAmelCase		)




  @slow
  @require_torch_gpu
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         _lowerCamelCase      ,   _lowerCamelCase						:    str         =			self.model_tester.prepare_config_and_inputs_for_common()
         for model_class in self.all_model_classes:
                # FlauBertForMultipleChoice behaves incorrectly in JIT environments.
                if model_class == FlaubertForMultipleChoice:
                       return

                _lowerCamelCase						:    Any         =			True
                _lowerCamelCase						:    int         =			model_class(config=__lowerCAmelCase		)

                _lowerCamelCase						:    List[str]         =			self._prepare_for_class(__lowerCAmelCase						, __lowerCAmelCase		)
                _lowerCamelCase						:    int         =			torch.jit.trace(
                    __lowerCAmelCase						, (inputs_dict['''input_ids'''].to('''cpu'''		), inputs_dict['''attention_mask'''].to('''cpu'''		))		)

                with tempfile.TemporaryDirectory() as tmp:
                       torch.jit.save(__lowerCAmelCase						, os.path.join(__lowerCAmelCase						, '''traced_model.pt'''		)		)
                       _lowerCamelCase						:    Union[str, Any]         =			torch.jit.load(os.path.join(__lowerCAmelCase						, '''traced_model.pt'''		)						, map_location=__lowerCAmelCase		)
                       loaded(inputs_dict['''input_ids'''].to(__lowerCAmelCase		)						, inputs_dict['''attention_mask'''].to(__lowerCAmelCase		)		)






@require_torch
class 	__snake_case	(     unittest.TestCase):




  @slow
  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[int]		):






         """simple docstring"""






         _lowerCamelCase						:    Union[str, Any]         =			FlaubertModel.from_pretrained('''flaubert/flaubert_base_cased'''		)
         _lowerCamelCase						:    Any         =			torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]		)
         with torch.no_grad():
                _lowerCamelCase						:    Any         =			model(__lowerCAmelCase		)[0]
         _lowerCamelCase						:    Optional[Any]         =			torch.Size((1, 1_1, 7_6_8)		)
         self.assertEqual(output.shape						, __lowerCAmelCase		)
         _lowerCamelCase						:    Optional[int]         =			torch.tensor(
             [[[-2.62_51, -1.42_98, -0.02_27], [-2.85_10, -1.63_87, 0.22_58], [-2.81_14, -1.18_32, -0.30_66]]]		)

         self.assertTrue(torch.allclose(output[:, :3, :3]						, __lowerCAmelCase						, atol=1E-4		)		)


