




"""simple docstring"""





import unittest

import numpy as np

from transformers import (
    GPTSanJapaneseConfig,
    GPTSanJapaneseForConditionalGeneration,
    GPTSanJapaneseModel,
    GPTSanJapaneseTokenizer,
    is_torch_available,
)
from transformers.generation import GenerationConfig
from transformers.testing_utils import require_torch, slow, tooslow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin

class  __lowercase    :


				'''simple docstring'''


				def __init__(				self				:       Optional[Any] ,SCREAMING_SNAKE_CASE_				:       Optional[Any] ,SCREAMING_SNAKE_CASE_				:       Union[str, Any]=3_6_0_0_0 ,SCREAMING_SNAKE_CASE_				:       List[str]=1_3 ,SCREAMING_SNAKE_CASE_				:       Optional[int]=7 ,SCREAMING_SNAKE_CASE_				:       Optional[Any]=True ,SCREAMING_SNAKE_CASE_				:       Optional[int]=3_2 ,SCREAMING_SNAKE_CASE_				:       Union[str, Any]=4_2 ,SCREAMING_SNAKE_CASE_				:       Dict=5 ,SCREAMING_SNAKE_CASE_				:       List[str]=2 ,SCREAMING_SNAKE_CASE_				:       Union[str, Any]=4 ,SCREAMING_SNAKE_CASE_				:       int=2 ,SCREAMING_SNAKE_CASE_				:       Dict=3_2 ,SCREAMING_SNAKE_CASE_				:       Any=8_0 ,SCREAMING_SNAKE_CASE_				:       List[str]=3_3 ,SCREAMING_SNAKE_CASE_				:       Optional[Any]=0.0 ,SCREAMING_SNAKE_CASE_				:       List[str]=1E-6 ,SCREAMING_SNAKE_CASE_				:       Dict=1_0_0 ,SCREAMING_SNAKE_CASE_				:       Any=0.0 ,):
										snake_case__    :   str							     =   vocab_size
										snake_case__    :   str							     =   parent
										snake_case__    :   List[str]							     =   batch_size
										snake_case__    :   str							     =   num_contexts
										# For common tests
										snake_case__    :   Dict							     =   self.num_contexts
										snake_case__    :   Dict							     =   is_training
										snake_case__    :   Any							     =   hidden_size
										snake_case__    :   Optional[int]							     =   num_ext_layers
										snake_case__    :   Optional[Any]							     =   ext_size
										snake_case__    :   Any							     =   num_hidden_layers
										snake_case__    :   Union[str, Any]							     =   num_attention_heads
										snake_case__    :   Optional[Any]							     =   num_experts
										snake_case__    :   Optional[Any]							     =   d_ff
										snake_case__    :   int							     =   d_ext
										snake_case__    :   Union[str, Any]							     =   d_spout
										snake_case__    :   List[Any]							     =   dropout_rate
										snake_case__    :   Optional[int]							     =   layer_norm_epsilon
										snake_case__    :   List[Any]							     =   expert_capacity
										snake_case__    :   Union[str, Any]							     =   router_jitter_noise

				def  _lowerCamelCase   (				self				:       Optional[int] ):
										return GPTSanJapaneseConfig.from_pretrained('''Tanrei/GPTSAN-japanese''' )

				def  _lowerCamelCase   (				self				:       List[Any] ):
										snake_case__    :   int							     =   ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size )

										snake_case__    :   Optional[int]							     =   self.get_config()

										return (config, input_ids)

				def  _lowerCamelCase   (				self				:       str ):
										snake_case__    :   Union[str, Any]							     =   ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size )

										snake_case__    :   Tuple							     =   self.get_config()

										return (config, {"input_ids": input_ids})

				def  _lowerCamelCase   (				self				:       Union[str, Any] ):
										return GPTSanJapaneseConfig(
										    vocab_size=self.vocab_size ,num_contexts=self.seq_length ,d_model=self.hidden_size ,d_ff=self.d_ff ,d_ext=self.d_ext ,d_spout=self.d_spout ,num_switch_layers=self.num_hidden_layers - self.num_ext_layers ,num_ext_layers=self.num_ext_layers ,num_heads=self.num_attention_heads ,num_experts=self.num_experts ,expert_capacity=self.expert_capacity ,dropout_rate=self.dropout_rate ,layer_norm_epsilon=self.layer_norm_epsilon ,router_jitter_noise=self.router_jitter_noise ,)





				def  _lowerCamelCase   (				self				:       Any ,SCREAMING_SNAKE_CASE_				:       List[str] ,SCREAMING_SNAKE_CASE_				:       Union[str, Any] ,):
										snake_case__    :   Tuple							     =   GPTSanJapaneseForConditionalGeneration(config=SCREAMING_SNAKE_CASE_ )
										model.to(SCREAMING_SNAKE_CASE_ )
										model.eval()
										snake_case__    :   int							     =   model(
										    input_ids=SCREAMING_SNAKE_CASE_ ,)
										self.parent.assertIsNotNone(SCREAMING_SNAKE_CASE_ )







@require_torch
class  __lowercase    (     snake_case_  ,				snake_case_  ,				unittest.TestCase		):


				'''simple docstring'''


				_lowerCAmelCase		     =(GPTSanJapaneseModel,) if is_torch_available() else ()
				_lowerCAmelCase		     =(
				    {
				        "conversational": GPTSanJapaneseForConditionalGeneration,
				        "feature-extraction": GPTSanJapaneseForConditionalGeneration,
				        "summarization": GPTSanJapaneseForConditionalGeneration,
				        "text2text-generation": GPTSanJapaneseForConditionalGeneration,
				        "translation": GPTSanJapaneseForConditionalGeneration,
				    }
				    if is_torch_available()
				    else {}
				)
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				# The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
				_lowerCAmelCase		     =[0.8, 0.9]

				def  _lowerCamelCase   (				self				:       str ,SCREAMING_SNAKE_CASE_				:       int ,SCREAMING_SNAKE_CASE_				:       List[Any] ,SCREAMING_SNAKE_CASE_				:       Optional[int] ,SCREAMING_SNAKE_CASE_				:       Any ,SCREAMING_SNAKE_CASE_				:       Union[str, Any] ):
										if pipeline_test_casse_name == "SummarizationPipelineTests":
																# TODO: fix `_reorder_cache` is not implemented for this model
																return True
										elif pipeline_test_casse_name == "Text2TextGenerationPipelineTests":
																# TODO: check this.
																return True

										return False

				def  _lowerCamelCase   (				self				:       Dict ):
										snake_case__    :   Tuple							     =   GPTSanJapaneseTester(self )
										snake_case__    :   Dict							     =   ConfigTester(self ,config_class=SCREAMING_SNAKE_CASE_ ,d_model=3_7 )

				def  _lowerCamelCase   (				self				:       Union[str, Any] ):
										GPTSanJapaneseConfig()

				def  _lowerCamelCase   (				self				:       Dict ):
										snake_case__    :   Union[str, Any]							     =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_ )

				@unittest.skip(
				    reason='''skip for now as the computed `max_memory` by `model_split_percents` in the test method will be changed inside `from_pretrained`''' )
				def  _lowerCamelCase   (				self				:       Optional[Any] ):
										super().test_model_parallelism()





				@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''' )
				def  _lowerCamelCase   (				self				:       str ):
										pass







@require_torch
class  __lowercase    (     snake_case_  ,				snake_case_  ,				unittest.TestCase		):


				'''simple docstring'''


				_lowerCAmelCase		     =(GPTSanJapaneseForConditionalGeneration,) if is_torch_available() else ()
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				_lowerCAmelCase		     =False
				# The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
				_lowerCAmelCase		     =[0.8, 0.9]

				def  _lowerCamelCase   (				self				:       Dict ):
										snake_case__    :   Optional[int]							     =   GPTSanJapaneseTester(self )
										snake_case__    :   List[Any]							     =   ConfigTester(self ,config_class=SCREAMING_SNAKE_CASE_ ,d_model=3_7 )

				def  _lowerCamelCase   (				self				:       Any ):
										GPTSanJapaneseConfig()

				def  _lowerCamelCase   (				self				:       Optional[int] ):
										snake_case__    :   int							     =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_ )

				@unittest.skip(
				    reason='''skip for now as the computed `max_memory` by `model_split_percents` in the test method will be changed inside `from_pretrained`''' )
				def  _lowerCamelCase   (				self				:       Union[str, Any] ):
										super().test_model_parallelism()

				@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''' )
				def  _lowerCamelCase   (				self				:       int ):
										pass

				@slow
				def  _lowerCamelCase   (				self				:       Dict ):
										snake_case__    :   List[Any]							     =   GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   List[str]							     =   GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   List[Any]							     =   tokenizer.encode('''武田信玄は''' ,return_tensors='''pt''' )
										snake_case__    :   str							     =   model(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   List[str]							     =   outputs.logits.detach().cpu().numpy()
										# Output of original model created with mesh-tensoflow
										snake_case__    :   str							     =   [
										    # fmt: off
										    [-12.037839889526367, -12.433061599731445, -14.333840370178223, -12.450345993041992, -11.1661376953125,
										    -11.930137634277344, -10.659740447998047, -12.909574508666992, -13.241043090820312, -13.398579597473145,
										    -11.107524871826172, -12.3685941696167, -22.97943115234375, -10.481067657470703, -12.484030723571777,
										    -12.807360649108887, -14.769700050354004, -12.233579635620117, -13.428145408630371, -22.624177932739258],
										    [-7.511149883270264, -8.281851768493652, -7.943127155303955, -7.55021333694458, -6.49869966506958,
										    -7.586796283721924, -6.978085994720459, -7.839145183563232, -8.21964168548584, -8.695091247558594,
										    -6.706910610198975, -6.6585798263549805, -19.565698623657227, -5.353842735290527, -8.350686073303223,
										    -8.039388656616211, -10.856569290161133, -7.75154447555542, -8.819022178649902, -19.51532745361328],
										    [-9.73066234588623, -10.223922729492188, -9.932981491088867, -11.857836723327637, -7.662626266479492,
										    -11.13529109954834, -7.765097618103027, -11.472923278808594, -9.543149948120117, -11.905633926391602,
										    -9.366164207458496, -11.5734281539917, -23.699003219604492, -9.429590225219727, -10.42839241027832,
										    -10.585240364074707, -10.94771957397461, -11.095416069030762, -10.390240669250488, -23.769372940063477],
										    [-9.728265762329102, -9.859712600708008, -10.09729290008545, -9.678522109985352, -6.879519939422607,
										    -9.68487548828125, -4.2803425788879395, -10.018914222717285, -9.308445930480957, -10.63394546508789,
										    -8.083646774291992, -9.06301498413086, -21.904266357421875, -8.90160846710205, -8.841876029968262,
										    -11.856719970703125, -12.079398155212402, -11.233753204345703, -10.177338600158691, -21.87256622314453],
										    [-9.669764518737793, -9.614198684692383, -9.814510345458984, -9.996501922607422, -11.375690460205078,
										    -10.113405227661133, -10.546867370605469, -10.04369068145752, -10.907809257507324, -10.504216194152832,
										    -11.129199028015137, -10.151124000549316, -21.96586799621582, -9.086349487304688, -11.730339050292969,
										    -10.460667610168457, -10.298049926757812, -10.784148216247559, -10.840693473815918, -22.03152847290039],
										    # fmt: on
										]
										snake_case__    :   Optional[Any]							     =   np.array(SCREAMING_SNAKE_CASE_ ).flatten()
										snake_case__    :   Any							     =   output_logits[0, :, :2_0].flatten()

										def check(SCREAMING_SNAKE_CASE_				:       str ,SCREAMING_SNAKE_CASE_				:       Tuple ,SCREAMING_SNAKE_CASE_				:       Optional[Any]=5E-4 ):
																return abs(a - b ) < epsilon * max(abs(SCREAMING_SNAKE_CASE_ ) ,abs(SCREAMING_SNAKE_CASE_ ) )

										self.assertTrue(np.all([check(target[i] ,predict[i] ) for i in range(len(SCREAMING_SNAKE_CASE_ ) )] ) )

				@slow
				def  _lowerCamelCase   (				self				:       Tuple ):
										snake_case__    :   Tuple							     =   GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   int							     =   GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										model.to(SCREAMING_SNAKE_CASE_ )

										# set deterministically
										snake_case__    :   str							     =   GenerationConfig.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   Any							     =   1

										# use different length sentences to test batching
										snake_case__    :   str							     =   [
										    '''甲斐なら武田と言うほど''',
										    '''織田信長は、''',
										]

										snake_case__    :   List[Any]							     =   '''left'''
										snake_case__    :   List[Any]							     =   tokenizer(SCREAMING_SNAKE_CASE_ ,return_tensors='''pt''' ,padding=SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Union[str, Any]							     =   inputs['''input_ids'''].to(SCREAMING_SNAKE_CASE_ )

										self.assertNotEqual(inputs['''attention_mask'''][0].numpy().tolist() ,inputs['''attention_mask'''][1].numpy().tolist() )

										snake_case__    :   Union[str, Any]							     =   model.generate(
										    input_ids=SCREAMING_SNAKE_CASE_ ,attention_mask=inputs['''attention_mask'''].to(SCREAMING_SNAKE_CASE_ ) ,max_new_tokens=3 ,generation_config=SCREAMING_SNAKE_CASE_ ,)

										snake_case__    :   Optional[int]							     =   tokenizer(sentences[0] ,return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   List[str]							     =   model.generate(
										    input_ids=SCREAMING_SNAKE_CASE_ ,max_new_tokens=3 ,generation_config=SCREAMING_SNAKE_CASE_ )

										snake_case__    :   str							     =   tokenizer(sentences[1] ,return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Tuple							     =   model.generate(input_ids=SCREAMING_SNAKE_CASE_ ,max_new_tokens=3 ,generation_config=SCREAMING_SNAKE_CASE_ )

										self.assertNotEqual(inputs_non_padded.shape ,inputs_padded.shape )

										snake_case__    :   Optional[Any]							     =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE_ ,skip_special_tokens=SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Optional[int]							     =   tokenizer.decode(output_non_padded[0] ,skip_special_tokens=SCREAMING_SNAKE_CASE_ )
										snake_case__    :   Tuple							     =   tokenizer.decode(output_padded[0] ,skip_special_tokens=SCREAMING_SNAKE_CASE_ )

										snake_case__    :   Any							     =   [
										    '''甲斐なら武田と言うほど甲斐の武田''',
										    '''織田信長は、このような''',
										]
										self.assertListEqual(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ )
										self.assertListEqual(SCREAMING_SNAKE_CASE_ ,[non_padded_sentence, padded_sentence] )

				@tooslow
				def  _lowerCamelCase   (				self				:       Any ):
										snake_case__    :   Optional[int]							     =   GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   Union[str, Any]							     =   GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										# Output of original model created with mesh-tensoflow
										snake_case__    :   Any							     =   [
										    ('''武田信玄は''', 3_5_6_7_5),
										    ('''武田信玄は、''', 4_5),
										    ('''武田信玄は、この''', 2_9),
										    ('''武田信玄は、このよう''', 3_0_6_4_2),
										    ('''武田信玄は、このような''', 3_5_6_8_0),
										    ('''武田信玄は、このような「''', 8_6_4_0),
										    ('''武田信玄は、このような「武田''', 3_1_6_1_7),
										    ('''武田信玄は、このような「武田家''', 3_0_6_4_6),
										    ('''武田信玄は、このような「武田家の''', 3_1_6_1_7),
										    ('''武田信玄は、このような「武田家の家''', 3_1_3_8_1),
										]
										for input, output in target:
																snake_case__    :   Optional[Any]							     =   tokenizer.encode(SCREAMING_SNAKE_CASE_ ,return_tensors='''pt''' )
																snake_case__    :   Optional[Any]							     =   model(SCREAMING_SNAKE_CASE_ )
																snake_case__    :   int							     =   outputs.logits.detach().cpu().numpy()[0]
																snake_case__    :   int							     =   np.argmax(output_logits[-1] )
																self.assertEqual(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ )

				@slow
				def  _lowerCamelCase   (				self				:       str ):
										snake_case__    :   List[str]							     =   GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   Tuple							     =   GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										model.to(SCREAMING_SNAKE_CASE_ )

										# set deterministically
										snake_case__    :   Optional[int]							     =   GenerationConfig.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   Union[str, Any]							     =   1

										snake_case__    :   Optional[Any]							     =   '''武田信玄は、'''
										snake_case__    :   Any							     =   tokenizer(SCREAMING_SNAKE_CASE_ ,return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE_ )
										snake_case__    :   str							     =   tokenizer([input_text, input_text] ,return_tensors='''pt''' ).input_ids.to(SCREAMING_SNAKE_CASE_ )

										# spout from uniform and one-hot
										snake_case__    :   List[Any]							     =   [
										    # fmt: off
										    [0.87882208, 0.38426396, 0.33220248, 0.43890406, 0.16562252,
										    0.04803985, 0.211572  , 0.23188473, 0.37153068, 0.7836377 ,
										    0.02160172, 0.38761719, 0.75290772, 0.90198857, 0.34365777,
										    0.64168169, 0.44318471, 0.14575746, 0.92562881, 0.40812148,
										    0.29019122, 0.88861599, 0.65524846, 0.43563456, 0.38177187,
										    0.70832965, 0.81527892, 0.68832812, 0.38833192, 0.4561522 ,
										    0.14828817, 0.47248213, 0.54357335, 0.82009566, 0.1338884 ,
										    0.02755417, 0.19764677, 0.2422084 , 0.04757674, 0.65409606,
										    0.0824589 , 0.03304383, 0.94387689, 0.98764509, 0.82433901,
										    0.27646741, 0.64907493, 0.76009406, 0.30087915, 0.17904689,
										    0.41601714, 0.67046398, 0.10422822, 0.08447374, 0.07354344,
										    0.61423565, 0.70284866, 0.7532333 , 0.1972038 , 0.29575659,
										    0.90583886, 0.29265307, 0.50000175, 0.70407655, 0.889363  ,
										    0.81904418, 0.66829128, 0.64468815, 0.56563723, 0.85601875,
										    0.94924672, 0.00166762, 0.25220643, 0.74540219, 0.67993247,
										    0.1549675 , 0.39385352, 0.92153607, 0.63745931, 0.27759043,
										    0.84702295, 0.65904271, 0.58676614, 0.8666936 , 0.39607438,
										    0.79954983, 0.42220697, 0.39650381, 0.7849864 , 0.56150201,
										    0.15678925, 0.14746032, 0.34542114, 0.47026783, 0.11956489,
										    0.25421435, 0.33788901, 0.68934842, 0.36424685, 0.71737898,
										    0.38983449, 0.94393779, 0.39575588, 0.36616553, 0.87104665,
										    0.64630203, 0.22516905, 0.88270804, 0.15031338, 0.75144345,
										    0.46459025, 0.85396454, 0.86355643, 0.65139851, 0.70266061,
										    0.30241389, 0.81056497, 0.88865969, 0.38773807, 0.70635849,
										    0.90718459, 0.43245789, 0.28000654, 0.45935562, 0.08773519,
										    0.9552151 , 0.93901511, 0.22489288], # uniform
										    [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
										     0., 0., 0., 0., 0., 0., 0., 0.],
										    # fmt: on
										]

										snake_case__    :   List[str]							     =   model.generate(
										    input_ids=SCREAMING_SNAKE_CASE_ ,spout=spouts[0] ,max_new_tokens=2_0 ,generation_config=SCREAMING_SNAKE_CASE_ ,)

										snake_case__    :   Any							     =   model.generate(
										    input_ids=SCREAMING_SNAKE_CASE_ ,spout=spouts[1] ,max_new_tokens=2_0 ,generation_config=SCREAMING_SNAKE_CASE_ ,)

										snake_case__    :   List[str]							     =   model.generate(
										    input_ids=SCREAMING_SNAKE_CASE_ ,spout=SCREAMING_SNAKE_CASE_ ,max_new_tokens=2_0 ,generation_config=SCREAMING_SNAKE_CASE_ ,)

										snake_case__    :   Any							     =   tokenizer.decode(outputa[0] )
										snake_case__    :   Union[str, Any]							     =   tokenizer.decode(outputa[0] )
										snake_case__    :   str							     =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE_ )

										snake_case__    :   Optional[int]							     =   [
										    '''武田信玄は、武田氏の滅亡後、武田氏の居城であった甲斐武田氏の居城である''',
										    '''武田信玄は、武田家の滅亡を防ぐため、武田家の家臣である武田信虎を討''',
										]
										self.assertListEqual(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ )
										self.assertListEqual(SCREAMING_SNAKE_CASE_ ,[outa_sentence, outa_sentence] )





				@slow
				def  _lowerCamelCase   (				self				:       Optional[int] ):
										snake_case__    :   Union[str, Any]							     =   GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   int							     =   GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										model.to(SCREAMING_SNAKE_CASE_ )

										# set deterministically
										snake_case__    :   Union[str, Any]							     =   GenerationConfig.from_pretrained('''Tanrei/GPTSAN-japanese''' )
										snake_case__    :   Union[str, Any]							     =   1

										snake_case__    :   Tuple							     =   '''武田信玄'''
										snake_case__    :   int							     =   '''織田信長'''
										snake_case__    :   List[Any]							     =   '''は、'''
										snake_case__    :   List[Any]							     =   '''が、'''
										snake_case__    :   Any							     =   tokenizer(SCREAMING_SNAKE_CASE_ ,prefix_text=SCREAMING_SNAKE_CASE_ ,return_tensors='''pt''' )
										snake_case__    :   int							     =   tokenizer(SCREAMING_SNAKE_CASE_ ,prefix_text=SCREAMING_SNAKE_CASE_ ,return_tensors='''pt''' )
										snake_case__    :   Tuple							     =   tokenizer([[prefix_text_a, input_text_a], [prefix_text_a, input_text_a]] ,return_tensors='''pt''' )

										snake_case__    :   Union[str, Any]							     =   model.generate(
										    input_ids=input_tok_a.input_ids.to(SCREAMING_SNAKE_CASE_ ) ,token_type_ids=input_tok_a.token_type_ids.to(SCREAMING_SNAKE_CASE_ ) ,max_new_tokens=2_0 ,generation_config=SCREAMING_SNAKE_CASE_ ,)

										snake_case__    :   Any							     =   model.generate(
										    input_ids=input_tok_a.input_ids.to(SCREAMING_SNAKE_CASE_ ) ,token_type_ids=input_tok_a.token_type_ids.to(SCREAMING_SNAKE_CASE_ ) ,max_new_tokens=2_0 ,generation_config=SCREAMING_SNAKE_CASE_ ,)

										snake_case__    :   Tuple							     =   model.generate(
										    input_ids=input_tok_a.input_ids.to(SCREAMING_SNAKE_CASE_ ) ,token_type_ids=input_tok_a.token_type_ids.to(SCREAMING_SNAKE_CASE_ ) ,attention_mask=input_tok_a.attention_mask.to(SCREAMING_SNAKE_CASE_ ) ,max_new_tokens=2_0 ,generation_config=SCREAMING_SNAKE_CASE_ ,)

										snake_case__    :   Tuple							     =   tokenizer.decode(outputa[0] )
										snake_case__    :   Optional[int]							     =   tokenizer.decode(outputa[0] )
										snake_case__    :   Any							     =   tokenizer.batch_decode(SCREAMING_SNAKE_CASE_ )

										snake_case__    :   List[Any]							     =   [
										    '''武田信玄は、武田氏の祖である武田信虎を、その子・武田信友を擁して''',
										    '''織田信長が、織田信長の妻・お市の方を妻として迎えたという逸話が残''',
										]
										self.assertListEqual(SCREAMING_SNAKE_CASE_ ,SCREAMING_SNAKE_CASE_ )
										self.assertListEqual(SCREAMING_SNAKE_CASE_ ,[outa_sentence, outa_sentence] )








"""simple docstring"""





import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing the checkpointing capability,
# and builds off the `nlp_example.py` script.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To help focus on the differences in the code, building `DataLoaders`
# was refactored into its own function.
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################

_lowercase		 =      16
_lowercase		 =      32

def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase = 16						):
						'''simple docstring'''


						snake_case__    :   int							     =   AutoTokenizer.from_pretrained('''bert-base-cased'''						)
						snake_case__    :   Dict							     =   load_dataset('''glue''' ,     '''mrpc'''						)

						def tokenize_function(lowerCamelCase						):
												# max_length=None => use the model max length (it's actually the default)
												snake_case__    :   Optional[int]							     =   tokenizer(examples['''sentence1'''] ,     examples['''sentence2'''] ,     truncation=lowerCamelCase ,     max_length=lowerCamelCase						)
												return outputs

						# Apply the method we just defined to all the examples in all the splits of the dataset
						# starting with the main process first:
						with accelerator.main_process_first():
												snake_case__    :   Tuple							     =   datasets.map(
												    lowerCamelCase ,     batched=lowerCamelCase ,     remove_columns=['''idx''', '''sentence1''', '''sentence2'''] ,     )

						# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
						# transformers library
						snake_case__    :   str							     =   tokenized_datasets.rename_column('''label''' ,     '''labels'''						)

						def collate_fn(lowerCamelCase						):
												# On TPU it's best to pad everything to the same length or training will be very slow.
												snake_case__    :   Tuple							     =   128 if accelerator.distributed_type == DistributedType.TPU else None
												# When using mixed precision we want round multiples of 8/16
												if accelerator.mixed_precision == "fp8":
																		snake_case__    :   Optional[Any]							     =   16
												elif accelerator.mixed_precision != "no":
																		snake_case__    :   Dict							     =   8
												else:
																		snake_case__    :   str							     =   None

												return tokenizer.pad(
												    lowerCamelCase ,     padding='''longest''' ,     max_length=lowerCamelCase ,     pad_to_multiple_of=lowerCamelCase ,     return_tensors='''pt''' ,     )

						# Instantiate dataloaders.
						snake_case__    :   List[Any]							     =   DataLoader(
						    tokenized_datasets['''train'''] ,     shuffle=lowerCamelCase ,     collate_fn=lowerCamelCase ,     batch_size=lowerCamelCase						)
						snake_case__    :   Tuple							     =   DataLoader(
						    tokenized_datasets['''validation'''] ,     shuffle=lowerCamelCase ,     collate_fn=lowerCamelCase ,     batch_size=lowerCamelCase						)

						return train_dataloader, eval_dataloader


# For testing only
if os.environ.get("""TESTING_MOCKED_DATALOADERS""", None) == "1":
							from accelerate.test_utils.training import mocked_dataloaders

							_lowercase		 =      mocked_dataloaders  # noqa: F811

def 						lowerCAmelCase							(	lowerCamelCase ,     lowerCamelCase						):
						'''simple docstring'''


						if os.environ.get('''TESTING_MOCKED_DATALOADERS''' ,     lowerCamelCase						) == "1":
												snake_case__    :   Any							     =   2
						# Initialize accelerator
						snake_case__    :   List[Any]							     =   Accelerator(cpu=args.cpu ,     mixed_precision=args.mixed_precision						)
						# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
						snake_case__    :   Any							     =   config['''lr''']
						snake_case__    :   int							     =   int(config['''num_epochs''']						)
						snake_case__    :   List[str]							     =   int(config['''seed''']						)
						snake_case__    :   Dict							     =   int(config['''batch_size''']						)

						# New Code #
						# Parse out whether we are saving every epoch or after a certain number of batches
						if hasattr(args.checkpointing_steps ,     '''isdigit'''						):
												if args.checkpointing_steps == "epoch":
																		snake_case__    :   Optional[Any]							     =   args.checkpointing_steps
												elif args.checkpointing_steps.isdigit():
																		snake_case__    :   Union[str, Any]							     =   int(args.checkpointing_steps						)
												else:
																		raise ValueError(
																		    f'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''						)
						else:
												snake_case__    :   Union[str, Any]							     =   None

						set_seed(lowerCamelCase						)

						snake_case__     ,		snake_case__    :   int							     =   get_dataloaders(lowerCamelCase ,     lowerCamelCase						)
						snake_case__    :   Optional[Any]							     =   evaluate.load('''glue''' ,     '''mrpc'''						)

						# If the batch size is too big we use gradient accumulation
						snake_case__    :   Optional[int]							     =   1
						if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
												snake_case__    :   Any							     =   batch_size // MAX_GPU_BATCH_SIZE
												snake_case__    :   Union[str, Any]							     =   MAX_GPU_BATCH_SIZE

						# Instantiate the model (we build the model here so that the seed also control new weights initialization)
						snake_case__    :   Tuple							     =   AutoModelForSequenceClassification.from_pretrained('''bert-base-cased''' ,     return_dict=lowerCamelCase						)

						# We could avoid this line since the accelerator is set with `device_placement=True` (default value).
						# Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
						# creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
						snake_case__    :   Dict							     =   model.to(accelerator.device						)

						# Instantiate optimizer
						snake_case__    :   Optional[int]							     =   AdamW(params=model.parameters() ,     lr=lowerCamelCase						)

						# Instantiate scheduler
						snake_case__    :   Optional[Any]							     =   get_linear_schedule_with_warmup(
						    optimizer=lowerCamelCase ,     num_warmup_steps=100 ,     num_training_steps=(len(lowerCamelCase						) * num_epochs) // gradient_accumulation_steps ,     )

						# Prepare everything
						# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
						# prepare method.
						snake_case__     ,		snake_case__     ,		snake_case__     ,		snake_case__     ,		snake_case__    :   Optional[Any]							     =   accelerator.prepare(
						    lowerCamelCase ,     lowerCamelCase ,     lowerCamelCase ,     lowerCamelCase ,     lowerCamelCase						)

						# New Code #
						# We need to keep track of how many total steps we have iterated over
						snake_case__    :   Any							     =   0
						# We also need to keep track of the stating epoch so files are named properly
						snake_case__    :   List[str]							     =   0

						# We need to load the checkpoint back in before training here with `load_state`
						# The total number of epochs is adjusted based on where the state is being loaded from,
						# as we assume continuation of the same training script
						if args.resume_from_checkpoint:
												if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
																		accelerator.print(f'''Resumed from checkpoint: {args.resume_from_checkpoint}'''						)
																		accelerator.load_state(args.resume_from_checkpoint						)
																		snake_case__    :   Optional[Any]							     =   os.path.basename(args.resume_from_checkpoint						)
												else:
																		# Get the most recent checkpoint
																		snake_case__    :   Any							     =   [f.name for f in os.scandir(os.getcwd()						) if f.is_dir()]
																		dirs.sort(key=os.path.getctime						)
																		snake_case__    :   str							     =   dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
												# Extract `epoch_{i}` or `step_{i}`
												snake_case__    :   List[Any]							     =   os.path.splitext(lowerCamelCase						)[0]

												if "epoch" in training_difference:
																		snake_case__    :   Optional[Any]							     =   int(training_difference.replace('''epoch_''' ,     ''''''						)						) + 1
																		snake_case__    :   List[str]							     =   None
												else:
																		snake_case__    :   int							     =   int(training_difference.replace('''step_''' ,     ''''''						)						)
																		snake_case__    :   Union[str, Any]							     =   resume_step // len(lowerCamelCase						)
																		resume_step -= starting_epoch * len(lowerCamelCase						)

    # Now we train the model
						for epoch in range(lowerCamelCase ,     lowerCamelCase						):
												model.train()
												# New Code #
												if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
																		# We need to skip steps until we reach the resumed step
																		snake_case__    :   int							     =   accelerator.skip_first_batches(lowerCamelCase ,     lowerCamelCase						)
																		overall_step += resume_step
												else:
																		# After the first iteration though, we need to go back to the original dataloader
																		snake_case__    :   Tuple							     =   train_dataloader
												for step, batch in enumerate(lowerCamelCase						):
																		# We could avoid this line since we set the accelerator with `device_placement=True`.
																		batch.to(accelerator.device						)
																		snake_case__    :   Tuple							     =   model(**lowerCamelCase						)
																		snake_case__    :   Optional[int]							     =   outputs.loss
																		snake_case__    :   Union[str, Any]							     =   loss / gradient_accumulation_steps
																		accelerator.backward(lowerCamelCase						)
																		if step % gradient_accumulation_steps == 0:
																								optimizer.step()
																								lr_scheduler.step()
																								optimizer.zero_grad()
																		# New Code #
																		overall_step += 1

																		# New Code #
																		# We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
																		# These are saved to folders named `step_{overall_step}`
																		# Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
																		# If mixed precision was used, will also save a "scalar.bin" file
																		if isinstance(lowerCamelCase ,     lowerCamelCase						):
																								snake_case__    :   Tuple							     =   f'''step_{overall_step}'''
																								if overall_step % checkpointing_steps == 0:
																														if args.output_dir is not None:
																																				snake_case__    :   Optional[int]							     =   os.path.join(args.output_dir ,     lowerCamelCase						)
																														accelerator.save_state(lowerCamelCase						)

												model.eval()
												for step, batch in enumerate(lowerCamelCase						):
																		# We could avoid this line since we set the accelerator with `device_placement=True` (the default).
																		batch.to(accelerator.device						)
																		with torch.no_grad():
																								snake_case__    :   List[str]							     =   model(**lowerCamelCase						)
																		snake_case__    :   Optional[int]							     =   outputs.logits.argmax(dim=-1						)
																		snake_case__     ,		snake_case__    :   Dict							     =   accelerator.gather_for_metrics((predictions, batch['''labels'''])						)
																		metric.add_batch(
																		    predictions=lowerCamelCase ,     references=lowerCamelCase ,     )

												snake_case__    :   str							     =   metric.compute()
												# Use accelerator.print to print only on the main process.
												accelerator.print(f'''epoch {epoch}:''' ,     lowerCamelCase						)

												# New Code #
												# We save the model, optimizer, lr_scheduler, and seed states by calling `save_state`
												# These are saved to folders named `epoch_{epoch}`
												# Will contain files: "pytorch_model.bin", "optimizer.bin", "scheduler.bin", and "random_states.pkl"
												# If mixed precision was used, will also save a "scalar.bin" file
												if checkpointing_steps == "epoch":
																		snake_case__    :   str							     =   f'''epoch_{epoch}'''
																		if args.output_dir is not None:
																								snake_case__    :   Optional[int]							     =   os.path.join(args.output_dir ,     lowerCamelCase						)
																		accelerator.save_state(lowerCamelCase						)

def 						lowerCAmelCase							(	):
						'''simple docstring'''


						snake_case__    :   Tuple							     =   argparse.ArgumentParser(description='''Simple example of training script.'''						)
						parser.add_argument(
						    '''--mixed_precision''' ,     type=lowerCamelCase ,     default=lowerCamelCase ,     choices=['''no''', '''fp16''', '''bf16''', '''fp8'''] ,     help='''Whether to use mixed precision. Choose'''
						    '''between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10.'''
						    '''and an Nvidia Ampere GPU.''' ,     )
						parser.add_argument('''--cpu''' ,     action='''store_true''' ,     help='''If passed, will train on the CPU.'''						)
						parser.add_argument(
						    '''--checkpointing_steps''' ,     type=lowerCamelCase ,     default=lowerCamelCase ,     help='''Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch.''' ,     )
						parser.add_argument(
						    '''--output_dir''' ,     type=lowerCamelCase ,     default='''.''' ,     help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.''' ,     )
						parser.add_argument(
						    '''--resume_from_checkpoint''' ,     type=lowerCamelCase ,     default=lowerCamelCase ,     help='''If the training should continue from a checkpoint folder.''' ,     )
						snake_case__    :   List[str]							     =   parser.parse_args()
						snake_case__    :   str							     =   {'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
						training_function(lowerCamelCase ,     lowerCamelCase						)


if __name__ == "__main__":
							main()


