
"""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 __lowerCAmelCase		:





     def __init__(     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase=13     ,					__UpperCAmelCase=7     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=99     ,					__UpperCAmelCase=32     ,					__UpperCAmelCase=5     ,					__UpperCAmelCase=4     ,					__UpperCAmelCase=4     ,					__UpperCAmelCase="gelu"     ,					__UpperCAmelCase=0.0     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=512     ,					__UpperCAmelCase=16     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=0.0_2     ,					__UpperCAmelCase=3     ,					__UpperCAmelCase=4     ,					__UpperCAmelCase=None     ,					):

      '''simple docstring'''




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





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




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

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

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

      __UpperCamelCase      =	self.get_config()

      return config, input_ids, input_mask, token_labels





     def 						UpperCAmelCase (     self				):

      '''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=__UpperCAmelCase     ,					initializer_range=self.initializer_range     ,					)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase      =	self.prepare_config_and_inputs()

      __UpperCamelCase      =	True

      return config, input_ids, input_mask, token_labels





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




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





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




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





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




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





     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase      =	True
      __UpperCamelCase      =	GPTNeoXJapaneseForCausalLM(config=__UpperCAmelCase				)
      model.to(__UpperCAmelCase				)
      model.eval()

      # first forward pass
      __UpperCamelCase      =	model(__UpperCAmelCase     ,					attention_mask=__UpperCAmelCase     ,					use_cache=__UpperCAmelCase				)
      __UpperCamelCase      =	outputs.past_key_values

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

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

      __UpperCamelCase      =	model(__UpperCAmelCase     ,					attention_mask=__UpperCAmelCase     ,					output_hidden_states=__UpperCAmelCase				)
      __UpperCamelCase      =	output_from_no_past['hidden_states'][0]
      __UpperCamelCase      =	model(
          __UpperCAmelCase     ,					attention_mask=__UpperCAmelCase     ,					past_key_values=__UpperCAmelCase     ,					output_hidden_states=__UpperCAmelCase     ,					)['hidden_states'][0]

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

      '''simple docstring'''




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


@require_torch
class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ,  __SCREAMING_SNAKE_CASE      ,  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				):

      '''simple docstring'''




      __UpperCamelCase      =	GPTNeoXJapaneseModelTester(self				)
      __UpperCamelCase      =	ConfigTester(self     ,					config_class=__UpperCAmelCase     ,					hidden_size=37				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      self.config_tester.run_common_tests()





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




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

      __UpperCamelCase      =	None

      self.model_tester.create_and_check_model_as_decoder(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_decoder_model_past_large_inputs(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase				)





     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      __UpperCamelCase      =	self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_causal_lm(*__UpperCAmelCase				)







     @slow
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




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

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

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

      __UpperCamelCase      =	GPTNeoXJapaneseTokenizer.from_pretrained(__UpperCAmelCase				)
      __UpperCamelCase      =	GPTNeoXJapaneseForCausalLM.from_pretrained(__UpperCAmelCase				)

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





"""simple docstring"""

from typing import Union

import fire
import torch
from tqdm import tqdm






def 							A						(    snake_case				:str    , snake_case				:str = "cpu"    , snake_case				:Union[str, None] = None		)				->     None:
 __UpperCamelCase      =	torch.load(snake_case    , map_location=snake_case		)
 for k, v in tqdm(state_dict.items()		):
  if not isinstance(snake_case    , torch.Tensor		):
   raise TypeError('FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin'		)
  __UpperCamelCase      =	v.half()
 if save_path is None:  # overwrite src_path
  __UpperCamelCase      =	src_path
 torch.save(snake_case    , snake_case		)


if __name__ == "__main__":
      fire.Fire(convert)



