





'''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 			UpperCAmelCase					:




      def __init__(      self      :			Optional[Any]      ,      __snake_case      :			Optional[int]      ,      __snake_case      :			Optional[Any]=13      ,      __snake_case      :			Any=7      ,      __snake_case      :			Optional[Any]=True      ,      __snake_case      :			Dict=True      ,      __snake_case      :			List[Any]=True      ,      __snake_case      :			Any=True      ,      __snake_case      :			int=99      ,      __snake_case      :			Union[str, Any]=32      ,      __snake_case      :			Union[str, Any]=5      ,      __snake_case      :			int=4      ,      __snake_case      :			str=4      ,      __snake_case      :			Any="gelu"      ,      __snake_case      :			Union[str, Any]=0.0      ,      __snake_case      :			Tuple=0.1      ,      __snake_case      :			Optional[Any]=True      ,      __snake_case      :			Union[str, Any]=5_12      ,      __snake_case      :			Tuple=16      ,      __snake_case      :			List[Any]=2      ,      __snake_case      :			List[Any]=0.02      ,      __snake_case      :			Union[str, Any]=3      ,      __snake_case      :			List[str]=4      ,      __snake_case      :			Union[str, Any]=None      ,      )    ->				Optional[int]:
            _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     lowercase__   (      self      :			List[str]  )    ->				Tuple:
            _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     lowercase__   (      self      :			Dict  )    ->				Any:
            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=UpperCamelCase__      ,      initializer_range=self.initializer_range      ,      )




      def     lowercase__   (      self      :			str  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							self.prepare_config_and_inputs()

            _lowerCAmelCase												=							True

            return config, input_ids, input_mask, token_labels




      def     lowercase__   (      self      :			List[Any]      ,      __snake_case      :			Tuple      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			List[str]  )    ->				Dict:
            _lowerCAmelCase												=							GPTNeoXJapaneseModel(config=UpperCamelCase__  )
            model.to(UpperCamelCase__  )
            model.eval()
            _lowerCAmelCase												=							model(UpperCamelCase__      ,      attention_mask=UpperCamelCase__  )
            _lowerCAmelCase												=							model(UpperCamelCase__  )
            self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)  )




      def     lowercase__   (      self      :			Tuple      ,      __snake_case      :			Optional[Any]      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Optional[Any]  )    ->				Tuple:
            _lowerCAmelCase												=							True
            _lowerCAmelCase												=							GPTNeoXJapaneseModel(UpperCamelCase__  )
            model.to(UpperCamelCase__  )
            model.eval()
            _lowerCAmelCase												=							model(UpperCamelCase__      ,      attention_mask=UpperCamelCase__  )
            self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)  )




      def     lowercase__   (      self      :			Dict      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Dict      ,      __snake_case      :			Optional[Any]      ,      __snake_case      :			List[str]  )    ->				Optional[int]:
            _lowerCAmelCase												=							GPTNeoXJapaneseForCausalLM(config=UpperCamelCase__  )
            model.to(UpperCamelCase__  )
            model.eval()
            _lowerCAmelCase												=							model(UpperCamelCase__      ,      attention_mask=UpperCamelCase__      ,      labels=UpperCamelCase__  )
            self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)  )




      def     lowercase__   (      self      :			Union[str, Any]      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Tuple  )    ->				List[Any]:
            _lowerCAmelCase												=							True
            _lowerCAmelCase												=							GPTNeoXJapaneseForCausalLM(config=UpperCamelCase__  )
            model.to(UpperCamelCase__  )
            model.eval()

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




      def     lowercase__   (      self      :			Optional[Any]  )    ->				Any:
            _lowerCAmelCase												=							self.prepare_config_and_inputs()
            _lowerCAmelCase												=							config_and_inputs
            _lowerCAmelCase												=							{'''input_ids''': input_ids, '''attention_mask''': input_mask}
            return config, inputs_dict




@require_torch
class 			UpperCAmelCase					(      __SCREAMING_SNAKE_CASE							,		__SCREAMING_SNAKE_CASE							,		unittest.TestCase       ):
      _lowercase:      Tuple											=       (GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
      _lowercase:      Dict											=       (GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
      _lowercase:      List[str]											=       (
          {'''feature-extraction''': GPTNeoXJapaneseModel, '''text-generation''': GPTNeoXJapaneseForCausalLM}
          if is_torch_available()
          else {}
      )
      _lowercase:      Union[str, Any]											=       False
      _lowercase:      List[Any]											=       False
      _lowercase:      str											=       False
      _lowercase:      Dict											=       False




      def     lowercase__   (      self      :			List[Any]  )    ->				Dict:
            _lowerCAmelCase												=							GPTNeoXJapaneseModelTester(self  )
            _lowerCAmelCase												=							ConfigTester(self      ,      config_class=UpperCamelCase__      ,      hidden_size=37  )




      def     lowercase__   (      self      :			Tuple  )    ->				Any:
            self.config_tester.run_common_tests()




      def     lowercase__   (      self      :			List[Any]  )    ->				Dict:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__  )




      def     lowercase__   (      self      :			List[Any]  )    ->				List[str]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs_for_decoder()
            self.model_tester.create_and_check_model_as_decoder(UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__  )




      def     lowercase__   (      self      :			Optional[int]  )    ->				Optional[Any]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs_for_decoder()

            _lowerCAmelCase												=							None

            self.model_tester.create_and_check_model_as_decoder(UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__  )




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				List[Any]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_decoder_model_past_large_inputs(UpperCamelCase__      ,      UpperCamelCase__      ,      UpperCamelCase__  )




      def     lowercase__   (      self      :			Any  )    ->				Optional[Any]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_causal_lm(*UpperCamelCase__  )




      @slow
      def     lowercase__   (      self      :			Optional[int]  )    ->				str:
            _lowerCAmelCase												=							'''abeja/gpt-neox-japanese-2.7b'''

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

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

            _lowerCAmelCase												=							GPTNeoXJapaneseTokenizer.from_pretrained(UpperCamelCase__  )
            _lowerCAmelCase												=							GPTNeoXJapaneseForCausalLM.from_pretrained(UpperCamelCase__  )

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












'''simple docstring'''






import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


A__    :       str              =5_00_00
A__    :       Optional[int]              =50_00

A__       ,						A__    :       Optional[int]              =os.path.split(__file__)
A__    :       Tuple              =os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      for i in range(lowerCAmelCase  ):
            _lowerCAmelCase												=							dataset[i]
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      for i in range(0							,     len(lowerCAmelCase  )							,     lowerCAmelCase  ):
            _lowerCAmelCase												=							dataset[i : i + batch_size]
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      with dataset.formatted_as(type=lowerCAmelCase  ):
            for i in range(lowerCAmelCase  ):
                  _lowerCAmelCase												=							dataset[i]
@get_duration
def        UpperCamelCase__				(		lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase							,     lowerCAmelCase  ):





      """simple docstring"""


      with dataset.formatted_as(type=lowerCAmelCase  ):
            for i in range(0							,     lowerCAmelCase							,     lowerCAmelCase  ):
                  _lowerCAmelCase												=							dataset[i : i + batch_size]
def        UpperCamelCase__				(		):





      """simple docstring"""


      _lowerCAmelCase												=							{"""num examples""": SPEED_TEST_N_EXAMPLES}
      _lowerCAmelCase												=							[
          (read, {"""length""": SMALL_TEST}),
          (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1_00}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10_00}),
          (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
          (read_formatted, {"""type""": """pandas""", """length""": SMALL_TEST}),
          (read_formatted, {"""type""": """torch""", """length""": SMALL_TEST}),
          (read_formatted, {"""type""": """tensorflow""", """length""": SMALL_TEST}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10_00}),
      ]

      _lowerCAmelCase												=							[
          (read, {"""length""": SMALL_TEST}),
          (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1_00}),
          (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10_00}),
          (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
          (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10_00}),
      ]
      with tempfile.TemporaryDirectory() as tmp_dir:
            print("""generating dataset"""  )
            _lowerCAmelCase												=							datasets.Features(
                {"""list""": datasets.Sequence(datasets.Value("""float32"""  )  ), """numbers""": datasets.Value("""float32"""  )}  )
            _lowerCAmelCase												=							generate_example_dataset(
                os.path.join(lowerCAmelCase							,     """dataset.arrow"""  )							,     lowerCAmelCase							,     num_examples=lowerCAmelCase							,     seq_shapes={"""list""": (1_00,)}							,     )
            print("""first set of iterations"""  )
            for func, kwargs in functions:
                  print(func.__name__							,     str(lowerCAmelCase  )  )
                  _lowerCAmelCase												=							func(lowerCAmelCase							,     **lowerCAmelCase  )

            print("""shuffling dataset"""  )
            _lowerCAmelCase												=							dataset.shuffle()
            print("""Second set of iterations (after shuffling"""  )
            for func, kwargs in functions_shuffled:
                  print("""shuffled """							,     func.__name__							,     str(lowerCAmelCase  )  )
                  _lowerCAmelCase												=							func(
                      lowerCAmelCase							,     **lowerCAmelCase  )

      with open(lowerCAmelCase							,     """wb"""  ) as f:
            f.write(json.dumps(lowerCAmelCase  ).encode("""utf-8"""  )  )


if __name__ == "__main__":  # useful to run the profiler
      benchmark_iterating()






