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  UpperCamelCase				:

     '''simple docstring'''

     def __init__(					self				,				A_				,				A_=13				,				A_=7				,				A_=True				,				A_=True				,				A_=True				,				A_=True				,				A_=99				,				A_=32				,				A_=5				,				A_=4				,				A_=4				,				A_="gelu"				,				A_=0.0				,				A_=0.1				,				A_=True				,				A_=5_12				,				A_=16				,				A_=2				,				A_=0.02				,				A_=3				,				A_=4				,				A_=None				,				) ->				Dict:






           """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						) ->				str:






           """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						) ->				Any:






           """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=A_				,				initializer_range=self.initializer_range				,				)



     def 		UpperCamelCase_						(					self						) ->				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				,				A_				,				A_				,				A_						) ->				Dict:






           """simple docstring"""
           _lowerCamelCase			    =				GPTNeoXJapaneseModel(config=A_						)
           model.to(A_						)
           model.eval()
           _lowerCamelCase			    =				model(A_				,				attention_mask=A_						)
           _lowerCamelCase			    =				model(A_						)
           self.parent.assertEqual(result.last_hidden_state.shape				,				(self.batch_size, self.seq_length, self.hidden_size)						)



     def 		UpperCamelCase_						(					self				,				A_				,				A_				,				A_						) ->				str:






           """simple docstring"""
           _lowerCamelCase			    =				True
           _lowerCamelCase			    =				GPTNeoXJapaneseModel(A_						)
           model.to(A_						)
           model.eval()
           _lowerCamelCase			    =				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				,				A_				,				A_				,				A_				,				A_						) ->				Tuple:






           """simple docstring"""
           _lowerCamelCase			    =				GPTNeoXJapaneseForCausalLM(config=A_						)
           model.to(A_						)
           model.eval()
           _lowerCamelCase			    =				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				,				A_				,				A_				,				A_						) ->				Optional[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				True
           _lowerCamelCase			    =				GPTNeoXJapaneseForCausalLM(config=A_						)
           model.to(A_						)
           model.eval()

           # first forward pass
           _lowerCamelCase			    =				model(A_				,				attention_mask=A_				,				use_cache=A_						)
           _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(A_				,				attention_mask=A_				,				output_hidden_states=A_						)
           _lowerCamelCase			    =				output_from_no_past['''hidden_states'''][0]
           _lowerCamelCase			    =				model(
               A_				,				attention_mask=A_				,				past_key_values=A_				,				output_hidden_states=A_				,				)['''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(A_				,				A_				,				atol=1E-3						)						)





     def 		UpperCamelCase_						(					self						) ->				Tuple:






           """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  UpperCamelCase				( __lowercase   ,  __lowercase   ,  unittest.TestCase ):

     '''simple docstring'''

     A_										=  (GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
     A_										=  (GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
     A_										=  (
         {'feature-extraction': GPTNeoXJapaneseModel, 'text-generation': GPTNeoXJapaneseForCausalLM}
         if is_torch_available()
         else {}
     )
     A_										=  False
     A_										=  False
     A_										=  False
     A_										=  False



     def 		UpperCamelCase_						(					self						) ->				Tuple:






           """simple docstring"""
           _lowerCamelCase			    =				GPTNeoXJapaneseModelTester(self						)
           _lowerCamelCase			    =				ConfigTester(self				,				config_class=A_				,				hidden_size=37						)



     def 		UpperCamelCase_						(					self						) ->				Optional[Any]:






           """simple docstring"""
           self.config_tester.run_common_tests()



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






           """simple docstring"""
           _lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(A_				,				A_				,				A_						)



     def 		UpperCamelCase_						(					self						) ->				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(A_				,				A_				,				A_						)



     def 		UpperCamelCase_						(					self						) ->				int:






           """simple docstring"""
           # This regression test was failing with PyTorch < 1.3
           _lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				self.model_tester.prepare_config_and_inputs_for_decoder()

           _lowerCamelCase			    =				None

           self.model_tester.create_and_check_model_as_decoder(A_				,				A_				,				A_						)



     def 		UpperCamelCase_						(					self						) ->				int:






           """simple docstring"""
           _lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase    ,					_lowerCamelCase			    =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_decoder_model_past_large_inputs(A_				,				A_				,				A_						)



     def 		UpperCamelCase_						(					self						) ->				Optional[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_causal_lm(*A_						)





     @slow
     def 		UpperCamelCase_						(					self						) ->				Optional[Any]:






           """simple docstring"""
           _lowerCamelCase			    =				'''abeja/gpt-neox-japanese-2.7b'''

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

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

           _lowerCamelCase			    =				GPTNeoXJapaneseTokenizer.from_pretrained(A_						)
           _lowerCamelCase			    =				GPTNeoXJapaneseForCausalLM.from_pretrained(A_						)

           _lowerCamelCase			    =				[]
           for prompt in prompts:
                 _lowerCamelCase			    =				tokenizer(A_				,				return_tensors='''pt'''						).input_ids
                 _lowerCamelCase			    =				model.generate(A_				,				max_length=50						)
                 _lowerCamelCase			    =				tokenizer.batch_decode(A_				,				skip_special_tokens=A_						)
                 predicted_outputs += generated_string
           self.assertListEqual(A_				,				A_						)
from typing import List

import numpy as np
def       __magic_name__(  __UpperCAmelCase		)						->    int:



      '''simple docstring'''

      _lowerCamelCase			    =				{key: len(__UpperCAmelCase		) for key, value in gen_kwargs.items() if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		)}
      if len(set(lists_lengths.values()		)		) > 1:
            raise RuntimeError(
                (
                    '''Sharding is ambiguous for this dataset: '''
                    + '''we found several data sources lists of different lengths, and we don\'t know over which list we should parallelize:\n'''
                    + '''\n'''.join(F'\t- key {key} has length {length}' for key, length in lists_lengths.items()		)
                    + '''\nTo fix this, check the \'gen_kwargs\' and make sure to use lists only for data sources, '''
                    + '''and use tuples otherwise. In the end there should only be one single list, or several lists with the same length.'''
                )		)
      _lowerCamelCase			    =				max(lists_lengths.values() ,						default=0		)
      return max(1 ,						__UpperCAmelCase		)
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    List[range]:



      '''simple docstring'''

      _lowerCamelCase			    =				[]
      for group_idx in range(__UpperCAmelCase		):
            _lowerCamelCase			    =				num_shards // max_num_jobs + (group_idx < (num_shards % max_num_jobs))
            if num_shards_to_add == 0:
                  break
            _lowerCamelCase			    =				shards_indices_per_group[-1].stop if shards_indices_per_group else 0
            _lowerCamelCase			    =				range(__UpperCAmelCase ,						start + num_shards_to_add		)
            shards_indices_per_group.append(__UpperCAmelCase		)
      return shards_indices_per_group
def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    List[dict]:



      '''simple docstring'''

      _lowerCamelCase			    =				_number_of_shards_in_gen_kwargs(__UpperCAmelCase		)
      if num_shards == 1:
            return [dict(__UpperCAmelCase		)]
      else:
            _lowerCamelCase			    =				_distribute_shards(num_shards=__UpperCAmelCase ,						max_num_jobs=__UpperCAmelCase		)
            return [
                {
                    key: [value[shard_idx] for shard_idx in shard_indices_per_group[group_idx]]
                    if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		)
                    else value
                    for key, value in gen_kwargs.items()
                }
                for group_idx in range(len(__UpperCAmelCase		)		)
            ]
def       __magic_name__(  __UpperCAmelCase		)						->    dict:



      '''simple docstring'''

      return {
          key: [value for gen_kwargs in gen_kwargs_list for value in gen_kwargs[key]]
          if isinstance(gen_kwargs_list[0][key] ,						__UpperCAmelCase		)
          else gen_kwargs_list[0][key]
          for key in gen_kwargs_list[0]
      }




def       __magic_name__(  __UpperCAmelCase ,						__UpperCAmelCase		)						->    dict:



      '''simple docstring'''

      _lowerCamelCase			    =				{len(__UpperCAmelCase		) for value in gen_kwargs.values() if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		)}
      _lowerCamelCase			    =				{}
      for size in list_sizes:
            _lowerCamelCase			    =				list(range(__UpperCAmelCase		)		)
            rng.shuffle(indices_per_size[size]		)
      # Now let's copy the gen_kwargs and shuffle the lists based on their sizes
      _lowerCamelCase			    =				dict(__UpperCAmelCase		)
      for key, value in shuffled_kwargs.items():
            if isinstance(__UpperCAmelCase ,						__UpperCAmelCase		):
                  _lowerCamelCase			    =				[value[i] for i in indices_per_size[len(__UpperCAmelCase		)]]
      return shuffled_kwargs