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





   """simple docstring"""




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



    """simple docstring"""



    _UpperCamelCase:     Dict							= parent
    _UpperCamelCase:     Tuple							= batch_size
    _UpperCamelCase:     int							= seq_length
    _UpperCamelCase:     List[Any]							= is_training
    _UpperCamelCase:     Dict							= use_input_mask
    _UpperCamelCase:     Tuple							= use_token_type_ids
    _UpperCamelCase:     Any							= use_labels
    _UpperCamelCase:     List[Any]							= vocab_size
    _UpperCamelCase:     Optional[int]							= hidden_size
    _UpperCamelCase:     List[str]							= num_hidden_layers
    _UpperCamelCase:     Optional[Any]							= num_attention_heads
    _UpperCamelCase:     Union[str, Any]							= intermediate_multiple_size
    _UpperCamelCase:     Any							= hidden_act
    _UpperCamelCase:     Any							= hidden_dropout
    _UpperCamelCase:     Dict							= attention_dropout
    _UpperCamelCase:     int							= weight_tying
    _UpperCamelCase:     Tuple							= max_position_embeddings
    _UpperCamelCase:     str							= type_vocab_size
    _UpperCamelCase:     Optional[int]							= type_sequence_label_size
    _UpperCamelCase:     Tuple							= initializer_range
    _UpperCamelCase:     Any							= num_labels
    _UpperCamelCase:     Union[str, Any]							= num_choices
    _UpperCamelCase:     Optional[Any]							= scope



   def   lowerCAmelCase						(	self					:  int							):



    """simple docstring"""



    _UpperCamelCase:     Optional[int]							= ids_tensor([self.batch_size, self.seq_length]		,  self.vocab_size							)

    _UpperCamelCase:     Optional[int]							= None
    if self.use_input_mask:
     _UpperCamelCase:     int							= random_attention_mask([self.batch_size, self.seq_length]							)

    _UpperCamelCase:     Union[str, Any]							= None
    if self.use_labels:
     _UpperCamelCase:     Dict							= ids_tensor([self.batch_size, self.seq_length]		,  self.num_labels							)

    _UpperCamelCase:     str							= self.get_config()

    return config, input_ids, input_mask, token_labels



   def   lowerCAmelCase						(	self					:  List[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=lowerCamelCase_		,  initializer_range=self.initializer_range		,  )



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase:     List[str]							= self.prepare_config_and_inputs()

    _UpperCamelCase:     Tuple							= True

    return config, input_ids, input_mask, token_labels



   def   lowerCAmelCase						(	self					:  List[Any]		,  _lowercase					:  Optional[Any]		,  _lowercase					:  int		,  _lowercase					:  List[str]							):



    """simple docstring"""



    _UpperCamelCase:     List[str]							= GPTNeoXJapaneseModel(config=lowerCamelCase_							)
    model.to(lowerCamelCase_							)
    model.eval()
    _UpperCamelCase:     str							= model(lowerCamelCase_		,  attention_mask=lowerCamelCase_							)
    _UpperCamelCase:     Any							= model(lowerCamelCase_							)
    self.parent.assertEqual(result.last_hidden_state.shape		,  (self.batch_size, self.seq_length, self.hidden_size)							)



   def   lowerCAmelCase						(	self					:  Optional[Any]		,  _lowercase					:  List[Any]		,  _lowercase					:  Any		,  _lowercase					:  Optional[int]							):



    """simple docstring"""



    _UpperCamelCase:     str							= True
    _UpperCamelCase:     Union[str, Any]							= GPTNeoXJapaneseModel(lowerCamelCase_							)
    model.to(lowerCamelCase_							)
    model.eval()
    _UpperCamelCase:     Optional[Any]							= model(lowerCamelCase_		,  attention_mask=lowerCamelCase_							)
    self.parent.assertEqual(result.last_hidden_state.shape		,  (self.batch_size, self.seq_length, self.hidden_size)							)



   def   lowerCAmelCase						(	self					:  int		,  _lowercase					:  Union[str, Any]		,  _lowercase					:  Optional[Any]		,  _lowercase					:  Dict		,  _lowercase					:  Any							):



    """simple docstring"""



    _UpperCamelCase:     Dict							= GPTNeoXJapaneseForCausalLM(config=lowerCamelCase_							)
    model.to(lowerCamelCase_							)
    model.eval()
    _UpperCamelCase:     Dict							= model(lowerCamelCase_		,  attention_mask=lowerCamelCase_		,  labels=lowerCamelCase_							)
    self.parent.assertEqual(result.logits.shape		,  (self.batch_size, self.seq_length, self.vocab_size)							)



   def   lowerCAmelCase						(	self					:  Union[str, Any]		,  _lowercase					:  Optional[int]		,  _lowercase					:  List[str]		,  _lowercase					:  Any							):



    """simple docstring"""



    _UpperCamelCase:     Union[str, Any]							= True
    _UpperCamelCase:     Union[str, Any]							= GPTNeoXJapaneseForCausalLM(config=lowerCamelCase_							)
    model.to(lowerCamelCase_							)
    model.eval()

    # first forward pass
    _UpperCamelCase:     Any							= model(lowerCamelCase_		,  attention_mask=lowerCamelCase_		,  use_cache=lowerCamelCase_							)
    _UpperCamelCase:     Optional[int]							= outputs.past_key_values

    # create hypothetical multiple next token and extent to next_input_ids
    _UpperCamelCase:     str							= ids_tensor((self.batch_size, 3)		,  config.vocab_size							)
    _UpperCamelCase:     List[Any]							= ids_tensor((self.batch_size, 3)		,  vocab_size=2							)

    # append to next input_ids and
    _UpperCamelCase:     List[Any]							= torch.cat([input_ids, next_tokens]		,  dim=-1							)
    _UpperCamelCase:     List[str]							= torch.cat([input_mask, next_mask]		,  dim=-1							)

    _UpperCamelCase:     Tuple							= model(lowerCamelCase_		,  attention_mask=lowerCamelCase_		,  output_hidden_states=lowerCamelCase_							)
    _UpperCamelCase:     List[Any]							= output_from_no_past['''hidden_states'''][0]
    _UpperCamelCase:     str							= model(
        lowerCamelCase_		,  attention_mask=lowerCamelCase_		,  past_key_values=lowerCamelCase_		,  output_hidden_states=lowerCamelCase_		,  )['''hidden_states'''][0]

    # select random slice
    _UpperCamelCase:     str							= ids_tensor((1,)		,  output_from_past.shape[-1]							).item()
    _UpperCamelCase:     Dict							= output_from_no_past[:, -3:, random_slice_idx].detach()
    _UpperCamelCase:     Union[str, Any]							= 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(lowerCamelCase_		,  lowerCamelCase_		,  atol=1E-3							)							)



   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



    _UpperCamelCase:     Tuple							= self.prepare_config_and_inputs()
    _UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase:     List[Any]							= config_and_inputs
    _UpperCamelCase:     List[Any]							= {'''input_ids''': input_ids, '''attention_mask''': input_mask}
    return config, inputs_dict




@require_torch
class  __magic_name__  (		a__      ,	a__      ,	unittest.TestCase    ):





   """simple docstring"""




   lowerCAmelCase      :    Optional[Any]  			=   (GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
   lowerCAmelCase      :    int  			=   (GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
   lowerCAmelCase      :    Tuple  			=   (
       {"feature-extraction": GPTNeoXJapaneseModel, "text-generation": GPTNeoXJapaneseForCausalLM}
       if is_torch_available()
       else {}
   )
   lowerCAmelCase      :    int  			=   False
   lowerCAmelCase      :    Union[str, Any]  			=   False
   lowerCAmelCase      :    Union[str, Any]  			=   False
   lowerCAmelCase      :    str  			=   False



   def   lowerCAmelCase						(	self					:  Optional[int]							):



    """simple docstring"""



    _UpperCamelCase:     List[Any]							= GPTNeoXJapaneseModelTester(self							)
    _UpperCamelCase:     Any							= ConfigTester(self		,  config_class=lowerCamelCase_		,  hidden_size=37							)



   def   lowerCAmelCase						(	self					:  Any							):



    """simple docstring"""



    self.config_tester.run_common_tests()



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase:     Union[str, Any]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(lowerCamelCase_		,  lowerCamelCase_		,  lowerCamelCase_							)



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase:     Optional[Any]							= self.model_tester.prepare_config_and_inputs_for_decoder()
    self.model_tester.create_and_check_model_as_decoder(lowerCamelCase_		,  lowerCamelCase_		,  lowerCamelCase_							)



   def   lowerCAmelCase						(	self					:  Dict							):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase:     Any							= self.model_tester.prepare_config_and_inputs_for_decoder()

    _UpperCamelCase:     Union[str, Any]							= None

    self.model_tester.create_and_check_model_as_decoder(lowerCamelCase_		,  lowerCamelCase_		,  lowerCamelCase_							)



   def   lowerCAmelCase						(	self					:  List[Any]							):



    """simple docstring"""



    _UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase	,							_UpperCamelCase:     Tuple							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_decoder_model_past_large_inputs(lowerCamelCase_		,  lowerCamelCase_		,  lowerCamelCase_							)



   def   lowerCAmelCase						(	self					:  Any							):



    """simple docstring"""



    _UpperCamelCase:     List[Any]							= self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_causal_lm(*lowerCamelCase_							)



   @slow
   def   lowerCAmelCase						(	self					:  Optional[Any]							):



    """simple docstring"""



    _UpperCamelCase:     List[str]							= '''abeja/gpt-neox-japanese-2.7b'''

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

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

    _UpperCamelCase:     Tuple							= GPTNeoXJapaneseTokenizer.from_pretrained(lowerCamelCase_							)
    _UpperCamelCase:     str							= GPTNeoXJapaneseForCausalLM.from_pretrained(lowerCamelCase_							)

    _UpperCamelCase:     Union[str, Any]							= []
    for prompt in prompts:
     _UpperCamelCase:     Union[str, Any]							= tokenizer(lowerCamelCase_		,  return_tensors='''pt'''							).input_ids
     _UpperCamelCase:     str							= model.generate(lowerCamelCase_		,  max_length=50							)
     _UpperCamelCase:     Any							= tokenizer.batch_decode(lowerCamelCase_		,  skip_special_tokens=lowerCamelCase_							)
     predicted_outputs += generated_string
    self.assertListEqual(lowerCamelCase_		,  lowerCamelCase_							)
# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


UpperCAmelCase_			    =							[
    '''EAGER''',
    '''AOT_EAGER''',
    '''INDUCTOR''',
    '''NVFUSER''',
    '''AOT_NVFUSER''',
    '''AOT_CUDAGRAPHS''',
    '''OFI''',
    '''FX2TRT''',
    '''ONNXRT''',
    '''IPEX''',
]




def       lowerCAmelCase_     (		lowercase:							str  ,					lowercase:							Optional[Any]=None  ,					lowercase:							List[str]=None  ,					lowercase:							List[str]=None		)    ->				Tuple:







 '''simple docstring'''
 _UpperCamelCase:     int							= True
 while ask_again:
  _UpperCamelCase:     Any							= input(lowercase		)
  try:
   if default is not None and len(lowercase		) == 0:
    return default
   return convert_value(lowercase		) if convert_value is not None else result
  except Exception:
   if error_message is not None:
    print(lowercase		)




def       lowerCAmelCase_     (		lowercase:							List[Any]  ,					lowercase:							Union[str, Any]=[]  ,					lowercase:							Union[str, Any]=None  ,					lowercase:							Union[str, Any]=0		)    ->				List[Any]:







 '''simple docstring'''
 _UpperCamelCase:     List[str]							= BulletMenu(lowercase  ,					lowercase		)
 _UpperCamelCase:     List[Any]							= menu.run(default_choice=lowercase		)
 return convert_value(lowercase		) if convert_value is not None else result




def       lowerCAmelCase_     (		lowercase:							Dict		)    ->				Optional[int]:







 '''simple docstring'''
 _UpperCamelCase:     int							= int(lowercase		)
 return ComputeEnvironment(['''LOCAL_MACHINE''', '''AMAZON_SAGEMAKER'''][value]		)




def       lowerCAmelCase_     (		lowercase:							Tuple		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     str							= int(lowercase		)
 return DistributedType(['''NO''', '''MULTI_CPU''', '''MULTI_XPU''', '''MULTI_GPU''', '''MULTI_NPU''', '''TPU'''][value]		)




def       lowerCAmelCase_     (		lowercase:							Union[str, Any]		)    ->				Optional[int]:







 '''simple docstring'''
 _UpperCamelCase:     Tuple							= int(lowercase		)
 return DynamoBackend(DYNAMO_BACKENDS[value]		).value




def       lowerCAmelCase_     (		lowercase:							int		)    ->				Any:







 '''simple docstring'''
 _UpperCamelCase:     int							= int(lowercase		)
 return PrecisionType(['''no''', '''fp16''', '''bf16''', '''fp8'''][value]		)




def       lowerCAmelCase_     (		lowercase:							List[str]		)    ->				List[str]:







 '''simple docstring'''
 _UpperCamelCase:     Dict							= int(lowercase		)
 return SageMakerDistributedType(['''NO''', '''DATA_PARALLEL''', '''MODEL_PARALLEL'''][value]		)




def       lowerCAmelCase_     (		lowercase:							Any		)    ->				List[str]:







 '''simple docstring'''
 return {"yes": True, "no": False}[value.lower()]





class  __magic_name__  (		argparse.RawDescriptionHelpFormatter    ):





   """simple docstring"""




   def   lowerCAmelCase						(	self					:  str		,  _lowercase					:  Optional[int]		,  _lowercase					:  List[Any]		,  _lowercase					:  Dict		,  _lowercase					:  Any							):



    """simple docstring"""



    _UpperCamelCase:     Optional[Any]							= super()._format_usage(_lowercase		,  _lowercase		,  _lowercase		,  _lowercase							)
    _UpperCamelCase:     int							= usage.replace('''<command> [<args>] '''		,  ''''''							)
    return usage