





import unittest

from parameterized import parameterized

from transformers import AutoTokenizer, GPTNeoXConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 (
          GPTNeoXForCausalLM,
          GPTNeoXForQuestionAnswering,
          GPTNeoXForSequenceClassification,
          GPTNeoXForTokenClassification,
          GPTNeoXModel,
      )







class    __snake_case	:





       '''simple docstring'''







       def __init__(  self			:   List[str]							,    A			:   Any							,    A			:   Dict=13							,    A			:   Any=7							,    A			:   Any=True							,    A			:   List[str]=True							,    A			:   Tuple=True							,    A			:   int=True							,    A			:   int=99							,    A			:   Optional[Any]=64							,    A			:   int=5							,    A			:   Optional[Any]=4							,    A			:   List[Any]=37							,    A			:   Tuple="gelu"							,    A			:   Optional[Any]=0.1							,    A			:   Dict=0.1							,    A			:   str=512							,    A			:   Union[str, Any]=16							,    A			:   Tuple=2							,    A			:   Any=0.02							,    A			:   str=3							,    A			:   List[str]=4							,    A			:   str=None							,    ):
         __snake_case:						List[str]			  =						parent
         __snake_case:						Optional[Any]			  =						batch_size
         __snake_case:						List[Any]			  =						seq_length
         __snake_case:						List[Any]			  =						is_training
         __snake_case:						Optional[int]			  =						use_input_mask
         __snake_case:						Dict			  =						use_token_type_ids
         __snake_case:						Tuple			  =						use_labels
         __snake_case:						str			  =						vocab_size
         __snake_case:						Dict			  =						hidden_size
         __snake_case:						Optional[Any]			  =						num_hidden_layers
         __snake_case:						str			  =						num_attention_heads
         __snake_case:						Optional[int]			  =						intermediate_size
         __snake_case:						List[Any]			  =						hidden_act
         __snake_case:						int			  =						hidden_dropout_prob
         __snake_case:						int			  =						attention_probs_dropout_prob
         __snake_case:						int			  =						max_position_embeddings
         __snake_case:						Optional[int]			  =						type_vocab_size
         __snake_case:						str			  =						type_sequence_label_size
         __snake_case:						Optional[int]			  =						initializer_range
         __snake_case:						Optional[Any]			  =						num_labels
         __snake_case:						Dict			  =						num_choices
         __snake_case:						Union[str, Any]			  =						scope
         __snake_case:						int			  =						vocab_size - 1




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

         __snake_case:						Optional[int]			  =						None
         if self.use_input_mask:
           __snake_case:						Dict			  =						random_attention_mask([self.batch_size, self.seq_length]		)

         __snake_case:						Any			  =						None
         if self.use_labels:
           __snake_case:						Tuple			  =						ids_tensor([self.batch_size, self.seq_length]							,    self.num_labels		)

         __snake_case:						Tuple			  =						self.get_config()

         return config, input_ids, input_mask, token_labels




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         return GPTNeoXConfig(
             vocab_size=self.vocab_size							,    hidden_size=self.hidden_size							,    num_hidden_layers=self.num_hidden_layers							,    num_attention_heads=self.num_attention_heads							,    intermediate_size=self.intermediate_size							,    hidden_act=self.hidden_act							,    hidden_dropout_prob=self.hidden_dropout_prob							,    attention_probs_dropout_prob=self.attention_probs_dropout_prob							,    max_position_embeddings=self.max_position_embeddings							,    type_vocab_size=self.type_vocab_size							,    is_decoder=A							,    initializer_range=self.initializer_range							,    pad_token_id=self.pad_token_id							,    )




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case  ,       __snake_case  ,       __snake_case  ,       __snake_case:						List[Any]			  =						self.prepare_config_and_inputs()

         __snake_case:						int			  =						True

         return config, input_ids, input_mask, token_labels




       def 							UpperCAmelCase__			(  self			:   Dict							,    A			:   Optional[int]							,    A			:   Optional[int]							,    A			:   Optional[int]		):
         __snake_case:						Any			  =						GPTNeoXModel(config=A		)
         model.to(A		)
         model.eval()
         __snake_case:						Any			  =						model(A							,    attention_mask=A		)
         __snake_case:						Optional[Any]			  =						model(A		)
         self.parent.assertEqual(result.last_hidden_state.shape							,    (self.batch_size, self.seq_length, self.hidden_size)		)




       def 							UpperCAmelCase__			(  self			:   int							,    A			:   Any							,    A			:   str							,    A			:   Any		):
         __snake_case:						Tuple			  =						True
         __snake_case:						Dict			  =						GPTNeoXModel(A		)
         model.to(A		)
         model.eval()
         __snake_case:						Any			  =						model(A							,    attention_mask=A		)
         self.parent.assertEqual(result.last_hidden_state.shape							,    (self.batch_size, self.seq_length, self.hidden_size)		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   Optional[int]							,    A			:   Optional[int]							,    A			:   List[Any]							,    A			:   List[Any]		):
         __snake_case:						List[Any]			  =						GPTNeoXForCausalLM(config=A		)
         model.to(A		)
         model.eval()
         __snake_case:						List[str]			  =						model(A							,    attention_mask=A							,    labels=A		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.vocab_size)		)




       def 							UpperCAmelCase__			(  self			:   int							,    A			:   List[Any]							,    A			:   Optional[Any]							,    A			:   int							,    A			:   Any		):
         __snake_case:						Any			  =						self.num_labels
         __snake_case:						List[str]			  =						GPTNeoXForQuestionAnswering(A		)
         model.to(A		)
         model.eval()
         __snake_case:						Tuple			  =						model(A							,    attention_mask=A		)
         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 							UpperCAmelCase__			(  self			:   List[Any]							,    A			:   Tuple							,    A			:   Optional[int]							,    A			:   Optional[int]							,    A			:   Union[str, Any]		):
         __snake_case:						Tuple			  =						self.num_labels
         __snake_case:						Tuple			  =						GPTNeoXForSequenceClassification(A		)
         model.to(A		)
         model.eval()
         __snake_case:						List[str]			  =						ids_tensor([self.batch_size]							,    self.type_sequence_label_size		)
         __snake_case:						Tuple			  =						model(A							,    attention_mask=A							,    labels=A		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_labels)		)




       def 							UpperCAmelCase__			(  self			:   List[Any]							,    A			:   str							,    A			:   List[str]							,    A			:   Any							,    A			:   int		):
         __snake_case:						Any			  =						self.num_labels
         __snake_case:						List[Any]			  =						GPTNeoXForTokenClassification(A		)
         model.to(A		)
         model.eval()
         __snake_case:						Optional[Any]			  =						model(A							,    attention_mask=A							,    labels=A		)
         self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.num_labels)		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]							,    A			:   Any							,    A			:   Union[str, Any]							,    A			:   Any		):
         __snake_case:						str			  =						True
         __snake_case:						Dict			  =						GPTNeoXForCausalLM(config=A		)
         model.to(A		)
         model.eval()

         # first forward pass
         __snake_case:						List[Any]			  =						model(A							,    attention_mask=A							,    use_cache=A		)
         __snake_case:						Any			  =						outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         __snake_case:						Any			  =						ids_tensor((self.batch_size, 3)							,    config.vocab_size		)
         __snake_case:						Any			  =						ids_tensor((self.batch_size, 3)							,    vocab_size=2		)

         # append to next input_ids and
         __snake_case:						Tuple			  =						torch.cat([input_ids, next_tokens]							,    dim=-1		)
         __snake_case:						str			  =						torch.cat([input_mask, next_mask]							,    dim=-1		)

         __snake_case:						str			  =						model(A							,    attention_mask=A							,    output_hidden_states=A		)
         __snake_case:						Tuple			  =						output_from_no_past["""hidden_states"""][0]
         __snake_case:						List[Any]			  =						model(
             A							,    attention_mask=A							,    past_key_values=A							,    output_hidden_states=A							,    )["""hidden_states"""][0]

         # select random slice
         __snake_case:						Optional[Any]			  =						ids_tensor((1,)							,    output_from_past.shape[-1]		).item()
         __snake_case:						int			  =						output_from_no_past[:, -3:, random_slice_idx].detach()
         __snake_case:						str			  =						output_from_past[:, :, random_slice_idx].detach()

         self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(A							,    A							,    atol=1E-3		)		)






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







@require_torch
class    __snake_case	(							__lowerCamelCase       ,      __lowerCamelCase       ,      __lowerCamelCase       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					(
           (
               GPTNeoXModel,
               GPTNeoXForCausalLM,
               GPTNeoXForQuestionAnswering,
               GPTNeoXForSequenceClassification,
               GPTNeoXForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       lowerCAmelCase__						=					(GPTNeoXForCausalLM,) if is_torch_available() else ()
       lowerCAmelCase__						=					(
           {
               """feature-extraction""": GPTNeoXModel,
               """question-answering""": GPTNeoXForQuestionAnswering,
               """text-classification""": GPTNeoXForSequenceClassification,
               """text-generation""": GPTNeoXForCausalLM,
               """token-classification""": GPTNeoXForTokenClassification,
               """zero-shot""": GPTNeoXForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False
       lowerCAmelCase__						=					False




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						List[str]			  =						GPTNeoXModelTester(self		)
         __snake_case:						Tuple			  =						ConfigTester(self							,    config_class=A							,    hidden_size=64							,    num_attention_heads=8		)




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         self.config_tester.run_common_tests()




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         __snake_case  ,       __snake_case  ,       __snake_case  ,       __snake_case:						str			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(A							,    A							,    A		)




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case  ,       __snake_case  ,       __snake_case  ,       __snake_case:						str			  =						self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(A							,    A							,    A		)




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         # This regression test was failing with PyTorch < 1.3
         __snake_case  ,       __snake_case  ,       __snake_case  ,       __snake_case:						Tuple			  =						self.model_tester.prepare_config_and_inputs_for_decoder()

         __snake_case:						List[str]			  =						None

         self.model_tester.create_and_check_model_as_decoder(A							,    A							,    A		)




       def 							UpperCAmelCase__			(  self			:   List[str]		):
         __snake_case  ,       __snake_case  ,       __snake_case  ,       __snake_case:						Tuple			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_decoder_model_past_large_inputs(A							,    A							,    A		)




       def 							UpperCAmelCase__			(  self			:   str		):
         __snake_case:						Optional[int]			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_causal_lm(*A		)




       def 							UpperCAmelCase__			(  self			:   int		):
         __snake_case:						str			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*A		)




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						str			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*A		)




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						str			  =						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*A		)




       @unittest.skip(reason="""Feed forward chunking is not implemented"""		)
       def 							UpperCAmelCase__			(  self			:   Dict		):
         pass






       @parameterized.expand([("""linear""",), ("""dynamic""",)]		)
       def 							UpperCAmelCase__			(  self			:   Tuple							,    A			:   Optional[Any]		):
         __snake_case  ,       __snake_case:						Any			  =						self.model_tester.prepare_config_and_inputs_for_common()
         __snake_case:						Optional[Any]			  =						ids_tensor([1, 10]							,    config.vocab_size		)
         __snake_case:						str			  =						ids_tensor([1, int(config.max_position_embeddings * 1.5		)]							,    config.vocab_size		)

         set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
         __snake_case:						Any			  =						GPTNeoXModel(A		)
         original_model.to(A		)
         original_model.eval()
         __snake_case:						str			  =						original_model(A		).last_hidden_state
         __snake_case:						Dict			  =						original_model(A		).last_hidden_state

         set_seed(42		)  # Fixed seed at init time so the two models get the same random weights
         __snake_case:						Dict			  =						{"""type""": scaling_type, """factor""": 10.0}
         __snake_case:						List[str]			  =						GPTNeoXModel(A		)
         scaled_model.to(A		)
         scaled_model.eval()
         __snake_case:						List[str]			  =						scaled_model(A		).last_hidden_state
         __snake_case:						Union[str, Any]			  =						scaled_model(A		).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
           self.assertTrue(torch.allclose(A							,    A							,    atol=1E-5		)		)
         else:
           self.assertFalse(torch.allclose(A							,    A							,    atol=1E-5		)		)

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(A							,    A							,    atol=1E-5		)		)







@require_torch
class    __snake_case	(							unittest.TestCase      ):





       '''simple docstring'''







       @slow
       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						Optional[Any]			  =						AutoTokenizer.from_pretrained("""EleutherAI/pythia-410m-deduped"""		)
         for checkpointing in [True, False]:
           __snake_case:						Tuple			  =						GPTNeoXForCausalLM.from_pretrained("""EleutherAI/pythia-410m-deduped"""		)

           if checkpointing:
             model.gradient_checkpointing_enable()
           else:
             model.gradient_checkpointing_disable()
           model.to(A		)

           __snake_case:						str			  =						tokenizer("""My favorite food is"""							,    return_tensors="""pt"""		).to(A		)
           # The hub repo. is updated on 2023-04-04, resulting in poor outputs.
           # See: https://github.com/huggingface/transformers/pull/24193
           __snake_case:						List[Any]			  =						"""My favorite food is a good old-fashioned, old-fashioned, old-fashioned.\n\nI'm not sure"""

           __snake_case:						Optional[int]			  =						model.generate(**A							,    do_sample=A							,    max_new_tokens=20		)
           __snake_case:						List[str]			  =						tokenizer.batch_decode(A		)[0]

           self.assertEqual(A							,    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		)






