
"""simple docstring"""






import unittest

from transformers import GPTNeoXJapaneseConfig, is_torch_available
from transformers.models.gpt_neox_japanese.tokenization_gpt_neox_japanese import GPTNeoXJapaneseTokenizer
from transformers.testing_utils import require_torch, 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 GPTNeoXJapaneseForCausalLM, GPTNeoXJapaneseModel



class 			lowercase					:






			def __init__(self  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	Any=13      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=7      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=True      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=True      ,SCREAMING_SNAKE_CASE_  :	int=True      ,SCREAMING_SNAKE_CASE_  :	Tuple=True      ,SCREAMING_SNAKE_CASE_  :	int=99      ,SCREAMING_SNAKE_CASE_  :	List[str]=32      ,SCREAMING_SNAKE_CASE_  :	Any=5      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=4      ,SCREAMING_SNAKE_CASE_  :	int=4      ,SCREAMING_SNAKE_CASE_  :	Tuple="gelu"      ,SCREAMING_SNAKE_CASE_  :	Dict=0.0      ,SCREAMING_SNAKE_CASE_  :	Any=0.1      ,SCREAMING_SNAKE_CASE_  :	str=True      ,SCREAMING_SNAKE_CASE_  :	int=512      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=16      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]=2      ,SCREAMING_SNAKE_CASE_  :	List[str]=0.02      ,SCREAMING_SNAKE_CASE_  :	str=3      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=4      ,SCREAMING_SNAKE_CASE_  :	str=None      ,)       ->	Any:





				"""simple docstring"""


				lowerCAmelCase												=  parent
				lowerCAmelCase												=  batch_size
				lowerCAmelCase												=  seq_length
				lowerCAmelCase												=  is_training
				lowerCAmelCase												=  use_input_mask
				lowerCAmelCase												=  use_token_type_ids
				lowerCAmelCase												=  use_labels
				lowerCAmelCase												=  vocab_size
				lowerCAmelCase												=  hidden_size
				lowerCAmelCase												=  num_hidden_layers
				lowerCAmelCase												=  num_attention_heads
				lowerCAmelCase												=  intermediate_multiple_size
				lowerCAmelCase												=  hidden_act
				lowerCAmelCase												=  hidden_dropout
				lowerCAmelCase												=  attention_dropout
				lowerCAmelCase												=  weight_tying
				lowerCAmelCase												=  max_position_embeddings
				lowerCAmelCase												=  type_vocab_size
				lowerCAmelCase												=  type_sequence_label_size
				lowerCAmelCase												=  initializer_range
				lowerCAmelCase												=  num_labels
				lowerCAmelCase												=  num_choices
				lowerCAmelCase												=  scope






			def 					UpperCAmelCase							(self  :	Optional[int]			)       ->	Tuple:





				"""simple docstring"""


				lowerCAmelCase												=  ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size			)

				lowerCAmelCase												=  None
				if self.use_input_mask:
					lowerCAmelCase												=  random_attention_mask([self.batch_size, self.seq_length]			)

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

				lowerCAmelCase												=  self.get_config()

				return config, input_ids, input_mask, token_labels






			def 					UpperCAmelCase							(self  :	str			)       ->	str:





				"""simple docstring"""


				return GPTNeoXJapaneseConfig(
				    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_multiple_size=self.intermediate_multiple_size      ,hidden_act=self.hidden_act      ,hidden_dropout=self.hidden_dropout      ,attention_dropout=self.attention_dropout      ,weight_tying=self.weight_tying      ,max_position_embeddings=self.max_position_embeddings      ,type_vocab_size=self.type_vocab_size      ,is_decoder=SCREAMING_SNAKE_CASE_      ,initializer_range=self.initializer_range      ,)






			def 					UpperCAmelCase							(self  :	str			)       ->	Any:





				"""simple docstring"""


				lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  self.prepare_config_and_inputs()

				lowerCAmelCase												=  True

				return config, input_ids, input_mask, token_labels






			def 					UpperCAmelCase							(self  :	Dict      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	str			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  GPTNeoXJapaneseModel(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.hidden_size)			)






			def 					UpperCAmelCase							(self  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	int      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[int]			)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  True
				lowerCAmelCase												=  GPTNeoXJapaneseModel(SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(result.last_hidden_state.shape      ,(self.batch_size, self.seq_length, self.hidden_size)			)






			def 					UpperCAmelCase							(self  :	Dict      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	int			)       ->	Optional[int]:





				"""simple docstring"""


				lowerCAmelCase												=  GPTNeoXJapaneseForCausalLM(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,labels=SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size)			)






			def 					UpperCAmelCase							(self  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	Any      ,SCREAMING_SNAKE_CASE_  :	Tuple			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  True
				lowerCAmelCase												=  GPTNeoXJapaneseForCausalLM(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()

				# first forward pass
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,use_cache=SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  outputs.past_key_values

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

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

				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,output_hidden_states=SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  output_from_no_past['''hidden_states'''][0]
				lowerCAmelCase												=  model(
				    SCREAMING_SNAKE_CASE_      ,attention_mask=SCREAMING_SNAKE_CASE_      ,past_key_values=SCREAMING_SNAKE_CASE_      ,output_hidden_states=SCREAMING_SNAKE_CASE_      ,)['''hidden_states'''][0]

				# select random slice
				lowerCAmelCase												=  ids_tensor((1,)      ,output_from_past.shape[-1]			).item()
				lowerCAmelCase												=  output_from_no_past[:, -3:, random_slice_idx].detach()
				lowerCAmelCase												=  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(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,atol=1e-3			)			)






			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  self.prepare_config_and_inputs()
				lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  config_and_inputs
				lowerCAmelCase												=  {'''input_ids''': input_ids, '''attention_mask''': input_mask}
				return config, inputs_dict





@require_torch
class 			lowercase					( lowercase__   ,lowercase__   ,unittest.TestCase		):
			lowercase        =     (GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
			lowercase        =     (GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
			lowercase        =     (
			    {'''feature-extraction''': GPTNeoXJapaneseModel, '''text-generation''': GPTNeoXJapaneseForCausalLM}
			    if is_torch_available()
			    else {}
			)
			lowercase        =     False
			lowercase        =     False
			lowercase        =     False
			lowercase        =     False






			def 					UpperCAmelCase							(self  :	int			)       ->	Any:





				"""simple docstring"""


				lowerCAmelCase												=  GPTNeoXJapaneseModelTester(self			)
				lowerCAmelCase												=  ConfigTester(self      ,config_class=SCREAMING_SNAKE_CASE_      ,hidden_size=37			)






			def 					UpperCAmelCase							(self  :	Any			)       ->	Union[str, Any]:





				"""simple docstring"""


				self.config_tester.run_common_tests()






			def 					UpperCAmelCase							(self  :	str			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	int			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_decoder()
				self.model_tester.create_and_check_model_as_decoder(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	int			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_decoder()

				lowerCAmelCase												=  None

				self.model_tester.create_and_check_model_as_decoder(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	str			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_decoder_model_past_large_inputs(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	List[str]			)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_causal_lm(*SCREAMING_SNAKE_CASE_			)






			@slow
			def 					UpperCAmelCase							(self  :	Dict			)       ->	Tuple:





				"""simple docstring"""


				lowerCAmelCase												=  '''abeja/gpt-neox-japanese-2.7b'''

				lowerCAmelCase												=  ['''データサイエンティストとは、''', '''100年後に必要とされる会社は、''', '''フルリモートの環境で働くために必要なことは、''', '''国境の長いトンネルを抜けると''', '''美味しい日本食といえば、''']

				lowerCAmelCase												=  [
				    '''データサイエンティストとは、データを分析し、ビジネスに役立つ知見を導き出す専門家のことです。''',
				    '''100年後に必要とされる会社は、「人」が中心の会社です。''',
				    '''フルリモートの環境で働くために必要なことは、「自分の時間をコントロールする」ことです。''',
				    '''国境の長いトンネルを抜けると、そこは雪国だった。''',
				    '''美味しい日本食といえば、やっぱりお寿司ですよね。''',
				]

				lowerCAmelCase												=  GPTNeoXJapaneseTokenizer.from_pretrained(SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  GPTNeoXJapaneseForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  []
				for prompt in prompts:
					lowerCAmelCase												=  tokenizer(SCREAMING_SNAKE_CASE_      ,return_tensors='''pt'''			).input_ids
					lowerCAmelCase												=  model.generate(SCREAMING_SNAKE_CASE_      ,max_length=50			)
					lowerCAmelCase												=  tokenizer.batch_decode(SCREAMING_SNAKE_CASE_      ,skip_special_tokens=SCREAMING_SNAKE_CASE_			)
					predicted_outputs += generated_string
				self.assertListEqual(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)




"""simple docstring"""






from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


UpperCAmelCase        =				logging.get_logger(__name__)

UpperCAmelCase        =				{
    """shi-labs/nat-mini-in1k-224""": """https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json""",
    # See all Nat models at https://huggingface.co/models?filter=nat
}



class 			lowercase					( lowercase__   ,lowercase__		):
			lowercase        =     '''nat'''

			lowercase        =     {
			    '''num_attention_heads''': '''num_heads''',
			    '''num_hidden_layers''': '''num_layers''',
			}






			def __init__(self  :	Union[str, Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=4      ,SCREAMING_SNAKE_CASE_  :	Dict=3      ,SCREAMING_SNAKE_CASE_  :	int=64      ,SCREAMING_SNAKE_CASE_  :	Tuple=[3, 4, 6, 5]      ,SCREAMING_SNAKE_CASE_  :	Dict=[2, 4, 8, 16]      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=7      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=3.0      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=True      ,SCREAMING_SNAKE_CASE_  :	Any=0.0      ,SCREAMING_SNAKE_CASE_  :	List[Any]=0.0      ,SCREAMING_SNAKE_CASE_  :	Dict=0.1      ,SCREAMING_SNAKE_CASE_  :	Tuple="gelu"      ,SCREAMING_SNAKE_CASE_  :	str=0.02      ,SCREAMING_SNAKE_CASE_  :	str=1e-5      ,SCREAMING_SNAKE_CASE_  :	int=0.0      ,SCREAMING_SNAKE_CASE_  :	List[Any]=None      ,SCREAMING_SNAKE_CASE_  :	Tuple=None      ,**SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,)       ->	Tuple:





				"""simple docstring"""


				super().__init__(**SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  patch_size
				lowerCAmelCase												=  num_channels
				lowerCAmelCase												=  embed_dim
				lowerCAmelCase												=  depths
				lowerCAmelCase												=  len(SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  num_heads
				lowerCAmelCase												=  kernel_size
				lowerCAmelCase												=  mlp_ratio
				lowerCAmelCase												=  qkv_bias
				lowerCAmelCase												=  hidden_dropout_prob
				lowerCAmelCase												=  attention_probs_dropout_prob
				lowerCAmelCase												=  drop_path_rate
				lowerCAmelCase												=  hidden_act
				lowerCAmelCase												=  layer_norm_eps
				lowerCAmelCase												=  initializer_range
				# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
				# this indicates the channel dimension after the last stage of the model
				lowerCAmelCase												=  int(embed_dim * 2 ** (len(SCREAMING_SNAKE_CASE_			) - 1)			)
				lowerCAmelCase												=  layer_scale_init_value
				lowerCAmelCase												=  ['''stem'''] + [F"""stage{idx}""" for idx in range(1      ,len(SCREAMING_SNAKE_CASE_			) + 1			)]
				lowerCAmelCase							,       lowerCAmelCase												=  get_aligned_output_features_output_indices(
				    out_features=SCREAMING_SNAKE_CASE_      ,out_indices=SCREAMING_SNAKE_CASE_      ,stage_names=self.stage_names			)


