






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


							def __init__(      self   :    Optional[int]      ,     snake_case   :    Optional[int]      ,     snake_case   :    List[str]=3_6000      ,     snake_case   :    Tuple=13      ,     snake_case   :    int=7      ,     snake_case   :    Dict=True      ,     snake_case   :    int=32      ,     snake_case   :    List[str]=42      ,     snake_case   :    List[Any]=5      ,     snake_case   :    Any=2      ,     snake_case   :    List[Any]=4      ,     snake_case   :    List[str]=2      ,     snake_case   :    Union[str, Any]=32      ,     snake_case   :    str=80      ,     snake_case   :    List[str]=33      ,     snake_case   :    Optional[Any]=0.0      ,     snake_case   :    str=1e-6      ,     snake_case   :    Dict=100      ,     snake_case   :    Tuple=0.0      ,     ):
													'''simple docstring'''





													A__    :	str			=	vocab_size
													A__    :	Optional[Any]			=	parent
													A__    :	Optional[int]			=	batch_size
													A__    :	Any			=	num_contexts
													# For common tests
													A__    :	str			=	self.num_contexts
													A__    :	int			=	is_training
													A__    :	Dict			=	hidden_size
													A__    :	Dict			=	num_ext_layers
													A__    :	str			=	ext_size
													A__    :	int			=	num_hidden_layers
													A__    :	List[Any]			=	num_attention_heads
													A__    :	Union[str, Any]			=	num_experts
													A__    :	Dict			=	d_ff
													A__    :	Union[str, Any]			=	d_ext
													A__    :	List[str]			=	d_spout
													A__    :	str			=	dropout_rate
													A__    :	str			=	layer_norm_epsilon
													A__    :	Any			=	expert_capacity
													A__    :	Optional[Any]			=	router_jitter_noise


							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													return GPTSanJapaneseConfig.from_pretrained("""Tanrei/GPTSAN-japanese"""				)


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





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

													A__    :	List[Any]			=	self.get_config()

													return (config, input_ids)


							def     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													A__    :	Any			=	ids_tensor([self.batch_size, self.seq_length]      ,     self.vocab_size				)

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

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


							def     _UpperCamelCase							(      self   :    Optional[int]				):
													'''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     _UpperCamelCase							(      self   :    Optional[Any]      ,     snake_case   :    int      ,     snake_case   :    List[str]      ,     ):
													'''simple docstring'''





													A__    :	Tuple			=	GPTSanJapaneseForConditionalGeneration(config=snake_case				)
													model.to(snake_case				)
													model.eval()
													A__    :	int			=	model(
													    input_ids=snake_case      ,     )
													self.parent.assertIsNotNone(snake_case				)






@require_torch
class 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase							,		UpperCamelCase							,		unittest.TestCase		):
							snake_case_  =				(GPTSanJapaneseModel,) if is_torch_available() else ()
							snake_case_  =				(
							    {
							        'conversational': GPTSanJapaneseForConditionalGeneration,
							        'feature-extraction': GPTSanJapaneseForConditionalGeneration,
							        'summarization': GPTSanJapaneseForConditionalGeneration,
							        'text2text-generation': GPTSanJapaneseForConditionalGeneration,
							        'translation': GPTSanJapaneseForConditionalGeneration,
							    }
							    if is_torch_available()
							    else {}
							)
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							# The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
							snake_case_  =				[0.8, 0.9]


							def     _UpperCamelCase							(      self   :    Dict      ,     snake_case   :    Optional[Any]      ,     snake_case   :    str      ,     snake_case   :    List[Any]      ,     snake_case   :    str      ,     snake_case   :    Optional[Any]				):
													'''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     _UpperCamelCase							(      self   :    Optional[int]				):
													'''simple docstring'''





													A__    :	List[str]			=	GPTSanJapaneseTester(self				)
													A__    :	List[Any]			=	ConfigTester(self      ,     config_class=snake_case      ,     d_model=37				)


							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													GPTSanJapaneseConfig()


							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__    :	str			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*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     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													super().test_model_parallelism()


							@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""				)
							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													pass






@require_torch
class 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase							,		UpperCamelCase							,		unittest.TestCase		):
							snake_case_  =				(GPTSanJapaneseForConditionalGeneration,) if is_torch_available() else ()
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							snake_case_  =				False
							# The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
							snake_case_  =				[0.8, 0.9]


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__    :	Tuple			=	GPTSanJapaneseTester(self				)
													A__    :	Optional[Any]			=	ConfigTester(self      ,     config_class=snake_case      ,     d_model=37				)


							def     _UpperCamelCase							(      self   :    Optional[Any]				):
													'''simple docstring'''





													GPTSanJapaneseConfig()


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__    :	Tuple			=	self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*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     _UpperCamelCase							(      self   :    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     _UpperCamelCase							(      self   :    Dict				):
													'''simple docstring'''





													pass


							@slow
							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__    :	List[Any]			=	GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													A__    :	List[str]			=	GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													A__    :	List[Any]			=	tokenizer.encode("""武田信玄は"""      ,     return_tensors="""pt"""				)
													A__    :	str			=	model(snake_case				)
													A__    :	Optional[Any]			=	outputs.logits.detach().cpu().numpy()
													# Output of original model created with mesh-tensoflow
													A__    :	Any			=	[
													    # 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
													]
													A__    :	int			=	np.array(snake_case				).flatten()
													A__    :	Optional[Any]			=	output_logits[0, :, :20].flatten()

													def check(snake_case   :    int      ,     snake_case   :    Tuple      ,     snake_case   :    Union[str, Any]=5e-4				):
																			return abs(a - b				) < epsilon * max(abs(snake_case				)      ,     abs(snake_case				)				)

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


							@slow
							def     _UpperCamelCase							(      self   :    Dict				):
													'''simple docstring'''





													A__    :	Any			=	GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													A__    :	int			=	GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													model.to(snake_case				)

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

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

													A__    :	int			=	"""left"""
													A__    :	Union[str, Any]			=	tokenizer(snake_case      ,     return_tensors="""pt"""      ,     padding=snake_case				)
													A__    :	List[Any]			=	inputs["""input_ids"""].to(snake_case				)

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

													A__    :	List[str]			=	model.generate(
													    input_ids=snake_case      ,     attention_mask=inputs["""attention_mask"""].to(snake_case				)      ,     max_new_tokens=3      ,     generation_config=snake_case      ,     )

													A__    :	List[Any]			=	tokenizer(sentences[0]      ,     return_tensors="""pt"""				).input_ids.to(snake_case				)
													A__    :	str			=	model.generate(
													    input_ids=snake_case      ,     max_new_tokens=3      ,     generation_config=snake_case				)

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

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

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

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


							@tooslow
							def     _UpperCamelCase							(      self   :    Tuple				):
													'''simple docstring'''





													A__    :	Optional[int]			=	GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													A__    :	str			=	GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													# Output of original model created with mesh-tensoflow
													A__    :	List[str]			=	[
													    ("""武田信玄は""", 3_5675),
													    ("""武田信玄は、""", 45),
													    ("""武田信玄は、この""", 29),
													    ("""武田信玄は、このよう""", 3_0642),
													    ("""武田信玄は、このような""", 3_5680),
													    ("""武田信玄は、このような「""", 8640),
													    ("""武田信玄は、このような「武田""", 3_1617),
													    ("""武田信玄は、このような「武田家""", 3_0646),
													    ("""武田信玄は、このような「武田家の""", 3_1617),
													    ("""武田信玄は、このような「武田家の家""", 3_1381),
													]
													for input, output in target:
																			A__    :	Any			=	tokenizer.encode(snake_case      ,     return_tensors="""pt"""				)
																			A__    :	Any			=	model(snake_case				)
																			A__    :	List[Any]			=	outputs.logits.detach().cpu().numpy()[0]
																			A__    :	str			=	np.argmax(output_logits[-1]				)
																			self.assertEqual(snake_case      ,     snake_case				)


							@slow
							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													A__    :	Optional[int]			=	GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													A__    :	Tuple			=	GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													model.to(snake_case				)

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

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

													# spout from uniform and one-hot
													A__    :	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
													]

													A__    :	List[Any]			=	model.generate(
													    input_ids=snake_case      ,     spout=spouts[0]      ,     max_new_tokens=20      ,     generation_config=snake_case      ,     )

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

													A__    :	str			=	model.generate(
													    input_ids=snake_case      ,     spout=snake_case      ,     max_new_tokens=20      ,     generation_config=snake_case      ,     )

													A__    :	int			=	tokenizer.decode(outputa[0]				)
													A__    :	int			=	tokenizer.decode(outputa[0]				)
													A__    :	Optional[int]			=	tokenizer.batch_decode(snake_case				)

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


							@slow
							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													A__    :	Union[str, Any]			=	GPTSanJapaneseForConditionalGeneration.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													A__    :	Optional[Any]			=	GPTSanJapaneseTokenizer.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													model.to(snake_case				)

													# set deterministically
													A__    :	Any			=	GenerationConfig.from_pretrained("""Tanrei/GPTSAN-japanese"""				)
													A__    :	List[str]			=	1

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

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

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

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

													A__    :	int			=	tokenizer.decode(outputa[0]				)
													A__    :	Any			=	tokenizer.decode(outputa[0]				)
													A__    :	Tuple			=	tokenizer.batch_decode(snake_case				)

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














"""simple docstring"""



import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

A_     =				logging.get_logger(__name__)




def   _lowerCAmelCase		(				UpperCAmelCase__							:       int,    UpperCAmelCase__							:       bool = True   )		->Tuple:
						A__    :	List[str]			=	{}
						for name in model.attn_processors.keys():
												A__    :	Optional[Any]			=	name.endswith("""attn2.processor"""   ) and not (
												    name.startswith("""transformer_in"""   ) or """temp_attentions""" in name.split("""."""   )
												)
												A__    :	Tuple			=	model.config.cross_attention_dim if has_cross_attention else None
												if name.startswith("""mid_block"""   ):
																		A__    :	Tuple			=	model.config.block_out_channels[-1]
												elif name.startswith("""up_blocks"""   ):
																		A__    :	str			=	int(name[len("""up_blocks."""   )]   )
																		A__    :	Tuple			=	list(reversed(model.config.block_out_channels   )   )[block_id]
												elif name.startswith("""down_blocks"""   ):
																		A__    :	List[str]			=	int(name[len("""down_blocks."""   )]   )
																		A__    :	Tuple			=	model.config.block_out_channels[block_id]
												elif name.startswith("""transformer_in"""   ):
																		# Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
																		A__    :	List[str]			=	8 * model.config.attention_head_dim

												A__    :	Optional[Any]			=	LoRAAttnProcessor(hidden_size=UpperCAmelCase__,    cross_attention_dim=UpperCAmelCase__   )
												A__    :	Optional[int]			=	lora_attn_procs[name].to(model.device   )

												if mock_weights:
																		# add 1 to weights to mock trained weights
																		with torch.no_grad():
																								lora_attn_procs[name].to_q_lora.up.weight += 1
																								lora_attn_procs[name].to_k_lora.up.weight += 1
																								lora_attn_procs[name].to_v_lora.up.weight += 1
																								lora_attn_procs[name].to_out_lora.up.weight += 1

						return lora_attn_procs






@skip_mps
class 				__SCREAMING_SNAKE_CASE    (					UpperCamelCase							,		UpperCamelCase							,		unittest.TestCase		):
							snake_case_  =				UNetaDConditionModel
							snake_case_  =				'sample'


							@property
							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__    :	Dict			=	4
													A__    :	str			=	4
													A__    :	str			=	4
													A__    :	Union[str, Any]			=	(32, 32)

													A__    :	List[Any]			=	floats_tensor((batch_size, num_channels, num_frames) + sizes				).to(snake_case				)
													A__    :	Any			=	torch.tensor([10]				).to(snake_case				)
													A__    :	Optional[Any]			=	floats_tensor((batch_size, 4, 32)				).to(snake_case				)

													return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}


							@property
							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													return (4, 4, 32, 32)


							@property
							def     _UpperCamelCase							(      self   :    Optional[int]				):
													'''simple docstring'''





													return (4, 4, 32, 32)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__    :	Tuple			=	{
													    """block_out_channels""": (32, 64),
													    """down_block_types""": (
													        """CrossAttnDownBlock3D""",
													        """DownBlock3D""",
													    ),
													    """up_block_types""": ("""UpBlock3D""", """CrossAttnUpBlock3D"""),
													    """cross_attention_dim""": 32,
													    """attention_head_dim""": 8,
													    """out_channels""": 4,
													    """in_channels""": 4,
													    """layers_per_block""": 1,
													    """sample_size""": 32,
													}
													A__    :	Optional[Any]			=	self.dummy_input
													return init_dict, inputs_dict


							@unittest.skipIf(
							    torch_device != """cuda""" or not is_xformers_available()      ,     reason="""XFormers attention is only available with CUDA and `xformers` installed"""      ,     )
							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()
													A__    :	Dict			=	self.model_class(**snake_case				)

													model.enable_xformers_memory_efficient_attention()

													assert (
													    model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
													    == "XFormersAttnProcessor"
													), "xformers is not enabled"


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	Optional[int]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Tuple			=	32

													A__    :	List[Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				)

																			if isinstance(snake_case      ,     snake_case				):
																									A__    :	int			=	output.sample

													self.assertIsNotNone(snake_case				)
													A__    :	Any			=	inputs_dict["""sample"""].shape
													self.assertEqual(output.shape      ,     snake_case      ,     """Input and output shapes do not match"""				)


							def     _UpperCamelCase							(      self   :    int				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()
													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													with torch.no_grad():
																			# Warmup pass when using mps (see #372)
																			if torch_device == "mps" and isinstance(snake_case      ,     snake_case				):
																									model(**self.dummy_input				)

																			A__    :	List[str]			=	model(**snake_case				)
																			if isinstance(snake_case      ,     snake_case				):
																									A__    :	List[Any]			=	first.sample

																			A__    :	List[str]			=	model(**snake_case				)
																			if isinstance(snake_case      ,     snake_case				):
																									A__    :	str			=	second.sample

													A__    :	Optional[Any]			=	first.cpu().numpy()
													A__    :	Tuple			=	second.cpu().numpy()
													A__    :	List[Any]			=	out_a[~np.isnan(snake_case				)]
													A__    :	Dict			=	out_a[~np.isnan(snake_case				)]
													A__    :	str			=	np.amax(np.abs(out_a - out_a				)				)
													self.assertLessEqual(snake_case      ,     1e-5				)


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													A__	,	A__    :	List[Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[str]			=	8

													A__    :	Optional[Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													model.set_attention_slice("""auto"""				)
													with torch.no_grad():
																			A__    :	List[str]			=	model(**snake_case				)
													assert output is not None

													model.set_attention_slice("""max"""				)
													with torch.no_grad():
																			A__    :	Union[str, Any]			=	model(**snake_case				)
													assert output is not None

													model.set_attention_slice(2				)
													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				)
													assert output is not None


							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[Any]			=	8

													A__    :	Union[str, Any]			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Optional[int]			=	model(**snake_case				).sample

													A__    :	Dict			=	create_lora_layers(snake_case				)

													# make sure we can set a list of attention processors
													model.set_attn_processor(snake_case				)
													model.to(snake_case				)

													# test that attn processors can be set to itself
													model.set_attn_processor(model.attn_processors				)

													with torch.no_grad():
																			A__    :	int			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.0}				).sample
																			A__    :	Optional[Any]			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample
																			A__    :	List[Any]			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample

													assert (samplea - samplea).abs().max() < 3e-3
													assert (samplea - samplea).abs().max() < 3e-3

													# sample 2 and sample 3 should be different
													assert (samplea - samplea).abs().max() > 3e-3


							def     _UpperCamelCase							(      self   :    str				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	str			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	model(**snake_case				).sample

													A__    :	str			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Union[str, Any]			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample

													with tempfile.TemporaryDirectory() as tmpdirname:
																			model.save_attn_procs(snake_case				)
																			self.assertTrue(os.path.isfile(os.path.join(snake_case      ,     """pytorch_lora_weights.bin"""				)				)				)
																			torch.manual_seed(0				)
																			A__    :	Union[str, Any]			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case				)

													with torch.no_grad():
																			A__    :	List[Any]			=	new_model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample

													assert (sample - new_sample).abs().max() < 1e-3

													# LoRA and no LoRA should NOT be the same
													assert (sample - old_sample).abs().max() > 1e-4


							def     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Tuple			=	model(**snake_case				).sample

													A__    :	Union[str, Any]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Tuple			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample

													with tempfile.TemporaryDirectory() as tmpdirname:
																			model.save_attn_procs(snake_case      ,     safe_serialization=snake_case				)
																			self.assertTrue(os.path.isfile(os.path.join(snake_case      ,     """pytorch_lora_weights.safetensors"""				)				)				)
																			torch.manual_seed(0				)
																			A__    :	Optional[Any]			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	new_model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.5}				).sample

													assert (sample - new_sample).abs().max() < 3e-3

													# LoRA and no LoRA should NOT be the same
													assert (sample - old_sample).abs().max() > 1e-4


							def     _UpperCamelCase							(      self   :    List[Any]				):
													'''simple docstring'''





													A__	,	A__    :	List[Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Any			=	8

													torch.manual_seed(0				)
													A__    :	str			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													A__    :	int			=	create_lora_layers(snake_case      ,     mock_weights=snake_case				)
													model.set_attn_processor(snake_case				)
													# Saving as torch, properly reloads with directly filename
													with tempfile.TemporaryDirectory() as tmpdirname:
																			model.save_attn_procs(snake_case				)
																			self.assertTrue(os.path.isfile(os.path.join(snake_case      ,     """pytorch_lora_weights.bin"""				)				)				)
																			torch.manual_seed(0				)
																			A__    :	str			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			new_model.load_attn_procs(snake_case      ,     weight_name="""pytorch_lora_weights.bin"""				)


							def     _UpperCamelCase							(      self   :    Union[str, Any]				):
													'''simple docstring'''





													A__	,	A__    :	Union[str, Any]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Optional[Any]			=	8

													torch.manual_seed(0				)
													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													A__    :	Dict			=	create_lora_layers(snake_case      ,     mock_weights=snake_case				)
													model.set_attn_processor(snake_case				)
													# Saving as torch, properly reloads with directly filename
													with tempfile.TemporaryDirectory() as tmpdirname:
																			model.save_attn_procs(snake_case				)
																			self.assertTrue(os.path.isfile(os.path.join(snake_case      ,     """pytorch_lora_weights.bin"""				)				)				)
																			torch.manual_seed(0				)
																			A__    :	Optional[Any]			=	self.model_class(**snake_case				)
																			new_model.to(snake_case				)
																			with self.assertRaises(snake_case				) as e:
																									new_model.load_attn_procs(snake_case      ,     use_safetensors=snake_case				)
																			self.assertIn("""Error no file named pytorch_lora_weights.safetensors"""      ,     str(e.exception				)				)


							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	Dict			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	Tuple			=	8

													torch.manual_seed(0				)
													A__    :	Optional[int]			=	self.model_class(**snake_case				)
													model.to(snake_case				)

													with torch.no_grad():
																			A__    :	Dict			=	model(**snake_case				).sample

													A__    :	Union[str, Any]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													with torch.no_grad():
																			A__    :	Any			=	model(**snake_case      ,     cross_attention_kwargs={"""scale""": 0.0}				).sample

													model.set_attn_processor(AttnProcessor()				)

													with torch.no_grad():
																			A__    :	str			=	model(**snake_case				).sample

													assert (sample - new_sample).abs().max() < 1e-4
													assert (sample - old_sample).abs().max() < 3e-3


							@unittest.skipIf(
							    torch_device != """cuda""" or not is_xformers_available()      ,     reason="""XFormers attention is only available with CUDA and `xformers` installed"""      ,     )
							def     _UpperCamelCase							(      self   :    List[str]				):
													'''simple docstring'''





													A__	,	A__    :	List[str]			=	self.prepare_init_args_and_inputs_for_common()

													A__    :	List[Any]			=	4

													torch.manual_seed(0				)
													A__    :	Tuple			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													A__    :	List[str]			=	create_lora_layers(snake_case				)
													model.set_attn_processor(snake_case				)

													# default
													with torch.no_grad():
																			A__    :	Optional[Any]			=	model(**snake_case				).sample

																			model.enable_xformers_memory_efficient_attention()
																			A__    :	List[str]			=	model(**snake_case				).sample

																			model.disable_xformers_memory_efficient_attention()
																			A__    :	Union[str, Any]			=	model(**snake_case				).sample

													assert (sample - on_sample).abs().max() < 1e-4
													assert (sample - off_sample).abs().max() < 1e-4


							def     _UpperCamelCase							(      self   :    Any				):
													'''simple docstring'''





													A__	,	A__    :	Optional[Any]			=	self.prepare_init_args_and_inputs_for_common()
													A__    :	List[Any]			=	32

													A__    :	Any			=	self.model_class(**snake_case				)
													model.to(snake_case				)
													model.eval()

													with torch.no_grad():
																			A__    :	List[Any]			=	model(**snake_case				)[0]

													model.enable_forward_chunking()
													with torch.no_grad():
																			A__    :	Union[str, Any]			=	model(**snake_case				)[0]

													self.assertEqual(output.shape      ,     output_a.shape      ,     """Shape doesn't match"""				)
													assert np.abs(output.cpu() - output_a.cpu()				).max() < 1e-2


# (todo: sayakpaul) implement SLOW tests.






