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





       """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:
          _SCREAMING_SNAKE_CASE      =				parent
          _SCREAMING_SNAKE_CASE      =				batch_size
          _SCREAMING_SNAKE_CASE      =				seq_length
          _SCREAMING_SNAKE_CASE      =				is_training
          _SCREAMING_SNAKE_CASE      =				use_input_mask
          _SCREAMING_SNAKE_CASE      =				use_token_type_ids
          _SCREAMING_SNAKE_CASE      =				use_labels
          _SCREAMING_SNAKE_CASE      =				vocab_size
          _SCREAMING_SNAKE_CASE      =				hidden_size
          _SCREAMING_SNAKE_CASE      =				num_hidden_layers
          _SCREAMING_SNAKE_CASE      =				num_attention_heads
          _SCREAMING_SNAKE_CASE      =				intermediate_multiple_size
          _SCREAMING_SNAKE_CASE      =				hidden_act
          _SCREAMING_SNAKE_CASE      =				hidden_dropout
          _SCREAMING_SNAKE_CASE      =				attention_dropout
          _SCREAMING_SNAKE_CASE      =				weight_tying
          _SCREAMING_SNAKE_CASE      =				max_position_embeddings
          _SCREAMING_SNAKE_CASE      =				type_vocab_size
          _SCREAMING_SNAKE_CASE      =				type_sequence_label_size
          _SCREAMING_SNAKE_CASE      =				initializer_range
          _SCREAMING_SNAKE_CASE      =				num_labels
          _SCREAMING_SNAKE_CASE      =				num_choices
          _SCREAMING_SNAKE_CASE      =				scope






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size      )

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

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

          _SCREAMING_SNAKE_CASE      =				self.get_config()

          return config, input_ids, input_mask, token_labels






       def        UpperCamelCase				(							self      )    ->			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=A__       ,							initializer_range=self.initializer_range       ,							)






       def        UpperCamelCase				(							self      )    ->			int:
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.prepare_config_and_inputs()

          _SCREAMING_SNAKE_CASE      =				True

          return config, input_ids, input_mask, token_labels






       def        UpperCamelCase				(							self       ,							A__       ,							A__       ,							A__      )    ->			Dict:
          _SCREAMING_SNAKE_CASE      =				GPTNeoXJapaneseModel(config=A__      )
          model.to(A__      )
          model.eval()
          _SCREAMING_SNAKE_CASE      =				model(A__       ,							attention_mask=A__      )
          _SCREAMING_SNAKE_CASE      =				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:
          _SCREAMING_SNAKE_CASE      =				True
          _SCREAMING_SNAKE_CASE      =				GPTNeoXJapaneseModel(A__      )
          model.to(A__      )
          model.eval()
          _SCREAMING_SNAKE_CASE      =				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__      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				GPTNeoXJapaneseForCausalLM(config=A__      )
          model.to(A__      )
          model.eval()
          _SCREAMING_SNAKE_CASE      =				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]:
          _SCREAMING_SNAKE_CASE      =				True
          _SCREAMING_SNAKE_CASE      =				GPTNeoXJapaneseForCausalLM(config=A__      )
          model.to(A__      )
          model.eval()

          # first forward pass
          _SCREAMING_SNAKE_CASE      =				model(A__       ,							attention_mask=A__       ,							use_cache=A__      )
          _SCREAMING_SNAKE_CASE      =				outputs.past_key_values

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

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

          _SCREAMING_SNAKE_CASE      =				model(A__       ,							attention_mask=A__       ,							output_hidden_states=A__      )
          _SCREAMING_SNAKE_CASE      =				output_from_no_past["""hidden_states"""][0]
          _SCREAMING_SNAKE_CASE      =				model(
              A__       ,							attention_mask=A__       ,							past_key_values=A__       ,							output_hidden_states=A__       ,							)["""hidden_states"""][0]

          # select random slice
          _SCREAMING_SNAKE_CASE      =				ids_tensor((1,)       ,							output_from_past.shape[-1]      ).item()
          _SCREAMING_SNAKE_CASE      =				output_from_no_past[:, -3:, random_slice_idx].detach()
          _SCREAMING_SNAKE_CASE      =				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      )    ->			int:
          _SCREAMING_SNAKE_CASE      =				self.prepare_config_and_inputs()
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				config_and_inputs
          _SCREAMING_SNAKE_CASE      =				{"""input_ids""": input_ids, """attention_mask""": input_mask}
          return config, inputs_dict






@require_torch
class    _a				(_lowerCamelCase  ,					_lowerCamelCase  ,					unittest.TestCase):





       """simple docstring"""




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






       def        UpperCamelCase				(							self      )    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				GPTNeoXJapaneseModelTester(self      )
          _SCREAMING_SNAKE_CASE      =				ConfigTester(self       ,							config_class=A__       ,							hidden_size=37      )






       def        UpperCamelCase				(							self      )    ->			Optional[int]:
          self.config_tester.run_common_tests()






       def        UpperCamelCase				(							self      )    ->			Dict:
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(A__       ,							A__       ,							A__      )






       def        UpperCamelCase				(							self      )    ->			int:
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_decoder()
          self.model_tester.create_and_check_model_as_decoder(A__       ,							A__       ,							A__      )






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          # This regression test was failing with PyTorch < 1.3
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs_for_decoder()

          _SCREAMING_SNAKE_CASE      =				None

          self.model_tester.create_and_check_model_as_decoder(A__       ,							A__       ,							A__      )






       def        UpperCamelCase				(							self      )    ->			Dict:
          _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_decoder_model_past_large_inputs(A__       ,							A__       ,							A__      )






       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_causal_lm(*A__      )






       @slow
       def        UpperCamelCase				(							self      )    ->			List[Any]:
          _SCREAMING_SNAKE_CASE      =				"""abeja/gpt-neox-japanese-2.7b"""

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

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

          _SCREAMING_SNAKE_CASE      =				GPTNeoXJapaneseTokenizer.from_pretrained(A__      )
          _SCREAMING_SNAKE_CASE      =				GPTNeoXJapaneseForCausalLM.from_pretrained(A__      )

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








'''simple docstring'''



import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger

class    _a				(_lowerCamelCase):





       """simple docstring"""




       def __init__(							self       ,							A__       ,							A__      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				params

          _SCREAMING_SNAKE_CASE      =				np.array(A__      )
          _SCREAMING_SNAKE_CASE      =				np.array([len(A__      ) for t in data]      )

          self.check()
          self.remove_long_sequences()
          self.remove_empty_sequences()
          self.remove_unknown_sequences()
          self.check()
          self.print_statistics()






       def __getitem__(							self       ,							A__      )    ->			Dict:
          return (self.token_ids[index], self.lengths[index])






       def __len__(							self      )    ->			Tuple:
          return len(self.lengths      )






       def        UpperCamelCase				(							self      )    ->			Dict:
          assert len(self.token_ids      ) == len(self.lengths      )
          assert all(self.lengths[i] == len(self.token_ids[i]      ) for i in range(len(self.lengths      )      )      )






       def        UpperCamelCase				(							self      )    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				self.params.max_model_input_size
          _SCREAMING_SNAKE_CASE      =				self.lengths > max_len
          logger.info(F"Splitting {sum(A__      )} too long sequences."      )

          def divide_chunks(A__       ,							A__      ):
             return [l[i : i + n] for i in range(0       ,							len(A__      )       ,							A__      )]

          _SCREAMING_SNAKE_CASE      =				[]
          _SCREAMING_SNAKE_CASE      =				[]
          if self.params.mlm:
             _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.params.special_tok_ids["""cls_token"""], self.params.special_tok_ids["""sep_token"""]
          else:
             _SCREAMING_SNAKE_CASE ,							_SCREAMING_SNAKE_CASE      =				self.params.special_tok_ids["""bos_token"""], self.params.special_tok_ids["""eos_token"""]

          for seq_, len_ in zip(self.token_ids       ,							self.lengths      ):
             assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
             if len_ <= max_len:
                new_tok_ids.append(seq_      )
                new_lengths.append(len_      )
             else:
                _SCREAMING_SNAKE_CASE      =				[]
                for sub_s in divide_chunks(seq_       ,							max_len - 2      ):
                   if sub_s[0] != cls_id:
                      _SCREAMING_SNAKE_CASE      =				np.insert(A__       ,							0       ,							A__      )
                   if sub_s[-1] != sep_id:
                      _SCREAMING_SNAKE_CASE      =				np.insert(A__       ,							len(A__      )       ,							A__      )
                   assert len(A__      ) <= max_len
                   assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
                   sub_seqs.append(A__      )

                new_tok_ids.extend(A__      )
                new_lengths.extend([len(A__      ) for l in sub_seqs]      )

          _SCREAMING_SNAKE_CASE      =				np.array(A__      )
          _SCREAMING_SNAKE_CASE      =				np.array(A__      )






       def        UpperCamelCase				(							self      )    ->			List[str]:
          _SCREAMING_SNAKE_CASE      =				len(self      )
          _SCREAMING_SNAKE_CASE      =				self.lengths > 11
          _SCREAMING_SNAKE_CASE      =				self.token_ids[indices]
          _SCREAMING_SNAKE_CASE      =				self.lengths[indices]
          _SCREAMING_SNAKE_CASE      =				len(self      )
          logger.info(F"Remove {init_size - new_size} too short (<=11 tokens) sequences."      )






       def        UpperCamelCase				(							self      )    ->			int:
          if "unk_token" not in self.params.special_tok_ids:
             return
          else:
             _SCREAMING_SNAKE_CASE      =				self.params.special_tok_ids["""unk_token"""]
          _SCREAMING_SNAKE_CASE      =				len(self      )
          _SCREAMING_SNAKE_CASE      =				np.array([np.count_nonzero(a == unk_token_id      ) for a in self.token_ids]      )
          _SCREAMING_SNAKE_CASE      =				(unk_occs / self.lengths) < 0.5
          _SCREAMING_SNAKE_CASE      =				self.token_ids[indices]
          _SCREAMING_SNAKE_CASE      =				self.lengths[indices]
          _SCREAMING_SNAKE_CASE      =				len(self      )
          logger.info(F"Remove {init_size - new_size} sequences with a high level of unknown tokens (50%)."      )






       def        UpperCamelCase				(							self      )    ->			Optional[Any]:
          if not self.params.is_master:
             return
          logger.info(F"{len(self      )} sequences"      )
          # data_len = sum(self.lengths)
          # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
          # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

          # unk_idx = self.params.special_tok_ids['unk_token']
          # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
          # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')






       def        UpperCamelCase				(							self       ,							A__      )    ->			Any:
          _SCREAMING_SNAKE_CASE      =				[t[0] for t in batch]
          _SCREAMING_SNAKE_CASE      =				[t[1] for t in batch]
          assert len(A__      ) == len(A__      )

          # Max for paddings
          _SCREAMING_SNAKE_CASE      =				max(A__      )

          # Pad token ids
          if self.params.mlm:
             _SCREAMING_SNAKE_CASE      =				self.params.special_tok_ids["""pad_token"""]
          else:
             _SCREAMING_SNAKE_CASE      =				self.params.special_tok_ids["""unk_token"""]
          _SCREAMING_SNAKE_CASE      =				[list(t.astype(A__      )      ) + [pad_idx] * (max_seq_len_ - len(A__      )) for t in token_ids]
          assert len(tk_      ) == len(A__      )
          assert all(len(A__      ) == max_seq_len_ for t in tk_      )

          _SCREAMING_SNAKE_CASE      =				torch.tensor(tk_      )  # (bs, max_seq_len_)
          _SCREAMING_SNAKE_CASE      =				torch.tensor(A__      )  # (bs)
          return tk_t, lg_t






