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



			"""simple docstring"""






			def __init__(       self							, snake_case_							, snake_case_=36_000							, snake_case_=13							, snake_case_=7							, snake_case_=True							, snake_case_=32							, snake_case_=42							, snake_case_=5							, snake_case_=2							, snake_case_=4							, snake_case_=2							, snake_case_=32							, snake_case_=80							, snake_case_=33							, snake_case_=0.0							, snake_case_=1e-6							, snake_case_=100							, snake_case_=0.0							, )     ->   Optional[Any]:
									snake_case   :					Optional[Any]							=						vocab_size
									snake_case   :					Any							=						parent
									snake_case   :					Any							=						batch_size
									snake_case   :					Union[str, Any]							=						num_contexts
									# For common tests
									snake_case   :					int							=						self.num_contexts
									snake_case   :					Optional[int]							=						is_training
									snake_case   :					str							=						hidden_size
									snake_case   :					str							=						num_ext_layers
									snake_case   :					str							=						ext_size
									snake_case   :					str							=						num_hidden_layers
									snake_case   :					Optional[Any]							=						num_attention_heads
									snake_case   :					Optional[int]							=						num_experts
									snake_case   :					Dict							=						d_ff
									snake_case   :					Optional[Any]							=						d_ext
									snake_case   :					Optional[Any]							=						d_spout
									snake_case   :					Dict							=						dropout_rate
									snake_case   :					Union[str, Any]							=						layer_norm_epsilon
									snake_case   :					str							=						expert_capacity
									snake_case   :					Optional[int]							=						router_jitter_noise



			def      lowerCamelCase   (       self				)     ->   int:
									return GPTSanJapaneseConfig.from_pretrained("Tanrei/GPTSAN-japanese"				)



			def      lowerCamelCase   (       self				)     ->   Optional[int]:
									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				)     ->   Optional[Any]:
									snake_case   :					Tuple							=						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				)     ->   Dict:
									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							, snake_case_							, snake_case_							, )     ->   List[Any]:
									snake_case   :					Any							=						GPTSanJapaneseForConditionalGeneration(config=lowerCamelCase__				)
									model.to(lowerCamelCase__				)
									model.eval()
									snake_case   :					Optional[int]							=						model(
									    input_ids=lowerCamelCase__							, )
									self.parent.assertIsNotNone(lowerCamelCase__				)







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



			"""simple docstring"""






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



			def      lowerCamelCase   (       self							, snake_case_							, snake_case_							, snake_case_							, snake_case_							, snake_case_				)     ->   Tuple:
									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   :					List[Any]							=						GPTSanJapaneseTester(self				)
									snake_case   :					int							=						ConfigTester(self							, config_class=lowerCamelCase__							, d_model=37				)



			def      lowerCamelCase   (       self				)     ->   Union[str, Any]:
									GPTSanJapaneseConfig()



			def      lowerCamelCase   (       self				)     ->   Tuple:
									snake_case   :					int							=						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowerCamelCase__				)



			@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				)     ->   List[str]:
									super().test_model_parallelism()






			@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."				)
			def      lowerCamelCase   (       self				)     ->   List[str]:
									pass







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



			"""simple docstring"""






			__magic_name__   :							int             =					(GPTSanJapaneseForConditionalGeneration,) if is_torch_available() else ()
			__magic_name__   :							Optional[Any]             =					False
			__magic_name__   :							Dict             =					False
			__magic_name__   :							Dict             =					False
			__magic_name__   :							Optional[Any]             =					False
			__magic_name__   :							int             =					False
			__magic_name__   :							int             =					False
			# The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
			__magic_name__   :							List[Any]             =					[0.8, 0.9]



			def      lowerCamelCase   (       self				)     ->   List[Any]:
									snake_case   :					str							=						GPTSanJapaneseTester(self				)
									snake_case   :					List[str]							=						ConfigTester(self							, config_class=lowerCamelCase__							, d_model=37				)



			def      lowerCamelCase   (       self				)     ->   List[str]:
									GPTSanJapaneseConfig()



			def      lowerCamelCase   (       self				)     ->   Optional[int]:
									snake_case   :					Any							=						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*lowerCamelCase__				)



			@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				)     ->   Dict:
									super().test_model_parallelism()



			@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."				)
			def      lowerCamelCase   (       self				)     ->   Optional[int]:
									pass



			@slow
			def      lowerCamelCase   (       self				)     ->   Optional[Any]:
									snake_case   :					str							=						GPTSanJapaneseForConditionalGeneration.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					Union[str, Any]							=						GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					List[str]							=						tokenizer.encode("武田信玄は"							, return_tensors="pt"				)
									snake_case   :					List[str]							=						model(lowerCamelCase__				)
									snake_case   :					Any							=						outputs.logits.detach().cpu().numpy()
									# Output of original model created with mesh-tensoflow
									snake_case   :					int							=						[
									    # fmt: off
									    [-12.037_839_889_526_367, -12.433_061_599_731_445, -14.333_840_370_178_223, -12.450_345_993_041_992, -11.1_661_376_953_125,
									    -11.930_137_634_277_344, -10.659_740_447_998_047, -12.909_574_508_666_992, -13.241_043_090_820_312, -13.398_579_597_473_145,
									    -11.107_524_871_826_172, -12.3_685_941_696_167, -22.97_943_115_234_375, -10.481_067_657_470_703, -12.484_030_723_571_777,
									    -12.807_360_649_108_887, -14.769_700_050_354_004, -12.233_579_635_620_117, -13.428_145_408_630_371, -22.624_177_932_739_258],
									    [-7.511_149_883_270_264, -8.281_851_768_493_652, -7.943_127_155_303_955, -7.55_021_333_694_458, -6.49_869_966_506_958,
									    -7.586_796_283_721_924, -6.978_085_994_720_459, -7.839_145_183_563_232, -8.21_964_168_548_584, -8.695_091_247_558_594,
									    -6.706_910_610_198_975, -6.6_585_798_263_549_805, -19.565_698_623_657_227, -5.353_842_735_290_527, -8.350_686_073_303_223,
									    -8.039_388_656_616_211, -10.856_569_290_161_133, -7.75_154_447_555_542, -8.819_022_178_649_902, -19.51_532_745_361_328],
									    [-9.73_066_234_588_623, -10.223_922_729_492_188, -9.932_981_491_088_867, -11.857_836_723_327_637, -7.662_626_266_479_492,
									    -11.13_529_109_954_834, -7.765_097_618_103_027, -11.472_923_278_808_594, -9.543_149_948_120_117, -11.905_633_926_391_602,
									    -9.366_164_207_458_496, -11.5_734_281_539_917, -23.699_003_219_604_492, -9.429_590_225_219_727, -10.42_839_241_027_832,
									    -10.585_240_364_074_707, -10.94_771_957_397_461, -11.095_416_069_030_762, -10.390_240_669_250_488, -23.769_372_940_063_477],
									    [-9.728_265_762_329_102, -9.859_712_600_708_008, -10.09_729_290_008_545, -9.678_522_109_985_352, -6.879_519_939_422_607,
									    -9.68_487_548_828_125, -4.2_803_425_788_879_395, -10.018_914_222_717_285, -9.308_445_930_480_957, -10.63_394_546_508_789,
									    -8.083_646_774_291_992, -9.06_301_498_413_086, -21.904_266_357_421_875, -8.90_160_846_710_205, -8.841_876_029_968_262,
									    -11.856_719_970_703_125, -12.079_398_155_212_402, -11.233_753_204_345_703, -10.177_338_600_158_691, -21.87_256_622_314_453],
									    [-9.669_764_518_737_793, -9.614_198_684_692_383, -9.814_510_345_458_984, -9.996_501_922_607_422, -11.375_690_460_205_078,
									    -10.113_405_227_661_133, -10.546_867_370_605_469, -10.04_369_068_145_752, -10.907_809_257_507_324, -10.504_216_194_152_832,
									    -11.129_199_028_015_137, -10.151_124_000_549_316, -21.96_586_799_621_582, -9.086_349_487_304_688, -11.730_339_050_292_969,
									    -10.460_667_610_168_457, -10.298_049_926_757_812, -10.784_148_216_247_559, -10.840_693_473_815_918, -22.03_152_847_290_039],
									    # fmt: on
									]
									snake_case   :					Union[str, Any]							=						np.array(lowerCamelCase__				).flatten()
									snake_case   :					int							=						output_logits[0, :, :20].flatten()

									def check(snake_case_							, snake_case_							, snake_case_=5e-4				):
															return abs(a - b				) < epsilon * max(abs(lowerCamelCase__				)							, abs(lowerCamelCase__				)				)

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



			@slow
			def      lowerCamelCase   (       self				)     ->   str:
									snake_case   :					Optional[int]							=						GPTSanJapaneseForConditionalGeneration.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					Tuple							=						GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese"				)
									model.to(lowerCamelCase__				)

									# set deterministically
									snake_case   :					Optional[int]							=						GenerationConfig.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					Dict							=						1

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

									snake_case   :					Tuple							=						"left"
									snake_case   :					Optional[Any]							=						tokenizer(lowerCamelCase__							, return_tensors="pt"							, padding=lowerCamelCase__				)
									snake_case   :					List[str]							=						inputs["input_ids"].to(lowerCamelCase__				)

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

									snake_case   :					Any							=						model.generate(
									    input_ids=lowerCamelCase__							, attention_mask=inputs["attention_mask"].to(lowerCamelCase__				)							, max_new_tokens=3							, generation_config=lowerCamelCase__							, )

									snake_case   :					Dict							=						tokenizer(sentences[0]							, return_tensors="pt"				).input_ids.to(lowerCamelCase__				)
									snake_case   :					Optional[Any]							=						model.generate(
									    input_ids=lowerCamelCase__							, max_new_tokens=3							, generation_config=lowerCamelCase__				)

									snake_case   :					List[str]							=						tokenizer(sentences[1]							, return_tensors="pt"				).input_ids.to(lowerCamelCase__				)
									snake_case   :					int							=						model.generate(input_ids=lowerCamelCase__							, max_new_tokens=3							, generation_config=lowerCamelCase__				)

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

									snake_case   :					Tuple							=						tokenizer.batch_decode(lowerCamelCase__							, skip_special_tokens=lowerCamelCase__				)
									snake_case   :					List[Any]							=						tokenizer.decode(output_non_padded[0]							, skip_special_tokens=lowerCamelCase__				)
									snake_case   :					Union[str, Any]							=						tokenizer.decode(output_padded[0]							, skip_special_tokens=lowerCamelCase__				)

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



			@tooslow
			def      lowerCamelCase   (       self				)     ->   List[Any]:
									snake_case   :					Dict							=						GPTSanJapaneseForConditionalGeneration.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					str							=						GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese"				)
									# Output of original model created with mesh-tensoflow
									snake_case   :					List[Any]							=						[
									    ("武田信玄は", 35_675),
									    ("武田信玄は、", 45),
									    ("武田信玄は、この", 29),
									    ("武田信玄は、このよう", 30_642),
									    ("武田信玄は、このような", 35_680),
									    ("武田信玄は、このような「", 8_640),
									    ("武田信玄は、このような「武田", 31_617),
									    ("武田信玄は、このような「武田家", 30_646),
									    ("武田信玄は、このような「武田家の", 31_617),
									    ("武田信玄は、このような「武田家の家", 31_381),
									]
									for input, output in target:
															snake_case   :					int							=						tokenizer.encode(lowerCamelCase__							, return_tensors="pt"				)
															snake_case   :					List[Any]							=						model(lowerCamelCase__				)
															snake_case   :					Optional[Any]							=						outputs.logits.detach().cpu().numpy()[0]
															snake_case   :					Tuple							=						np.argmax(output_logits[-1]				)
															self.assertEqual(lowerCamelCase__							, lowerCamelCase__				)



			@slow
			def      lowerCamelCase   (       self				)     ->   Union[str, Any]:
									snake_case   :					str							=						GPTSanJapaneseForConditionalGeneration.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					Union[str, Any]							=						GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese"				)
									model.to(lowerCamelCase__				)

									# set deterministically
									snake_case   :					Tuple							=						GenerationConfig.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					Any							=						1

									snake_case   :					Optional[int]							=						"武田信玄は、"
									snake_case   :					Optional[Any]							=						tokenizer(lowerCamelCase__							, return_tensors="pt"				).input_ids.to(lowerCamelCase__				)
									snake_case   :					Dict							=						tokenizer([input_text, input_text]							, return_tensors="pt"				).input_ids.to(lowerCamelCase__				)

									# spout from uniform and one-hot
									snake_case   :					List[str]							=						[
									    # fmt: off
									    [0.87_882_208, 0.38_426_396, 0.33_220_248, 0.43_890_406, 0.16_562_252,
									    0.04_803_985, 0.211_572  , 0.23_188_473, 0.37_153_068, 0.7_836_377 ,
									    0.02_160_172, 0.38_761_719, 0.75_290_772, 0.90_198_857, 0.34_365_777,
									    0.64_168_169, 0.44_318_471, 0.14_575_746, 0.92_562_881, 0.40_812_148,
									    0.29_019_122, 0.88_861_599, 0.65_524_846, 0.43_563_456, 0.38_177_187,
									    0.70_832_965, 0.81_527_892, 0.68_832_812, 0.38_833_192, 0.4_561_522 ,
									    0.14_828_817, 0.47_248_213, 0.54_357_335, 0.82_009_566, 0.1_338_884 ,
									    0.02_755_417, 0.19_764_677, 0.2_422_084 , 0.04_757_674, 0.65_409_606,
									    0.0_824_589 , 0.03_304_383, 0.94_387_689, 0.98_764_509, 0.82_433_901,
									    0.27_646_741, 0.64_907_493, 0.76_009_406, 0.30_087_915, 0.17_904_689,
									    0.41_601_714, 0.67_046_398, 0.10_422_822, 0.08_447_374, 0.07_354_344,
									    0.61_423_565, 0.70_284_866, 0.7_532_333 , 0.1_972_038 , 0.29_575_659,
									    0.90_583_886, 0.29_265_307, 0.50_000_175, 0.70_407_655, 0.889_363  ,
									    0.81_904_418, 0.66_829_128, 0.64_468_815, 0.56_563_723, 0.85_601_875,
									    0.94_924_672, 0.00_166_762, 0.25_220_643, 0.74_540_219, 0.67_993_247,
									    0.1_549_675 , 0.39_385_352, 0.92_153_607, 0.63_745_931, 0.27_759_043,
									    0.84_702_295, 0.65_904_271, 0.58_676_614, 0.8_666_936 , 0.39_607_438,
									    0.79_954_983, 0.42_220_697, 0.39_650_381, 0.7_849_864 , 0.56_150_201,
									    0.15_678_925, 0.14_746_032, 0.34_542_114, 0.47_026_783, 0.11_956_489,
									    0.25_421_435, 0.33_788_901, 0.68_934_842, 0.36_424_685, 0.71_737_898,
									    0.38_983_449, 0.94_393_779, 0.39_575_588, 0.36_616_553, 0.87_104_665,
									    0.64_630_203, 0.22_516_905, 0.88_270_804, 0.15_031_338, 0.75_144_345,
									    0.46_459_025, 0.85_396_454, 0.86_355_643, 0.65_139_851, 0.70_266_061,
									    0.30_241_389, 0.81_056_497, 0.88_865_969, 0.38_773_807, 0.70_635_849,
									    0.90_718_459, 0.43_245_789, 0.28_000_654, 0.45_935_562, 0.08_773_519,
									    0.9_552_151 , 0.93_901_511, 0.22_489_288], # 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   :					str							=						model.generate(
									    input_ids=lowerCamelCase__							, spout=spouts[0]							, max_new_tokens=20							, generation_config=lowerCamelCase__							, )

									snake_case   :					int							=						model.generate(
									    input_ids=lowerCamelCase__							, spout=spouts[1]							, max_new_tokens=20							, generation_config=lowerCamelCase__							, )

									snake_case   :					List[str]							=						model.generate(
									    input_ids=lowerCamelCase__							, spout=lowerCamelCase__							, max_new_tokens=20							, generation_config=lowerCamelCase__							, )

									snake_case   :					Union[str, Any]							=						tokenizer.decode(outputa[0]				)
									snake_case   :					Optional[Any]							=						tokenizer.decode(outputa[0]				)
									snake_case   :					Optional[Any]							=						tokenizer.batch_decode(lowerCamelCase__				)

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






			@slow
			def      lowerCamelCase   (       self				)     ->   str:
									snake_case   :					Optional[Any]							=						GPTSanJapaneseForConditionalGeneration.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					List[Any]							=						GPTSanJapaneseTokenizer.from_pretrained("Tanrei/GPTSAN-japanese"				)
									model.to(lowerCamelCase__				)

									# set deterministically
									snake_case   :					Any							=						GenerationConfig.from_pretrained("Tanrei/GPTSAN-japanese"				)
									snake_case   :					Tuple							=						1

									snake_case   :					Optional[Any]							=						"武田信玄"
									snake_case   :					int							=						"織田信長"
									snake_case   :					Any							=						"は、"
									snake_case   :					int							=						"が、"
									snake_case   :					str							=						tokenizer(lowerCamelCase__							, prefix_text=lowerCamelCase__							, return_tensors="pt"				)
									snake_case   :					List[str]							=						tokenizer(lowerCamelCase__							, prefix_text=lowerCamelCase__							, return_tensors="pt"				)
									snake_case   :					Dict							=						tokenizer([[prefix_text_a, input_text_a], [prefix_text_a, input_text_a]]							, return_tensors="pt"				)

									snake_case   :					Dict							=						model.generate(
									    input_ids=input_tok_a.input_ids.to(lowerCamelCase__				)							, token_type_ids=input_tok_a.token_type_ids.to(lowerCamelCase__				)							, max_new_tokens=20							, generation_config=lowerCamelCase__							, )

									snake_case   :					List[str]							=						model.generate(
									    input_ids=input_tok_a.input_ids.to(lowerCamelCase__				)							, token_type_ids=input_tok_a.token_type_ids.to(lowerCamelCase__				)							, max_new_tokens=20							, generation_config=lowerCamelCase__							, )

									snake_case   :					Optional[int]							=						model.generate(
									    input_ids=input_tok_a.input_ids.to(lowerCamelCase__				)							, token_type_ids=input_tok_a.token_type_ids.to(lowerCamelCase__				)							, attention_mask=input_tok_a.attention_mask.to(lowerCamelCase__				)							, max_new_tokens=20							, generation_config=lowerCamelCase__							, )

									snake_case   :					List[Any]							=						tokenizer.decode(outputa[0]				)
									snake_case   :					Any							=						tokenizer.decode(outputa[0]				)
									snake_case   :					Union[str, Any]							=						tokenizer.batch_decode(lowerCamelCase__				)

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






'''simple docstring'''

import argparse
import json
import os.path
from collections import OrderedDict

import numpy as np
import requests
import torch
from flax.training.checkpoints import restore_checkpoint
from huggingface_hub import hf_hub_download

from transformers import VivitConfig, VivitForVideoClassification, VivitImageProcessor
from transformers.image_utils import PILImageResampling


def 					UpperCAmelCase__      (							_lowerCAmelCase						: Optional[Any])						->		List[Any]:



						"""simple docstring"""

						snake_case   :					List[str]							=						"https://storage.googleapis.com/scenic-bucket/vivit/kinetics_400/vivit_base_16x2_unfactorized/checkpoint"

						with open(_lowerCAmelCase,			"wb") as f:
												with requests.get(_lowerCAmelCase,			stream=_lowerCAmelCase) as req:
																		for chunk in req.iter_content(chunk_size=2048):
																								f.write(_lowerCAmelCase)


def 					UpperCAmelCase__      (							)						->		VivitConfig:



						"""simple docstring"""

						snake_case   :					Any							=						VivitConfig()

						snake_case   :					Optional[int]							=						400
						snake_case   :					Tuple							=						"huggingface/label-files"
						snake_case   :					Tuple							=						"kinetics400-id2label.json"

						snake_case   :					str							=						json.load(open(hf_hub_download(_lowerCAmelCase,			_lowerCAmelCase,			repo_type="dataset"),			"r"))
						snake_case   :					Optional[int]							=						{int(_lowerCAmelCase): v for k, v in idalabel.items()}
						snake_case   :					Union[str, Any]							=						idalabel
						snake_case   :					Tuple							=						{v: k for k, v in idalabel.items()}
						return config


def 					UpperCAmelCase__      (							)						->		int:



						"""simple docstring"""

						snake_case   :					List[Any]							=						hf_hub_download(
						    repo_id="hf-internal-testing/spaghetti-video",			filename="eating_spaghetti_32_frames.npy",			repo_type="dataset")
						snake_case   :					int							=						np.load(_lowerCAmelCase)
						return list(_lowerCAmelCase)


def 					UpperCAmelCase__      (							_lowerCAmelCase						: np.ndarray)						->		str:



						"""simple docstring"""

						if np.ndim(_lowerCAmelCase) == 2:
												return transform_attention_bias(_lowerCAmelCase)

						elif np.ndim(_lowerCAmelCase) == 3:
												return transform_attention_kernel(_lowerCAmelCase)

						else:
												raise Exception(f"""Invalid number of dimesions: {np.ndim(_lowerCAmelCase)}""")


def 					UpperCAmelCase__      (							_lowerCAmelCase						: np.ndarray)						->		Optional[Any]:



						"""simple docstring"""

						return current.flatten()


def 					UpperCAmelCase__      (							_lowerCAmelCase						: np.ndarray)						->		List[Any]:



						"""simple docstring"""

						return np.reshape(_lowerCAmelCase,			(current.shape[0], current.shape[1] * current.shape[2])).T


def 					UpperCAmelCase__      (							_lowerCAmelCase						: np.ndarray)						->		Optional[Any]:



						"""simple docstring"""

						return np.reshape(_lowerCAmelCase,			(current.shape[0] * current.shape[1], current.shape[2])).T


def 					UpperCAmelCase__      (							_lowerCAmelCase						: Union[str, Any],			_lowerCAmelCase						: List[Any])						->		Union[str, Any]:



						"""simple docstring"""

						snake_case   :					Union[str, Any]							=						state_dict["optimizer"]["target"]["Transformer"][f"""encoderblock_{i}"""]

						snake_case   :					Any							=						f"""encoder.layer.{i}."""
						snake_case   :					Optional[Any]							=						{
						    prefix + "intermediate.dense.bias": state["MlpBlock_0"]["Dense_0"]["bias"],
						    prefix + "intermediate.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_0"]["kernel"]),
						    prefix + "output.dense.bias": state["MlpBlock_0"]["Dense_1"]["bias"],
						    prefix + "output.dense.weight": np.transpose(state["MlpBlock_0"]["Dense_1"]["kernel"]),
						    prefix + "layernorm_before.bias": state["LayerNorm_0"]["bias"],
						    prefix + "layernorm_before.weight": state["LayerNorm_0"]["scale"],
						    prefix + "layernorm_after.bias": state["LayerNorm_1"]["bias"],
						    prefix + "layernorm_after.weight": state["LayerNorm_1"]["scale"],
						    prefix
						    + "attention.attention.query.bias": transform_attention(
						        state["MultiHeadDotProductAttention_0"]["query"]["bias"]),
						    prefix
						    + "attention.attention.query.weight": transform_attention(
						        state["MultiHeadDotProductAttention_0"]["query"]["kernel"]),
						    prefix
						    + "attention.attention.key.bias": transform_attention(state["MultiHeadDotProductAttention_0"]["key"]["bias"]),
						    prefix
						    + "attention.attention.key.weight": transform_attention(
						        state["MultiHeadDotProductAttention_0"]["key"]["kernel"]),
						    prefix
						    + "attention.attention.value.bias": transform_attention(
						        state["MultiHeadDotProductAttention_0"]["value"]["bias"]),
						    prefix
						    + "attention.attention.value.weight": transform_attention(
						        state["MultiHeadDotProductAttention_0"]["value"]["kernel"]),
						    prefix + "attention.output.dense.bias": state["MultiHeadDotProductAttention_0"]["out"]["bias"],
						    prefix
						    + "attention.output.dense.weight": transform_attention_output_weight(
						        state["MultiHeadDotProductAttention_0"]["out"]["kernel"]),
						}

						return new_state


def 					UpperCAmelCase__      (							_lowerCAmelCase						: List[Any])						->		List[Any]:



						"""simple docstring"""

						return sum([1 if "encoderblock_" in k else 0 for k in state_dict["optimizer"]["target"]["Transformer"].keys()])


def 					UpperCAmelCase__      (							_lowerCAmelCase						: List[Any],			_lowerCAmelCase						: List[Any]=False)						->		List[Any]:



						"""simple docstring"""

						snake_case   :					Dict							=						get_n_layers(_lowerCAmelCase)

						snake_case   :					Tuple							=						OrderedDict()

						snake_case   :					List[str]							=						state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["bias"]
						snake_case   :					List[str]							=						state_dict["optimizer"]["target"]["Transformer"]["encoder_norm"]["scale"]

						snake_case   :					Tuple							=						np.transpose(
						    state_dict["optimizer"]["target"]["embedding"]["kernel"],			(4, 3, 0, 1, 2))
						snake_case   :					Optional[Any]							=						state_dict["optimizer"]["target"]["embedding"]["bias"]

						snake_case   :					List[Any]							=						state_dict["optimizer"]["target"]["cls"]
						snake_case   :					Tuple							=						state_dict["optimizer"]["target"]["Transformer"]["posembed_input"][
						    "pos_embedding"
						]

						for i in range(_lowerCAmelCase):
												new_state.update(transform_state_encoder_block(_lowerCAmelCase,			_lowerCAmelCase))

						if classification_head:
												snake_case   :					Tuple							=						{"vivit." + k: v for k, v in new_state.items()}
												snake_case   :					Union[str, Any]							=						np.transpose(state_dict["optimizer"]["target"]["output_projection"]["kernel"])
												snake_case   :					Optional[Any]							=						np.transpose(state_dict["optimizer"]["target"]["output_projection"]["bias"])

						return {k: torch.tensor(_lowerCAmelCase) for k, v in new_state.items()}


def 					UpperCAmelCase__      (							)						->		VivitImageProcessor:



						"""simple docstring"""

						snake_case   :					int							=						VivitImageProcessor()

						assert extractor.do_resize is True
						assert extractor.size == {"shortest_edge": 256}
						assert extractor.do_center_crop is True
						assert extractor.crop_size == {"width": 224, "height": 224}
						assert extractor.resample == PILImageResampling.BILINEAR

						# here: https://github.com/deepmind/dmvr/blob/master/dmvr/modalities.py
						# one can seen that add_image has default values for normalization_mean and normalization_std set to 0 and 1
						# which effectively means no normalization (and ViViT does not overwrite those when calling this func)
						assert extractor.do_normalize is False
						assert extractor.do_rescale is True
						assert extractor.rescale_factor == 1 / 255

						# zero-centering = True in original implementation
						assert extractor.do_zero_centering is True

						return extractor


def 					UpperCAmelCase__      (							_lowerCAmelCase						: str)						->		Union[str, Any]:



						"""simple docstring"""

						snake_case   :					str							=						"checkpoint"

						if not os.path.exists(_lowerCAmelCase):
												download_checkpoint(_lowerCAmelCase)

						snake_case   :					Union[str, Any]							=						restore_checkpoint(_lowerCAmelCase,			_lowerCAmelCase)
						snake_case   :					Optional[Any]							=						transform_state(_lowerCAmelCase,			classification_head=_lowerCAmelCase)

						snake_case   :					List[Any]							=						get_vivit_config()

						assert config.image_size == 224
						assert config.num_frames == 32

						snake_case   :					Union[str, Any]							=						VivitForVideoClassification(_lowerCAmelCase)
						model.load_state_dict(_lowerCAmelCase)
						model.eval()

						snake_case   :					Any							=						get_processor()

						snake_case   :					Dict							=						prepare_video()
						snake_case   :					Union[str, Any]							=						extractor(_lowerCAmelCase,			return_tensors="pt")

						snake_case   :					Tuple							=						model(**_lowerCAmelCase)

						snake_case   :					Any							=						torch.Size([1, 400])
						snake_case   :					int							=						torch.tensor([-1.0543, 2.0764, -0.2104, 0.4439, -0.9658])

						assert outputs.logits.shape == expected_shape
						assert torch.allclose(outputs.logits[0, :5],			_lowerCAmelCase,			atol=1E-4), outputs.logits[0, :5]

						model.save_pretrained(_lowerCAmelCase)
						extractor.save_pretrained(_lowerCAmelCase)


if __name__ == "__main__":
						A_        =   argparse.ArgumentParser()

						parser.add_argument("--output_model_name", "-o", type=str, help="Output path for the converted HuggingFace model")

						A_        =   parser.parse_args()
						convert(args.output_model_name)





