




import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
   import jax
   import jax.numpy as jnp

   from transformers.modeling_flax_pytorch_utils import (
       convert_pytorch_state_dict_to_flax,
       load_flax_weights_in_pytorch_model,
   )
   from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
   import torch

class   UpperCAmelCase    :






   def __init__(						self :							Tuple    ,       UpperCamelCase_ :							List[str]    ,       UpperCamelCase_ :							Dict=14    ,       UpperCamelCase_ :							str=7    ,       UpperCamelCase_ :							int=True    ,       UpperCamelCase_ :							List[str]=True    ,       UpperCamelCase_ :							Dict=False    ,       UpperCamelCase_ :							List[Any]=True    ,       UpperCamelCase_ :							int=99    ,       UpperCamelCase_ :							str=32    ,       UpperCamelCase_ :							Optional[Any]=4    ,       UpperCamelCase_ :							List[Any]=4    ,       UpperCamelCase_ :							List[str]=[[["global", "local"], 2]]    ,       UpperCamelCase_ :							str=37    ,       UpperCamelCase_ :							List[Any]="gelu"    ,       UpperCamelCase_ :							List[Any]=0.1    ,       UpperCamelCase_ :							Optional[int]=0.1    ,       UpperCamelCase_ :							Optional[Any]=5_12    ,       UpperCamelCase_ :							List[str]=7    ,       UpperCamelCase_ :							Optional[Any]=0.02    ,       )->      Optional[Any]:
     _lowercase:			str      =     parent
     _lowercase:			int      =     batch_size
     _lowercase:			str      =     seq_length
     _lowercase:			List[str]      =     is_training
     _lowercase:			Any      =     use_input_mask
     _lowercase:			List[Any]      =     use_token_type_ids
     _lowercase:			Any      =     use_labels
     _lowercase:			Optional[Any]      =     vocab_size
     _lowercase:			Dict      =     hidden_size
     _lowercase:			List[Any]      =     num_hidden_layers
     _lowercase:			Any      =     num_attention_heads
     _lowercase:			Tuple      =     attention_types
     _lowercase:			int      =     intermediate_size
     _lowercase:			Union[str, Any]      =     hidden_act
     _lowercase:			int      =     hidden_dropout_prob
     _lowercase:			Tuple      =     attention_probs_dropout_prob
     _lowercase:			str      =     max_position_embeddings
     _lowercase:			Optional[Any]      =     window_size
     _lowercase:			Dict      =     initializer_range
     _lowercase:			Dict      =     None
     _lowercase:			Optional[Any]      =     vocab_size - 1
     _lowercase:			Optional[Any]      =     vocab_size - 1
     _lowercase:			List[str]      =     vocab_size - 1






   def 			_a    (						self :							Union[str, Any]  )->      Dict:
     _lowercase:			int      =     ids_tensor([self.batch_size, self.seq_length]    ,       self.vocab_size  )

     _lowercase:			List[str]      =     None
     if self.use_input_mask:
       _lowercase:			Tuple      =     random_attention_mask([self.batch_size, self.seq_length]  )

     _lowercase:			Optional[int]      =     GPTNeoConfig(
         vocab_size=self.vocab_size    ,       hidden_size=self.hidden_size    ,       num_layers=self.num_hidden_layers    ,       num_heads=self.num_attention_heads    ,       max_position_embeddings=self.max_position_embeddings    ,       use_cache=UpperCamelCase_    ,       bos_token_id=self.bos_token_id    ,       eos_token_id=self.eos_token_id    ,       pad_token_id=self.pad_token_id    ,       window_size=self.window_size    ,       attention_types=self.attention_types    ,       )

     return (config, input_ids, input_mask)






   def 			_a    (						self :							Optional[int]  )->      Tuple:
     _lowercase:			Optional[int]      =     self.prepare_config_and_inputs()
     _lowercase:			Union[str, Any]      =     config_and_inputs
     _lowercase:			Optional[Any]      =     {'input_ids': input_ids, 'attention_mask': attention_mask}
     return config, inputs_dict






   def 			_a    (						self :							Tuple    ,       UpperCamelCase_ :							Union[str, Any]    ,       UpperCamelCase_ :							Any    ,       UpperCamelCase_ :							Tuple    ,       UpperCamelCase_ :							Tuple  )->      Optional[Any]:
     _lowercase:			Optional[int]      =     20
     _lowercase:			int      =     model_class_name(UpperCamelCase_  )

     _lowercase:			Tuple      =     model.init_cache(input_ids.shape[0]    ,       UpperCamelCase_  )
     _lowercase:			Dict      =     jnp.ones((input_ids.shape[0], max_decoder_length)    ,       dtype='i4'  )

     _lowercase:			List[str]      =     jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1  )[None, :]    ,       (input_ids.shape[0], input_ids.shape[-1] - 1)  )
     _lowercase:			Dict      =     model(
         input_ids[:, :-1]    ,       attention_mask=UpperCamelCase_    ,       past_key_values=UpperCamelCase_    ,       position_ids=UpperCamelCase_    ,       )

     _lowercase:			str      =     jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,       dtype='i4'  )
     _lowercase:			str      =     model(
         input_ids[:, -1:]    ,       attention_mask=UpperCamelCase_    ,       past_key_values=outputs_cache.past_key_values    ,       position_ids=UpperCamelCase_    ,       )

     _lowercase:			Any      =     model(UpperCamelCase_  )

     _lowercase:			List[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 			_a    (						self :							List[str]    ,       UpperCamelCase_ :							Tuple    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Union[str, Any]    ,       UpperCamelCase_ :							str  )->      int:
     _lowercase:			Union[str, Any]      =     20
     _lowercase:			str      =     model_class_name(UpperCamelCase_  )

     _lowercase:			Optional[Any]      =     jnp.concatenate(
         [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])  )]    ,       axis=-1    ,       )

     _lowercase:			int      =     model.init_cache(input_ids.shape[0]    ,       UpperCamelCase_  )
     _lowercase:			str      =     jnp.broadcast_to(
         jnp.arange(input_ids.shape[-1] - 1  )[None, :]    ,       (input_ids.shape[0], input_ids.shape[-1] - 1)  )

     _lowercase:			Union[str, Any]      =     model(
         input_ids[:, :-1]    ,       attention_mask=UpperCamelCase_    ,       past_key_values=UpperCamelCase_    ,       position_ids=UpperCamelCase_    ,       )
     _lowercase:			Any      =     jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]    ,       dtype='i4'  )
     _lowercase:			List[str]      =     model(
         input_ids[:, -1:]    ,       past_key_values=outputs_cache.past_key_values    ,       attention_mask=UpperCamelCase_    ,       position_ids=UpperCamelCase_    ,       )

     _lowercase:			List[str]      =     model(UpperCamelCase_    ,       attention_mask=UpperCamelCase_  )

     _lowercase:			Optional[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}'  )

@require_flax
class   UpperCAmelCase    (      a_ ,       a_ ,       unittest.TestCase ):
   __SCREAMING_SNAKE_CASE		:				List[str]      					=       (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
   __SCREAMING_SNAKE_CASE		:				Optional[Any]      					=       (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()






   def 			_a    (						self :							Dict  )->      List[Any]:
     _lowercase:			Union[str, Any]      =     FlaxGPTNeoModelTester(self  )






   def 			_a    (						self :							Dict  )->      Optional[Any]:
     for model_class_name in self.all_model_classes:
       _lowercase:			List[str]      =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward(UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_  )






   def 			_a    (						self :							int  )->      Dict:
     for model_class_name in self.all_model_classes:
       _lowercase:			List[str]      =     self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward_with_attn_mask(
           UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_  )






   @slow
   def 			_a    (						self :							List[Any]  )->      List[Any]:
     _lowercase:			List[str]      =     GPTaTokenizer.from_pretrained('gpt2'    ,       pad_token='<|endoftext|>'    ,       padding_side='left'  )
     _lowercase:			str      =     tokenizer(['Hello this is a long string', 'Hey']    ,       return_tensors='np'    ,       padding=UpperCamelCase_    ,       truncation=UpperCamelCase_  )

     _lowercase:			int      =     FlaxGPTNeoForCausalLM.from_pretrained('EleutherAI/gpt-neo-125M'  )
     _lowercase:			List[str]      =     False
     _lowercase:			List[str]      =     model.config.eos_token_id

     _lowercase:			Tuple      =     jax.jit(model.generate  )

     _lowercase:			Dict      =     jit_generate(
         inputs['input_ids']    ,       attention_mask=inputs['attention_mask']    ,       pad_token_id=tokenizer.pad_token_id  ).sequences

     _lowercase:			Optional[Any]      =     tokenizer.batch_decode(UpperCamelCase_    ,       skip_special_tokens=UpperCamelCase_  )

     _lowercase:			int      =     [
         'Hello this is a long string of text.\n\nI\'m trying to get the text of the',
         'Hey, I\'m a little late to the party. I\'m going to',
     ]

     self.assertListEqual(UpperCamelCase_    ,       UpperCamelCase_  )






   @is_pt_flax_cross_test
   def 			_a    (						self :							Optional[int]  )->      int:
     _lowercase:			Optional[Any]      =     self.model_tester.prepare_config_and_inputs_for_common()

     for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__  ):
         # prepare inputs
         _lowercase:			Tuple      =     self._prepare_for_class(UpperCamelCase_    ,       UpperCamelCase_  )
         _lowercase:			Union[str, Any]      =     {k: torch.tensor(v.tolist()  ) for k, v in prepared_inputs_dict.items()}

         # load corresponding PyTorch class
         _lowercase:			Dict      =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
         _lowercase:			int      =     getattr(UpperCamelCase_    ,       UpperCamelCase_  )

         _lowercase:			str      =     pt_inputs['input_ids'].shape
         _lowercase:			Optional[int]      =     np.random.randint(0    ,       seq_length - 1    ,       size=(batch_size,)  )
         for batch_idx, start_index in enumerate(UpperCamelCase_  ):
           _lowercase:			Optional[Any]      =     0
           _lowercase:			str      =     1
           _lowercase:			Union[str, Any]      =     0
           _lowercase:			Dict      =     1
         _lowercase:			Optional[int]      =     pt_model_class(UpperCamelCase_  ).eval()
         _lowercase:			Dict      =     model_class(UpperCamelCase_    ,       dtype=jnp.floataa  )

         _lowercase:			int      =     convert_pytorch_state_dict_to_flax(pt_model.state_dict()    ,       UpperCamelCase_  )
         _lowercase:			Dict      =     fx_state

         with torch.no_grad():
           _lowercase:			Dict      =     pt_model(**UpperCamelCase_  ).to_tuple()

         _lowercase:			List[str]      =     fx_model(**UpperCamelCase_  ).to_tuple()
         self.assertEqual(len(UpperCamelCase_  )    ,       len(UpperCamelCase_  )    ,       'Output lengths differ between Flax and PyTorch'  )
         for fx_output, pt_output in zip(UpperCamelCase_    ,       UpperCamelCase_  ):
           self.assert_almost_equals(fx_output[:, -1]    ,       pt_output[:, -1].numpy()    ,       4E-2  )

         with tempfile.TemporaryDirectory() as tmpdirname:
           pt_model.save_pretrained(UpperCamelCase_  )
           _lowercase:			Tuple      =     model_class.from_pretrained(UpperCamelCase_    ,       from_pt=UpperCamelCase_  )

         _lowercase:			List[Any]      =     fx_model_loaded(**UpperCamelCase_  ).to_tuple()
         self.assertEqual(
             len(UpperCamelCase_  )    ,       len(UpperCamelCase_  )    ,       'Output lengths differ between Flax and PyTorch'  )
         for fx_output_loaded, pt_output in zip(UpperCamelCase_    ,       UpperCamelCase_  ):
           self.assert_almost_equals(fx_output_loaded[:, -1]    ,       pt_output[:, -1].numpy()    ,       4E-2  )






   @is_pt_flax_cross_test
   def 			_a    (						self :							List[str]  )->      Optional[Any]:
     _lowercase:			Optional[int]      =     self.model_tester.prepare_config_and_inputs_for_common()
     for model_class in self.all_model_classes:
       with self.subTest(model_class.__name__  ):
         # prepare inputs
         _lowercase:			Dict      =     self._prepare_for_class(UpperCamelCase_    ,       UpperCamelCase_  )
         _lowercase:			List[str]      =     {k: torch.tensor(v.tolist()  ) for k, v in prepared_inputs_dict.items()}

         # load corresponding PyTorch class
         _lowercase:			Dict      =     model_class.__name__[4:]  # Skip the "Flax" at the beginning
         _lowercase:			Any      =     getattr(UpperCamelCase_    ,       UpperCamelCase_  )

         _lowercase:			Union[str, Any]      =     pt_model_class(UpperCamelCase_  ).eval()
         _lowercase:			Optional[Any]      =     model_class(UpperCamelCase_    ,       dtype=jnp.floataa  )

         _lowercase:			Any      =     load_flax_weights_in_pytorch_model(UpperCamelCase_    ,       fx_model.params  )
         _lowercase:			Optional[int]      =     pt_inputs['input_ids'].shape
         _lowercase:			Optional[Any]      =     np.random.randint(0    ,       seq_length - 1    ,       size=(batch_size,)  )
         for batch_idx, start_index in enumerate(UpperCamelCase_  ):
           _lowercase:			Optional[int]      =     0
           _lowercase:			List[Any]      =     1
           _lowercase:			Optional[int]      =     0
           _lowercase:			int      =     1

         # make sure weights are tied in PyTorch
         pt_model.tie_weights()

         with torch.no_grad():
           _lowercase:			Tuple      =     pt_model(**UpperCamelCase_  ).to_tuple()

         _lowercase:			Optional[int]      =     fx_model(**UpperCamelCase_  ).to_tuple()
         self.assertEqual(len(UpperCamelCase_  )    ,       len(UpperCamelCase_  )    ,       'Output lengths differ between Flax and PyTorch'  )
         for fx_output, pt_output in zip(UpperCamelCase_    ,       UpperCamelCase_  ):
           self.assert_almost_equals(fx_output[:, -1]    ,       pt_output[:, -1].numpy()    ,       4E-2  )

         with tempfile.TemporaryDirectory() as tmpdirname:
           fx_model.save_pretrained(UpperCamelCase_  )
           _lowercase:			Tuple      =     pt_model_class.from_pretrained(UpperCamelCase_    ,       from_flax=UpperCamelCase_  )

         with torch.no_grad():
           _lowercase:			Tuple      =     pt_model_loaded(**UpperCamelCase_  ).to_tuple()

         self.assertEqual(
             len(UpperCamelCase_  )    ,       len(UpperCamelCase_  )    ,       'Output lengths differ between Flax and PyTorch'  )
         for fx_output, pt_output in zip(UpperCamelCase_    ,       UpperCamelCase_  ):
           self.assert_almost_equals(fx_output[:, -1]    ,       pt_output[:, -1].numpy()    ,       4E-2  )






   @slow
   def 			_a    (						self :							Optional[Any]  )->      List[Any]:
     for model_class_name in self.all_model_classes:
       _lowercase:			Any      =     model_class_name.from_pretrained('EleutherAI/gpt-neo-125M'  )
       _lowercase:			int      =     model(np.ones((1, 1)  )  )
       self.assertIsNotNone(UpperCamelCase_  )

from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase      		=					logging.get_logger(__name__)

UpperCAmelCase      		=					{
    '''flaubert/flaubert_small_cased''': '''https://huggingface.co/flaubert/flaubert_small_cased/resolve/main/config.json''',
    '''flaubert/flaubert_base_uncased''': '''https://huggingface.co/flaubert/flaubert_base_uncased/resolve/main/config.json''',
    '''flaubert/flaubert_base_cased''': '''https://huggingface.co/flaubert/flaubert_base_cased/resolve/main/config.json''',
    '''flaubert/flaubert_large_cased''': '''https://huggingface.co/flaubert/flaubert_large_cased/resolve/main/config.json''',
}




class  __lowerCAmelCase	(							a_    ):



 '''simple docstring'''
 _lowerCamelCase    :						List[Any]					     =			'flaubert'
 _lowerCamelCase    :						List[str]					     =			{
     'hidden_size': 'emb_dim',
     'num_attention_heads': 'n_heads',
     'num_hidden_layers': 'n_layers',
     'n_words': 'vocab_size',  # For backward compatibility
 }

 def __init__(      self							:		List[Any]							,      lowerCamelCase							:		Tuple=False							,      lowerCamelCase							:		str=0.0							,      lowerCamelCase							:		int=3_0_1_4_5							,      lowerCamelCase							:		List[str]=2_0_4_8							,      lowerCamelCase							:		List[str]=1_2							,      lowerCamelCase							:		Dict=1_6							,      lowerCamelCase							:		Any=0.1							,      lowerCamelCase							:		List[str]=0.1							,      lowerCamelCase							:		List[str]=True							,      lowerCamelCase							:		Any=False							,      lowerCamelCase							:		Optional[int]=False							,      lowerCamelCase							:		int=False							,      lowerCamelCase							:		Optional[Any]=1							,      lowerCamelCase							:		int=True							,      lowerCamelCase							:		str=5_1_2							,      lowerCamelCase							:		Tuple=2_0_4_8**-0.5							,      lowerCamelCase							:		str=1e-12							,      lowerCamelCase							:		Any=0.0_2							,      lowerCamelCase							:		Dict=0							,      lowerCamelCase							:		Union[str, Any]=1							,      lowerCamelCase							:		List[Any]=2							,      lowerCamelCase							:		Union[str, Any]=3							,      lowerCamelCase							:		List[Any]=5							,      lowerCamelCase							:		int=True							,      lowerCamelCase							:		Dict="first"							,      lowerCamelCase							:		Optional[int]=True							,      lowerCamelCase							:		Any=None							,      lowerCamelCase							:		Optional[int]=True							,      lowerCamelCase							:		List[str]=0.1							,      lowerCamelCase							:		List[Any]=5							,      lowerCamelCase							:		Optional[Any]=5							,      lowerCamelCase							:		str=0							,      lowerCamelCase							:		Optional[int]=0							,      lowerCamelCase							:		List[str]=2							,      lowerCamelCase							:		Any=0							,      **lowerCamelCase							:		Tuple							,      )			->				Any:
        _UpperCamelCase  :							Optional[Any]	    =	pre_norm
        _UpperCamelCase  :							Optional[int]	    =	layerdrop
        _UpperCamelCase  :							Dict	    =	vocab_size
        _UpperCamelCase  :							Dict	    =	emb_dim
        _UpperCamelCase  :							int	    =	n_layers
        _UpperCamelCase  :							Tuple	    =	n_heads
        _UpperCamelCase  :							Tuple	    =	dropout
        _UpperCamelCase  :							Any	    =	attention_dropout
        _UpperCamelCase  :							Union[str, Any]	    =	gelu_activation
        _UpperCamelCase  :							Dict	    =	sinusoidal_embeddings
        _UpperCamelCase  :							str	    =	causal
        _UpperCamelCase  :							str	    =	asm
        _UpperCamelCase  :							Tuple	    =	n_langs
        _UpperCamelCase  :							Optional[int]	    =	use_lang_emb
        _UpperCamelCase  :							Optional[Any]	    =	layer_norm_eps
        _UpperCamelCase  :							Optional[Any]	    =	bos_index
        _UpperCamelCase  :							Dict	    =	eos_index
        _UpperCamelCase  :							Union[str, Any]	    =	pad_index
        _UpperCamelCase  :							Union[str, Any]	    =	unk_index
        _UpperCamelCase  :							Optional[Any]	    =	mask_index
        _UpperCamelCase  :							Any	    =	is_encoder
        _UpperCamelCase  :							List[str]	    =	max_position_embeddings
        _UpperCamelCase  :							Union[str, Any]	    =	embed_init_std
        _UpperCamelCase  :							Optional[int]	    =	init_std
        _UpperCamelCase  :							Any	    =	summary_type
        _UpperCamelCase  :							Any	    =	summary_use_proj
        _UpperCamelCase  :							Optional[int]	    =	summary_activation
        _UpperCamelCase  :							Union[str, Any]	    =	summary_proj_to_labels
        _UpperCamelCase  :							Dict	    =	summary_first_dropout
        _UpperCamelCase  :							Union[str, Any]	    =	start_n_top
        _UpperCamelCase  :							Any	    =	end_n_top
        _UpperCamelCase  :							List[Any]	    =	mask_token_id
        _UpperCamelCase  :							str	    =	lang_id

        if "n_words" in kwargs:
               _UpperCamelCase  :							Tuple	    =	kwargs['n_words']

        super().__init__(pad_token_id=lowerCamelCase							,      bos_token_id=lowerCamelCase							,      **lowerCamelCase	)







class  __lowerCAmelCase	(							a_    ):



 '''simple docstring'''

 @property
 def        __snake_case       (      self							:		str	)			->				Mapping[str, Mapping[int, str]]:
        if self.task == "multiple-choice":
               _UpperCamelCase  :							int	    =	{0: 'batch', 1: 'choice', 2: 'sequence'}
        else:
               _UpperCamelCase  :							Dict	    =	{0: 'batch', 1: 'sequence'}
        return OrderedDict(
            [
                ('input_ids', dynamic_axis),
                ('attention_mask', dynamic_axis),
            ]	)



