





import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

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.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
      import torch







class    __snake_case	:





       '''simple docstring'''







       def __init__(  self			:   Optional[int]							,    A			:   Optional[Any]							,    A			:   str=14							,    A			:   str=7							,    A			:   int=True							,    A			:   List[str]=True							,    A			:   Any=False							,    A			:   Tuple=True							,    A			:   Optional[Any]=99							,    A			:   Dict=32							,    A			:   Dict=4							,    A			:   List[str]=4							,    A			:   str=4							,    A			:   List[str]=37							,    A			:   Dict="gelu"							,    A			:   Tuple=0.1							,    A			:   Any=0.1							,    A			:   Union[str, Any]=512							,    A			:   List[str]=0.02							,    ):
         __snake_case:						str			  =						parent
         __snake_case:						Tuple			  =						batch_size
         __snake_case:						Dict			  =						seq_length
         __snake_case:						List[Any]			  =						is_training
         __snake_case:						List[str]			  =						use_input_mask
         __snake_case:						Union[str, Any]			  =						use_token_type_ids
         __snake_case:						Any			  =						use_labels
         __snake_case:						int			  =						vocab_size
         __snake_case:						Any			  =						hidden_size
         __snake_case:						List[str]			  =						rotary_dim
         __snake_case:						Any			  =						num_hidden_layers
         __snake_case:						Dict			  =						num_attention_heads
         __snake_case:						int			  =						intermediate_size
         __snake_case:						Optional[Any]			  =						hidden_act
         __snake_case:						Tuple			  =						hidden_dropout_prob
         __snake_case:						Tuple			  =						attention_probs_dropout_prob
         __snake_case:						int			  =						max_position_embeddings
         __snake_case:						str			  =						initializer_range
         __snake_case:						int			  =						None
         __snake_case:						Any			  =						vocab_size - 1
         __snake_case:						List[Any]			  =						vocab_size - 1
         __snake_case:						Union[str, Any]			  =						vocab_size - 1




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						List[Any]			  =						ids_tensor([self.batch_size, self.seq_length]							,    self.vocab_size		)

         __snake_case:						Any			  =						None
         if self.use_input_mask:
           __snake_case:						Tuple			  =						random_attention_mask([self.batch_size, self.seq_length]		)

         __snake_case:						int			  =						GPTJConfig(
             vocab_size=self.vocab_size							,    n_embd=self.hidden_size							,    n_layer=self.num_hidden_layers							,    n_head=self.num_attention_heads							,    n_positions=self.max_position_embeddings							,    use_cache=A							,    bos_token_id=self.bos_token_id							,    eos_token_id=self.eos_token_id							,    pad_token_id=self.pad_token_id							,    rotary_dim=self.rotary_dim							,    )

         return (config, input_ids, input_mask)




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         __snake_case:						Union[str, Any]			  =						self.prepare_config_and_inputs()
         __snake_case  ,       __snake_case  ,       __snake_case:						str			  =						config_and_inputs
         __snake_case:						int			  =						{"""input_ids""": input_ids, """attention_mask""": attention_mask}
         return config, inputs_dict




       def 							UpperCAmelCase__			(  self			:   Optional[Any]							,    A			:   Optional[Any]							,    A			:   Union[str, Any]							,    A			:   Optional[Any]							,    A			:   Any		):
         __snake_case:						List[Any]			  =						20
         __snake_case:						Optional[int]			  =						model_class_name(A		)

         __snake_case:						int			  =						model.init_cache(input_ids.shape[0]							,    A		)
         __snake_case:						Dict			  =						jnp.ones((input_ids.shape[0], max_decoder_length)							,    dtype="""i4"""		)

         __snake_case:						str			  =						jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1		)[None, :]							,    (input_ids.shape[0], input_ids.shape[-1] - 1)		)
         __snake_case:						int			  =						model(
             input_ids[:, :-1]							,    attention_mask=A							,    past_key_values=A							,    position_ids=A							,    )

         __snake_case:						Optional[int]			  =						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							,    dtype="""i4"""		)
         __snake_case:						List[str]			  =						model(
             input_ids[:, -1:]							,    attention_mask=A							,    past_key_values=outputs_cache.past_key_values							,    position_ids=A							,    )

         __snake_case:						int			  =						model(A		)

         __snake_case:						Union[str, 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 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   List[str]							,    A			:   Tuple							,    A			:   int							,    A			:   Optional[Any]		):
         __snake_case:						Optional[int]			  =						20
         __snake_case:						Dict			  =						model_class_name(A		)

         __snake_case:						Union[str, Any]			  =						jnp.concatenate(
             [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])		)]							,    axis=-1							,    )

         __snake_case:						Dict			  =						model.init_cache(input_ids.shape[0]							,    A		)
         __snake_case:						Tuple			  =						jnp.broadcast_to(
             jnp.arange(input_ids.shape[-1] - 1		)[None, :]							,    (input_ids.shape[0], input_ids.shape[-1] - 1)		)

         __snake_case:						Tuple			  =						model(
             input_ids[:, :-1]							,    attention_mask=A							,    past_key_values=A							,    position_ids=A							,    )
         __snake_case:						List[Any]			  =						jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]							,    dtype="""i4"""		)
         __snake_case:						Tuple			  =						model(
             input_ids[:, -1:]							,    past_key_values=outputs_cache.past_key_values							,    attention_mask=A							,    position_ids=A							,    )

         __snake_case:						Any			  =						model(A							,    attention_mask=A		)

         __snake_case:						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    __snake_case	(							__lowerCamelCase       ,      __lowerCamelCase       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
       lowerCAmelCase__						=					(FlaxGPTJForCausalLM,) if is_flax_available() else ()




       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						int			  =						FlaxGPTJModelTester(self		)




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         for model_class_name in self.all_model_classes:
           __snake_case  ,       __snake_case  ,       __snake_case:						List[Any]			  =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.check_use_cache_forward(A							,    A							,    A							,    A		)




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         for model_class_name in self.all_model_classes:
           __snake_case  ,       __snake_case  ,       __snake_case:						List[str]			  =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.check_use_cache_forward_with_attn_mask(
               A							,    A							,    A							,    A		)




       @tooslow
       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						Dict			  =						GPTaTokenizer.from_pretrained("""gpt2"""							,    pad_token="""<|endoftext|>"""							,    padding_side="""left"""		)
         __snake_case:						Optional[int]			  =						tokenizer(["""Hello this is a long string""", """Hey"""]							,    return_tensors="""np"""							,    padding=A							,    truncation=A		)

         __snake_case:						Any			  =						FlaxGPTJForCausalLM.from_pretrained("""EleutherAI/gpt-j-6B"""		)
         __snake_case:						List[str]			  =						False
         __snake_case:						Any			  =						model.config.eos_token_id

         __snake_case:						Dict			  =						jax.jit(model.generate		)

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

         __snake_case:						int			  =						tokenizer.batch_decode(A							,    skip_special_tokens=A		)

         __snake_case:						Any			  =						[
             """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(A							,    A		)




       @is_pt_flax_cross_test
       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case  ,       __snake_case:						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
             __snake_case:						Optional[int]			  =						self._prepare_for_class(A							,    A		)
             __snake_case:						Union[str, Any]			  =						{k: torch.tensor(v.tolist()		) for k, v in prepared_inputs_dict.items()}

             # load corresponding PyTorch class
             __snake_case:						int			  =						model_class.__name__[4:]  # Skip the "Flax" at the beginning
             __snake_case:						List[str]			  =						getattr(A							,    A		)

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

             __snake_case:						str			  =						convert_pytorch_state_dict_to_flax(pt_model.state_dict()							,    A		)
             __snake_case:						List[Any]			  =						fx_state

             with torch.no_grad():
               __snake_case:						str			  =						pt_model(**A		).to_tuple()

             __snake_case:						Tuple			  =						fx_model(**A		).to_tuple()
             self.assertEqual(len(A		)							,    len(A		)							,    """Output lengths differ between Flax and PyTorch"""		)
             for fx_output, pt_output in zip(A							,    A		):
               self.assert_almost_equals(fx_output[:, -1]							,    pt_output[:, -1].numpy()							,    4E-2		)

             with tempfile.TemporaryDirectory() as tmpdirname:
               pt_model.save_pretrained(A		)
               __snake_case:						List[Any]			  =						model_class.from_pretrained(A							,    from_pt=A		)

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




       @is_pt_flax_cross_test
       def 							UpperCAmelCase__			(  self			:   int		):
         __snake_case  ,       __snake_case:						Tuple			  =						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
             __snake_case:						Tuple			  =						self._prepare_for_class(A							,    A		)
             __snake_case:						List[str]			  =						{k: torch.tensor(v.tolist()		) for k, v in prepared_inputs_dict.items()}

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

             __snake_case:						str			  =						pt_model_class(A		).eval()
             __snake_case:						List[str]			  =						model_class(A							,    dtype=jnp.floataa		)

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

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

             with torch.no_grad():
               __snake_case:						Optional[int]			  =						pt_model(**A		).to_tuple()

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

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

             with torch.no_grad():
               __snake_case:						Optional[Any]			  =						pt_model_loaded(**A		).to_tuple()

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






       @tooslow
       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         for model_class_name in self.all_model_classes:
           __snake_case:						List[str]			  =						model_class_name.from_pretrained("""EleutherAI/gpt-j-6B"""		)
           __snake_case:						Union[str, Any]			  =						model(np.ones((1, 1)		)		)
           self.assertIsNotNone(A		)













import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCAmelCase	:       Optional[Any]    	=							logging.get_logger(__name__)

__UpperCAmelCase	:       Union[str, Any]    	=							{
    "asapp/sew-d-tiny-100k": "https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json",
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}







class    __snake_case	(							__lowerCamelCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					"""sew-d"""




       def __init__(  self			:   Dict							,    A			:   Any=32							,    A			:   Dict=768							,    A			:   Optional[Any]=12							,    A			:   Union[str, Any]=12							,    A			:   Union[str, Any]=3_072							,    A			:   Optional[Any]=2							,    A			:   Union[str, Any]=512							,    A			:   List[Any]=256							,    A			:   Dict=True							,    A			:   Union[str, Any]=True							,    A			:   Optional[int]=("p2c", "c2p")							,    A			:   str="layer_norm"							,    A			:   Dict="gelu_python"							,    A			:   Tuple=0.1							,    A			:   Any=0.1							,    A			:   Tuple=0.1							,    A			:   Optional[int]=0.0							,    A			:   Any=0.1							,    A			:   Any=0.02							,    A			:   Dict=1E-7							,    A			:   str=1E-5							,    A			:   int="group"							,    A			:   int="gelu"							,    A			:   str=(64, 128, 128, 128, 128, 256, 256, 256, 256, 512, 512, 512, 512)							,    A			:   Union[str, Any]=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1)							,    A			:   List[Any]=(10, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1)							,    A			:   Optional[int]=False							,    A			:   int=128							,    A			:   int=16							,    A			:   Optional[Any]=True							,    A			:   List[Any]=0.05							,    A			:   Any=10							,    A			:   Dict=2							,    A			:   List[Any]=0.0							,    A			:   Union[str, Any]=10							,    A			:   int=0							,    A			:   List[Any]="mean"							,    A			:   Union[str, Any]=False							,    A			:   Any=False							,    A			:   Optional[int]=256							,    A			:   List[Any]=0							,    A			:   Any=1							,    A			:   List[Any]=2							,    **A			:   List[Any]							,    ):
         super().__init__(**A							,    pad_token_id=A							,    bos_token_id=A							,    eos_token_id=A		)
         __snake_case:						Optional[int]			  =						hidden_size
         __snake_case:						str			  =						feat_extract_norm
         __snake_case:						int			  =						feat_extract_activation
         __snake_case:						str			  =						list(A		)
         __snake_case:						Any			  =						list(A		)
         __snake_case:						str			  =						list(A		)
         __snake_case:						Union[str, Any]			  =						conv_bias
         __snake_case:						int			  =						num_conv_pos_embeddings
         __snake_case:						str			  =						num_conv_pos_embedding_groups
         __snake_case:						List[Any]			  =						len(self.conv_dim		)
         __snake_case:						List[str]			  =						num_hidden_layers
         __snake_case:						Union[str, Any]			  =						intermediate_size
         __snake_case:						Dict			  =						squeeze_factor
         __snake_case:						List[Any]			  =						max_position_embeddings
         __snake_case:						List[Any]			  =						position_buckets
         __snake_case:						List[str]			  =						share_att_key
         __snake_case:						int			  =						relative_attention
         __snake_case:						Union[str, Any]			  =						norm_rel_ebd
         __snake_case:						List[str]			  =						list(A		)
         __snake_case:						Tuple			  =						hidden_act
         __snake_case:						List[Any]			  =						num_attention_heads
         __snake_case:						str			  =						hidden_dropout
         __snake_case:						int			  =						attention_dropout
         __snake_case:						Dict			  =						activation_dropout
         __snake_case:						Any			  =						feat_proj_dropout
         __snake_case:						int			  =						final_dropout
         __snake_case:						List[Any]			  =						layer_norm_eps
         __snake_case:						List[str]			  =						feature_layer_norm_eps
         __snake_case:						List[Any]			  =						initializer_range
         __snake_case:						List[Any]			  =						vocab_size

         if (
             (len(self.conv_stride		) != self.num_feat_extract_layers)
             or (len(self.conv_kernel		) != self.num_feat_extract_layers)
             or (len(self.conv_dim		) != self.num_feat_extract_layers)
         ):
           raise ValueError(
               """Configuration for convolutional layers is incorrect."""
               """It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,"""
               f'''but is `len(config.conv_dim) = {len(self.conv_dim		)}`, `len(config.conv_stride)'''
               f'''= {len(self.conv_stride		)}`, `len(config.conv_kernel) = {len(self.conv_kernel		)}`.'''		)

         # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
         __snake_case:						List[Any]			  =						apply_spec_augment
         __snake_case:						List[Any]			  =						mask_time_prob
         __snake_case:						str			  =						mask_time_length
         __snake_case:						List[str]			  =						mask_time_min_masks
         __snake_case:						str			  =						mask_feature_prob
         __snake_case:						Optional[int]			  =						mask_feature_length
         __snake_case:						Dict			  =						mask_feature_min_masks

         # ctc loss
         __snake_case:						Any			  =						ctc_loss_reduction
         __snake_case:						str			  =						ctc_zero_infinity

         # sequence classification
         __snake_case:						Optional[Any]			  =						use_weighted_layer_sum
         __snake_case:						List[Any]			  =						classifier_proj_size






       @property
       def 							UpperCAmelCase__			(  self			:   int		):
         return functools.reduce(operator.mul							,    self.conv_stride							,    1		)






