code: 
  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   snake_case__					:
  
  
  
  
  
  
  
  		"""simple docstring"""
  
  		def __init__(			self							: Optional[Any]     ,  UpperCamelCase__							: str     ,  UpperCamelCase__							: int=3_6000     ,  UpperCamelCase__							: int=13     ,  UpperCamelCase__							: List[Any]=7     ,  UpperCamelCase__							: List[str]=True     ,  UpperCamelCase__							: str=32     ,  UpperCamelCase__							: Optional[Any]=42     ,  UpperCamelCase__							: str=5     ,  UpperCamelCase__							: List[str]=2     ,  UpperCamelCase__							: str=4     ,  UpperCamelCase__							: Tuple=2     ,  UpperCamelCase__							: Optional[Any]=32     ,  UpperCamelCase__							: List[str]=80     ,  UpperCamelCase__							: Tuple=33     ,  UpperCamelCase__							: Dict=0.0     ,  UpperCamelCase__							: Any=1e-6     ,  UpperCamelCase__							: str=100     ,  UpperCamelCase__							: Tuple=0.0     ,  )	-> Tuple:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					str     =	vocab_size
  				snake_case     :					List[Any]     =	parent
  				snake_case     :					Dict     =	batch_size
  				snake_case     :					Tuple     =	num_contexts
  				# For common tests
  				snake_case     :					List[Any]     =	self.num_contexts
  				snake_case     :					Any     =	is_training
  				snake_case     :					Union[str, Any]     =	hidden_size
  				snake_case     :					Dict     =	num_ext_layers
  				snake_case     :					Union[str, Any]     =	ext_size
  				snake_case     :					str     =	num_hidden_layers
  				snake_case     :					Optional[Any]     =	num_attention_heads
  				snake_case     :					List[Any]     =	num_experts
  				snake_case     :					List[str]     =	d_ff
  				snake_case     :					Any     =	d_ext
  				snake_case     :					Optional[int]     =	d_spout
  				snake_case     :					str     =	dropout_rate
  				snake_case     :					Optional[Any]     =	layer_norm_epsilon
  				snake_case     :					str     =	expert_capacity
  				snake_case     :					str     =	router_jitter_noise
  
  
  
  		def 			lowerCAmelCase							(			self							: str			)	-> List[str]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				return GPTSanJapaneseConfig.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  
  
  
  		def 			lowerCAmelCase							(			self							: Any			)	-> Dict:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Any     =	ids_tensor([self.batch_size, self.seq_length]     ,  self.vocab_size			)
  
  				snake_case     :					List[Any]     =	self.get_config()
  
  				return (config, input_ids)
  
  
  
  		def 			lowerCAmelCase							(			self							: List[str]			)	-> Dict:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Optional[int]     =	ids_tensor([self.batch_size, self.seq_length]     ,  self.vocab_size			)
  
  				snake_case     :					str     =	self.get_config()
  
  				return (config, {"input_ids": input_ids})
  
  
  
  		def 			lowerCAmelCase							(			self							: Any			)	-> str:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				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							: Union[str, Any]     ,  UpperCamelCase__							: Dict     ,  UpperCamelCase__							: Optional[int]     ,  )	-> List[Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					List[str]     =	GPTSanJapaneseForConditionalGeneration(config=UpperCamelCase__			)
  				model.to(UpperCamelCase__			)
  				model.eval()
  				snake_case     :					List[str]     =	model(
  				    input_ids=UpperCamelCase__     ,  )
  				self.parent.assertIsNotNone(UpperCamelCase__			)
  
  
  
  
  
  
  
  @require_torch
  class   snake_case__					(			__SCREAMING_SNAKE_CASE						,				__SCREAMING_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							: Optional[int]     ,  UpperCamelCase__							: Optional[int]     ,  UpperCamelCase__							: Optional[Any]     ,  UpperCamelCase__							: List[str]     ,  UpperCamelCase__							: str     ,  UpperCamelCase__							: Tuple			)	-> str:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				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							: Union[str, Any]			)	-> Tuple:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					int     =	GPTSanJapaneseTester(self			)
  				snake_case     :					Tuple     =	ConfigTester(self     ,  config_class=UpperCamelCase__     ,  d_model=37			)
  
  
  
  		def 			lowerCAmelCase							(			self							: Any			)	-> Optional[Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				GPTSanJapaneseConfig()
  
  
  
  		def 			lowerCAmelCase							(			self							: Optional[Any]			)	-> str:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					str     =	self.model_tester.prepare_config_and_inputs()
  				self.model_tester.create_and_check_model(*UpperCamelCase__			)
  
  
  
  		@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]			)	-> List[str]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				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[Any]			)	-> Dict:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				pass
  
  
  
  
  
  
  
  @require_torch
  class   snake_case__					(			__SCREAMING_SNAKE_CASE						,				__SCREAMING_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							: Any			)	-> str:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Any     =	GPTSanJapaneseTester(self			)
  				snake_case     :					Optional[Any]     =	ConfigTester(self     ,  config_class=UpperCamelCase__     ,  d_model=37			)
  
  
  
  		def 			lowerCAmelCase							(			self							: Dict			)	-> Optional[int]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				GPTSanJapaneseConfig()
  
  
  
  		def 			lowerCAmelCase							(			self							: Tuple			)	-> str:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					str     =	self.model_tester.prepare_config_and_inputs()
  				self.model_tester.create_and_check_model(*UpperCamelCase__			)
  
  
  
  		@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]			)	-> Union[str, Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				super().test_model_parallelism()
  
  
  
  		@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''			)
  		def 			lowerCAmelCase							(			self							: Union[str, Any]			)	-> Any:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				pass
  
  
  
  		@slow
  		def 			lowerCAmelCase							(			self							: Union[str, Any]			)	-> Optional[int]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Union[str, Any]     =	GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					str     =	GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					Union[str, Any]     =	tokenizer.encode('''武田信玄は'''     ,  return_tensors='''pt'''			)
  				snake_case     :					List[Any]     =	model(UpperCamelCase__			)
  				snake_case     :					Optional[int]     =	outputs.logits.detach().cpu().numpy()
  				# Output of original model created with mesh-tensoflow
  				snake_case     :					Union[str, Any]     =	[
  				    # 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     :					Dict     =	np.array(UpperCamelCase__			).flatten()
  				snake_case     :					Union[str, Any]     =	output_logits[0, :, :20].flatten()
  
  				def check(UpperCamelCase__							: Optional[Any]     ,  UpperCamelCase__							: List[Any]     ,  UpperCamelCase__							: List[Any]=5e-4			):
  						return abs(a - b			) < epsilon * max(abs(UpperCamelCase__			)     ,  abs(UpperCamelCase__			)			)
  
  				self.assertTrue(np.all([check(target[i]     ,  predict[i]			) for i in range(len(UpperCamelCase__			)			)]			)			)
  
  
  
  		@slow
  		def 			lowerCAmelCase							(			self							: Any			)	-> Dict:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Tuple     =	GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					Any     =	GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				model.to(UpperCamelCase__			)
  
  				# set deterministically
  				snake_case     :					Any     =	GenerationConfig.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					List[str]     =	1
  
  				# use different length sentences to test batching
  				snake_case     :					Tuple     =	[
  				    '''甲斐なら武田と言うほど''',
  				    '''織田信長は、''',
  				]
  
  				snake_case     :					Union[str, Any]     =	'''left'''
  				snake_case     :					Union[str, Any]     =	tokenizer(UpperCamelCase__     ,  return_tensors='''pt'''     ,  padding=UpperCamelCase__			)
  				snake_case     :					List[Any]     =	inputs['''input_ids'''].to(UpperCamelCase__			)
  
  				self.assertNotEqual(inputs['''attention_mask'''][0].numpy().tolist()     ,  inputs['''attention_mask'''][1].numpy().tolist()			)
  
  				snake_case     :					Optional[Any]     =	model.generate(
  				    input_ids=UpperCamelCase__     ,  attention_mask=inputs['''attention_mask'''].to(UpperCamelCase__			)     ,  max_new_tokens=3     ,  generation_config=UpperCamelCase__     ,  )
  
  				snake_case     :					Tuple     =	tokenizer(sentences[0]     ,  return_tensors='''pt'''			).input_ids.to(UpperCamelCase__			)
  				snake_case     :					Union[str, Any]     =	model.generate(
  				    input_ids=UpperCamelCase__     ,  max_new_tokens=3     ,  generation_config=UpperCamelCase__			)
  
  				snake_case     :					Optional[int]     =	tokenizer(sentences[1]     ,  return_tensors='''pt'''			).input_ids.to(UpperCamelCase__			)
  				snake_case     :					str     =	model.generate(input_ids=UpperCamelCase__     ,  max_new_tokens=3     ,  generation_config=UpperCamelCase__			)
  
  				self.assertNotEqual(inputs_non_padded.shape     ,  inputs_padded.shape			)
  
  				snake_case     :					int     =	tokenizer.batch_decode(UpperCamelCase__     ,  skip_special_tokens=UpperCamelCase__			)
  				snake_case     :					List[Any]     =	tokenizer.decode(output_non_padded[0]     ,  skip_special_tokens=UpperCamelCase__			)
  				snake_case     :					Dict     =	tokenizer.decode(output_padded[0]     ,  skip_special_tokens=UpperCamelCase__			)
  
  				snake_case     :					Optional[int]     =	[
  				    '''甲斐なら武田と言うほど甲斐の武田''',
  				    '''織田信長は、このような''',
  				]
  				self.assertListEqual(UpperCamelCase__     ,  UpperCamelCase__			)
  				self.assertListEqual(UpperCamelCase__     ,  [non_padded_sentence, padded_sentence]			)
  
  
  
  		@tooslow
  		def 			lowerCAmelCase							(			self							: int			)	-> int:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					int     =	GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					Tuple     =	GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				# Output of original model created with mesh-tensoflow
  				snake_case     :					Dict     =	[
  				    ('''武田信玄は''', 3_5675),
  				    ('''武田信玄は、''', 45),
  				    ('''武田信玄は、この''', 29),
  				    ('''武田信玄は、このよう''', 3_0642),
  				    ('''武田信玄は、このような''', 3_5680),
  				    ('''武田信玄は、このような「''', 8640),
  				    ('''武田信玄は、このような「武田''', 3_1617),
  				    ('''武田信玄は、このような「武田家''', 3_0646),
  				    ('''武田信玄は、このような「武田家の''', 3_1617),
  				    ('''武田信玄は、このような「武田家の家''', 3_1381),
  				]
  				for input, output in target:
  						snake_case     :					Union[str, Any]     =	tokenizer.encode(UpperCamelCase__     ,  return_tensors='''pt'''			)
  						snake_case     :					Any     =	model(UpperCamelCase__			)
  						snake_case     :					Any     =	outputs.logits.detach().cpu().numpy()[0]
  						snake_case     :					Any     =	np.argmax(output_logits[-1]			)
  						self.assertEqual(UpperCamelCase__     ,  UpperCamelCase__			)
  
  
  
  		@slow
  		def 			lowerCAmelCase							(			self							: Tuple			)	-> Tuple:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					List[str]     =	GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					List[Any]     =	GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				model.to(UpperCamelCase__			)
  
  				# set deterministically
  				snake_case     :					Optional[Any]     =	GenerationConfig.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					int     =	1
  
  				snake_case     :					Optional[int]     =	'''武田信玄は、'''
  				snake_case     :					Optional[int]     =	tokenizer(UpperCamelCase__     ,  return_tensors='''pt'''			).input_ids.to(UpperCamelCase__			)
  				snake_case     :					Optional[int]     =	tokenizer([input_text, input_text]     ,  return_tensors='''pt'''			).input_ids.to(UpperCamelCase__			)
  
  				# spout from uniform and one-hot
  				snake_case     :					Union[str, Any]     =	[
  				    # 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     :					Optional[int]     =	model.generate(
  				    input_ids=UpperCamelCase__     ,  spout=spouts[0]     ,  max_new_tokens=20     ,  generation_config=UpperCamelCase__     ,  )
  
  				snake_case     :					str     =	model.generate(
  				    input_ids=UpperCamelCase__     ,  spout=spouts[1]     ,  max_new_tokens=20     ,  generation_config=UpperCamelCase__     ,  )
  
  				snake_case     :					str     =	model.generate(
  				    input_ids=UpperCamelCase__     ,  spout=UpperCamelCase__     ,  max_new_tokens=20     ,  generation_config=UpperCamelCase__     ,  )
  
  				snake_case     :					int     =	tokenizer.decode(outputa[0]			)
  				snake_case     :					Optional[int]     =	tokenizer.decode(outputa[0]			)
  				snake_case     :					int     =	tokenizer.batch_decode(UpperCamelCase__			)
  
  				snake_case     :					Tuple     =	[
  				    '''武田信玄は、武田氏の滅亡後、武田氏の居城であった甲斐武田氏の居城である''',
  				    '''武田信玄は、武田家の滅亡を防ぐため、武田家の家臣である武田信虎を討''',
  				]
  				self.assertListEqual(UpperCamelCase__     ,  UpperCamelCase__			)
  				self.assertListEqual(UpperCamelCase__     ,  [outa_sentence, outa_sentence]			)
  
  
  
  		@slow
  		def 			lowerCAmelCase							(			self							: str			)	-> Optional[Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Optional[Any]     =	GPTSanJapaneseForConditionalGeneration.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					Tuple     =	GPTSanJapaneseTokenizer.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				model.to(UpperCamelCase__			)
  
  				# set deterministically
  				snake_case     :					Dict     =	GenerationConfig.from_pretrained('''Tanrei/GPTSAN-japanese'''			)
  				snake_case     :					int     =	1
  
  				snake_case     :					Dict     =	'''武田信玄'''
  				snake_case     :					str     =	'''織田信長'''
  				snake_case     :					Optional[int]     =	'''は、'''
  				snake_case     :					int     =	'''が、'''
  				snake_case     :					str     =	tokenizer(UpperCamelCase__     ,  prefix_text=UpperCamelCase__     ,  return_tensors='''pt'''			)
  				snake_case     :					Dict     =	tokenizer(UpperCamelCase__     ,  prefix_text=UpperCamelCase__     ,  return_tensors='''pt'''			)
  				snake_case     :					List[str]     =	tokenizer([[prefix_text_a, input_text_a], [prefix_text_a, input_text_a]]     ,  return_tensors='''pt'''			)
  
  				snake_case     :					List[Any]     =	model.generate(
  				    input_ids=input_tok_a.input_ids.to(UpperCamelCase__			)     ,  token_type_ids=input_tok_a.token_type_ids.to(UpperCamelCase__			)     ,  max_new_tokens=20     ,  generation_config=UpperCamelCase__     ,  )
  
  				snake_case     :					Dict     =	model.generate(
  				    input_ids=input_tok_a.input_ids.to(UpperCamelCase__			)     ,  token_type_ids=input_tok_a.token_type_ids.to(UpperCamelCase__			)     ,  max_new_tokens=20     ,  generation_config=UpperCamelCase__     ,  )
  
  				snake_case     :					Union[str, Any]     =	model.generate(
  				    input_ids=input_tok_a.input_ids.to(UpperCamelCase__			)     ,  token_type_ids=input_tok_a.token_type_ids.to(UpperCamelCase__			)     ,  attention_mask=input_tok_a.attention_mask.to(UpperCamelCase__			)     ,  max_new_tokens=20     ,  generation_config=UpperCamelCase__     ,  )
  
  				snake_case     :					List[Any]     =	tokenizer.decode(outputa[0]			)
  				snake_case     :					List[str]     =	tokenizer.decode(outputa[0]			)
  				snake_case     :					Union[str, Any]     =	tokenizer.batch_decode(UpperCamelCase__			)
  
  				snake_case     :					Any     =	[
  				    '''武田信玄は、武田氏の祖である武田信虎を、その子・武田信友を擁して''',
  				    '''織田信長が、織田信長の妻・お市の方を妻として迎えたという逸話が残''',
  				]
  				self.assertListEqual(UpperCamelCase__     ,  UpperCamelCase__			)
  				self.assertListEqual(UpperCamelCase__     ,  [outa_sentence, outa_sentence]			)
  
  
  
style_context: 
  '''simple docstring'''
  from typing import Callable, List, Optional, Union
  
  import torch
  
  from ...models import UNetaDModel
  from ...schedulers import CMStochasticIterativeScheduler
  from ...utils import (
      is_accelerate_available,
      is_accelerate_version,
      logging,
      randn_tensor,
      replace_example_docstring,
  )
  from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
  
  
  lowercase__						=			logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  lowercase__						=			"\n    Examples:\n        ```py\n        >>> import torch\n\n        >>> from diffusers import ConsistencyModelPipeline\n\n        >>> device = \"cuda\"\n        >>> # Load the cd_imagenet64_l2 checkpoint.\n        >>> model_id_or_path = \"openai/diffusers-cd_imagenet64_l2\"\n        >>> pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)\n        >>> pipe.to(device)\n\n        >>> # Onestep Sampling\n        >>> image = pipe(num_inference_steps=1).images[0]\n        >>> image.save(\"cd_imagenet64_l2_onestep_sample.png\")\n\n        >>> # Onestep sampling, class-conditional image generation\n        >>> # ImageNet-64 class label 145 corresponds to king penguins\n        >>> image = pipe(num_inference_steps=1, class_labels=145).images[0]\n        >>> image.save(\"cd_imagenet64_l2_onestep_sample_penguin.png\")\n\n        >>> # Multistep sampling, class-conditional image generation\n        >>> # Timesteps can be explicitly specified; the particular timesteps below are from the original Github repo:\n        >>> # https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77\n        >>> image = pipe(num_inference_steps=None, timesteps=[22, 0], class_labels=145).images[0]\n        >>> image.save(\"cd_imagenet64_l2_multistep_sample_penguin.png\")\n        ```\n"
  
  
  
  
  
  
  class   snake_case__					(			__SCREAMING_SNAKE_CASE				):
  
  
  
  
  
  
  
  		"""simple docstring"""
  
  		def __init__(			self							: List[Any]     ,  UpperCamelCase__							: UNetaDModel     ,  UpperCamelCase__							: CMStochasticIterativeScheduler			)	-> None:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				super().__init__()
  
  				self.register_modules(
  				    unet=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  )
  
  				snake_case     :					str     =	None
  
  
  
  		def 			lowerCAmelCase							(			self							: Tuple     ,  UpperCamelCase__							: str=0			)	-> int:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				if is_accelerate_available() and is_accelerate_version('''>='''     ,  '''0.14.0'''			):
  						from accelerate import cpu_offload
  				else:
  						raise ImportError('''`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher'''			)
  
  				snake_case     :					Optional[Any]     =	torch.device(f'cuda:{gpu_id}'			)
  
  				if self.device.type != "cpu":
  						self.to('''cpu'''     ,  silence_dtype_warnings=UpperCamelCase__			)
  						torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
  
  				for cpu_offloaded_model in [self.unet]:
  						cpu_offload(UpperCamelCase__     ,  UpperCamelCase__			)
  
  				if self.safety_checker is not None:
  						cpu_offload(self.safety_checker     ,  execution_device=UpperCamelCase__     ,  offload_buffers=UpperCamelCase__			)
  
  
  
  		def 			lowerCAmelCase							(			self							: Optional[int]     ,  UpperCamelCase__							: int=0			)	-> Dict:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				if is_accelerate_available() and is_accelerate_version('''>='''     ,  '''0.17.0.dev0'''			):
  						from accelerate import cpu_offload_with_hook
  				else:
  						raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''			)
  
  				snake_case     :					Union[str, Any]     =	torch.device(f'cuda:{gpu_id}'			)
  
  				if self.device.type != "cpu":
  						self.to('''cpu'''     ,  silence_dtype_warnings=UpperCamelCase__			)
  						torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
  
  				snake_case     :					List[Any]     =	None
  				for cpu_offloaded_model in [self.unet]:
  						snake_case ,snake_case     :					Union[str, Any]     =	cpu_offload_with_hook(UpperCamelCase__     ,  UpperCamelCase__     ,  prev_module_hook=UpperCamelCase__			)
  
  				if self.safety_checker is not None:
  						snake_case ,snake_case     :					Optional[Any]     =	cpu_offload_with_hook(self.safety_checker     ,  UpperCamelCase__     ,  prev_module_hook=UpperCamelCase__			)
  
  				# We'll offload the last model manually.
  				snake_case     :					Optional[Any]     =	hook
  
  
  
  		@property
  		# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  		def 			lowerCAmelCase							(			self							: int			)	-> List[str]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				if not hasattr(self.unet     ,  '''_hf_hook'''			):
  						return self.device
  				for module in self.unet.modules():
  						if (
  						    hasattr(UpperCamelCase__     ,  '''_hf_hook'''			)
  						    and hasattr(module._hf_hook     ,  '''execution_device'''			)
  						    and module._hf_hook.execution_device is not None
  						):
  								return torch.device(module._hf_hook.execution_device			)
  				return self.device
  
  
  
  		def 			lowerCAmelCase							(			self							: int     ,  UpperCamelCase__							: Optional[Any]     ,  UpperCamelCase__							: Dict     ,  UpperCamelCase__							: int     ,  UpperCamelCase__							: List[str]     ,  UpperCamelCase__							: int     ,  UpperCamelCase__							: Optional[int]     ,  UpperCamelCase__							: Any     ,  UpperCamelCase__							: Optional[int]=None			)	-> Tuple:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Union[str, Any]     =	(batch_size, num_channels, height, width)
  				if isinstance(UpperCamelCase__     ,  UpperCamelCase__			) and len(UpperCamelCase__			) != batch_size:
  						raise ValueError(
  						    f'You have passed a list of generators of length {len(UpperCamelCase__			)}, but requested an effective batch'
  						    f' size of {batch_size}. Make sure the batch size matches the length of the generators.'			)
  
  				if latents is None:
  						snake_case     :					List[Any]     =	randn_tensor(UpperCamelCase__     ,  generator=UpperCamelCase__     ,  device=UpperCamelCase__     ,  dtype=UpperCamelCase__			)
  				else:
  						snake_case     :					Tuple     =	latents.to(device=UpperCamelCase__     ,  dtype=UpperCamelCase__			)
  
  				# scale the initial noise by the standard deviation required by the scheduler
  				snake_case     :					str     =	latents * self.scheduler.init_noise_sigma
  				return latents
  
  
  
  		def 			lowerCAmelCase							(			self							: Any     ,  UpperCamelCase__							: torch.FloatTensor     ,  UpperCamelCase__							: str = "pil"			)	-> Union[str, Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				if output_type not in ["pt", "np", "pil"]:
  						raise ValueError(
  						    f'output_type={output_type} is not supported. Make sure to choose one of [\'pt\', \'np\', or \'pil\']'			)
  
  				# Equivalent to diffusers.VaeImageProcessor.denormalize
  				snake_case     :					Any     =	(sample / 2 + 0.5).clamp(0     ,  1			)
  				if output_type == "pt":
  						return sample
  
  				# Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
  				snake_case     :					Optional[int]     =	sample.cpu().permute(0     ,  2     ,  3     ,  1			).numpy()
  				if output_type == "np":
  						return sample
  
  				# Output_type must be 'pil'
  				snake_case     :					Tuple     =	self.numpy_to_pil(UpperCamelCase__			)
  				return sample
  
  
  
  		def 			lowerCAmelCase							(			self							: Tuple     ,  UpperCamelCase__							: List[str]     ,  UpperCamelCase__							: Dict     ,  UpperCamelCase__							: List[str]=None			)	-> List[str]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				if self.unet.config.num_class_embeds is not None:
  						if isinstance(UpperCamelCase__     ,  UpperCamelCase__			):
  								snake_case     :					str     =	torch.tensor(UpperCamelCase__     ,  dtype=torch.int			)
  						elif isinstance(UpperCamelCase__     ,  UpperCamelCase__			):
  								assert batch_size == 1, "Batch size must be 1 if classes is an int"
  								snake_case     :					Dict     =	torch.tensor([class_labels]     ,  dtype=torch.int			)
  						elif class_labels is None:
  								# Randomly generate batch_size class labels
  								# TODO: should use generator here? int analogue of randn_tensor is not exposed in ...utils
  								snake_case     :					List[Any]     =	torch.randint(0     ,  self.unet.config.num_class_embeds     ,  size=(batch_size,)			)
  						snake_case     :					Dict     =	class_labels.to(UpperCamelCase__			)
  				else:
  						snake_case     :					int     =	None
  				return class_labels
  
  
  
  		def 			lowerCAmelCase							(			self							: List[Any]     ,  UpperCamelCase__							: int     ,  UpperCamelCase__							: Any     ,  UpperCamelCase__							: List[Any]     ,  UpperCamelCase__							: List[Any]     ,  UpperCamelCase__							: List[str]     ,  UpperCamelCase__							: Union[str, Any]			)	-> Union[str, Any]:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				if num_inference_steps is None and timesteps is None:
  						raise ValueError('''Exactly one of `num_inference_steps` or `timesteps` must be supplied.'''			)
  
  				if num_inference_steps is not None and timesteps is not None:
  						logger.warning(
  						    f'Both `num_inference_steps`: {num_inference_steps} and `timesteps`: {timesteps} are supplied;'
  						    ''' `timesteps` will be used over `num_inference_steps`.'''			)
  
  				if latents is not None:
  						snake_case     :					List[Any]     =	(batch_size, 3, img_size, img_size)
  						if latents.shape != expected_shape:
  								raise ValueError(f'The shape of latents is {latents.shape} but is expected to be {expected_shape}.'			)
  
  				if (callback_steps is None) or (
  				    callback_steps is not None and (not isinstance(UpperCamelCase__     ,  UpperCamelCase__			) or callback_steps <= 0)
  				):
  						raise ValueError(
  						    f'`callback_steps` has to be a positive integer but is {callback_steps} of type'
  						    f' {type(UpperCamelCase__			)}.'			)
  
  
  
  		@torch.no_grad()
  		@replace_example_docstring(UpperCamelCase__			)
  		def __call__(			self							: Optional[int]     ,  UpperCamelCase__							: int = 1     ,  UpperCamelCase__							: Optional[Union[torch.Tensor, List[int], int]] = None     ,  UpperCamelCase__							: int = 1     ,  UpperCamelCase__							: List[int] = None     ,  UpperCamelCase__							: Optional[Union[torch.Generator, List[torch.Generator]]] = None     ,  UpperCamelCase__							: Optional[torch.FloatTensor] = None     ,  UpperCamelCase__							: Optional[str] = "pil"     ,  UpperCamelCase__							: bool = True     ,  UpperCamelCase__							: Optional[Callable[[int, int, torch.FloatTensor], None]] = None     ,  UpperCamelCase__							: int = 1     ,  )	-> str:
  
  
  
  
  
  
  				"""simple docstring"""
  
  
  
  
  
  
  
  				snake_case     :					Optional[int]     =	self.unet.config.sample_size
  				snake_case     :					List[Any]     =	self._execution_device
  
  				# 1. Check inputs
  				self.check_inputs(UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__			)
  
  				# 2. Prepare image latents
  				# Sample image latents x_0 ~ N(0, sigma_0^2 * I)
  				snake_case     :					List[str]     =	self.prepare_latents(
  				    batch_size=UpperCamelCase__     ,  num_channels=self.unet.config.in_channels     ,  height=UpperCamelCase__     ,  width=UpperCamelCase__     ,  dtype=self.unet.dtype     ,  device=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  latents=UpperCamelCase__     ,  )
  
  				# 3. Handle class_labels for class-conditional models
  				snake_case     :					List[str]     =	self.prepare_class_labels(UpperCamelCase__     ,  UpperCamelCase__     ,  class_labels=UpperCamelCase__			)
  
  				# 4. Prepare timesteps
  				if timesteps is not None:
  						self.scheduler.set_timesteps(timesteps=UpperCamelCase__     ,  device=UpperCamelCase__			)
  						snake_case     :					Union[str, Any]     =	self.scheduler.timesteps
  						snake_case     :					Optional[int]     =	len(UpperCamelCase__			)
  				else:
  						self.scheduler.set_timesteps(UpperCamelCase__			)
  						snake_case     :					List[str]     =	self.scheduler.timesteps
  
  				# 5. Denoising loop
  				# Multistep sampling: implements Algorithm 1 in the paper
  				with self.progress_bar(total=UpperCamelCase__			) as progress_bar:
  						for i, t in enumerate(UpperCamelCase__			):
  								snake_case     :					Any     =	self.scheduler.scale_model_input(UpperCamelCase__     ,  UpperCamelCase__			)
  								snake_case     :					int     =	self.unet(UpperCamelCase__     ,  UpperCamelCase__     ,  class_labels=UpperCamelCase__     ,  return_dict=UpperCamelCase__			)[0]
  
  								snake_case     :					Dict     =	self.scheduler.step(UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  generator=UpperCamelCase__			)[0]
  
  								# call the callback, if provided
  								progress_bar.update()
  								if callback is not None and i % callback_steps == 0:
  										callback(UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__			)
  
          # 6. Post-process image sample
  				snake_case     :					Any     =	self.postprocess_image(UpperCamelCase__     ,  output_type=UpperCamelCase__			)
  
  				# Offload last model to CPU
  				if hasattr(self     ,  '''final_offload_hook'''			) and self.final_offload_hook is not None:
  						self.final_offload_hook.offload()
  
  				if not return_dict:
  						return (image,)
  
  				return ImagePipelineOutput(images=UpperCamelCase__			)
  
  
  