'''simple docstring'''




import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY





@is_pipeline_test
class     __magic_name__						(    unittest.TestCase):
			UpperCamelCase__				=       MODEL_FOR_MASKED_LM_MAPPING
			UpperCamelCase__				=       TF_MODEL_FOR_MASKED_LM_MAPPING







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										super().tearDown()
										# clean-up as much as possible GPU memory occupied by PyTorch
										gc.collect()
										if is_torch_available():
																	import torch

																	torch.cuda.empty_cache()







			@require_tf
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										lowercase_					:			Union[str, Any]										=       pipeline(task="""fill-mask"""				,			model="""sshleifer/tiny-distilroberta-base"""				,			top_k=2				,			framework="""tf"""    )
										lowercase_					:			Optional[int]										=       unmasker("""My name is <mask>"""    )
										self.assertEqual(
										    nested_simplify(lowercase_				,			decimals=6    )				,			[
										        {"""sequence""": """My name is grouped""", """score""": 2.1E-05, """token""": 38015, """token_str""": """ grouped"""},
										        {"""sequence""": """My name is accuser""", """score""": 2.1E-05, """token""": 25506, """token_str""": """ accuser"""},
										    ]				,			)

										lowercase_					:			List[Any]										=       unmasker("""The largest city in France is <mask>"""    )
										self.assertEqual(
										    nested_simplify(lowercase_				,			decimals=6    )				,			[
										        {
										            """sequence""": """The largest city in France is grouped""",
										            """score""": 2.1E-05,
										            """token""": 38015,
										            """token_str""": """ grouped""",
										        },
										        {
										            """sequence""": """The largest city in France is accuser""",
										            """score""": 2.1E-05,
										            """token""": 25506,
										            """token_str""": """ accuser""",
										        },
										    ]				,			)

										lowercase_					:			str										=       unmasker("""My name is <mask>"""				,			targets=[""" Patrick""", """ Clara""", """ Teven"""]				,			top_k=3    )
										self.assertEqual(
										    nested_simplify(lowercase_				,			decimals=6    )				,			[
										        {"""sequence""": """My name is Clara""", """score""": 2E-05, """token""": 13606, """token_str""": """ Clara"""},
										        {"""sequence""": """My name is Patrick""", """score""": 2E-05, """token""": 3499, """token_str""": """ Patrick"""},
										        {"""sequence""": """My name is Te""", """score""": 1.9E-05, """token""": 2941, """token_str""": """ Te"""},
										    ]				,			)







			@require_torch
			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Optional[Any]										=       pipeline(task="""fill-mask"""				,			model="""sshleifer/tiny-distilroberta-base"""				,			top_k=2				,			framework="""pt"""    )

										lowercase_					:			Union[str, Any]										=       unmasker("""My name is <mask>"""    )
										self.assertEqual(
										    nested_simplify(lowercase_				,			decimals=6    )				,			[
										        {"""sequence""": """My name is Maul""", """score""": 2.2E-05, """token""": 35676, """token_str""": """ Maul"""},
										        {"""sequence""": """My name isELS""", """score""": 2.2E-05, """token""": 16416, """token_str""": """ELS"""},
										    ]				,			)

										lowercase_					:			Union[str, Any]										=       unmasker("""The largest city in France is <mask>"""    )
										self.assertEqual(
										    nested_simplify(lowercase_				,			decimals=6    )				,			[
										        {
										            """sequence""": """The largest city in France is Maul""",
										            """score""": 2.2E-05,
										            """token""": 35676,
										            """token_str""": """ Maul""",
										        },
										        {"""sequence""": """The largest city in France isELS""", """score""": 2.2E-05, """token""": 16416, """token_str""": """ELS"""},
										    ]				,			)

										lowercase_					:			int										=       unmasker("""My name is <mask>"""				,			targets=[""" Patrick""", """ Clara""", """ Teven"""]				,			top_k=3    )
										self.assertEqual(
										    nested_simplify(lowercase_				,			decimals=6    )				,			[
										        {"""sequence""": """My name is Patrick""", """score""": 2.1E-05, """token""": 3499, """token_str""": """ Patrick"""},
										        {"""sequence""": """My name is Te""", """score""": 2E-05, """token""": 2941, """token_str""": """ Te"""},
										        {"""sequence""": """My name is Clara""", """score""": 2E-05, """token""": 13606, """token_str""": """ Clara"""},
										    ]				,			)

										lowercase_					:			Optional[Any]										=       unmasker("""My name is <mask> <mask>"""				,			top_k=2    )

										self.assertEqual(
										    nested_simplify(lowercase_				,			decimals=6    )				,			[
										        [
										            {
										                """score""": 2.2E-05,
										                """token""": 35676,
										                """token_str""": """ Maul""",
										                """sequence""": """<s>My name is Maul<mask></s>""",
										            },
										            {"""score""": 2.2E-05, """token""": 16416, """token_str""": """ELS""", """sequence""": """<s>My name isELS<mask></s>"""},
										        ],
										        [
										            {
										                """score""": 2.2E-05,
										                """token""": 35676,
										                """token_str""": """ Maul""",
										                """sequence""": """<s>My name is<mask> Maul</s>""",
										            },
										            {"""score""": 2.2E-05, """token""": 16416, """token_str""": """ELS""", """sequence""": """<s>My name is<mask>ELS</s>"""},
										        ],
										    ]				,			)







			@require_torch_gpu
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			List[Any]										=       pipeline("""fill-mask"""				,			model="""hf-internal-testing/tiny-random-distilbert"""				,			device=0				,			framework="""pt"""    )

										# convert model to fp16
										pipe.model.half()

										lowercase_					:			Dict										=       pipe("""Paris is the [MASK] of France."""    )
										# We actually don't care about the result, we just want to make sure
										# it works, meaning the float16 tensor got casted back to float32
										# for postprocessing.
										self.assertIsInstance(lowercase_				,			lowercase_    )







			@slow
			@require_torch
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										lowercase_					:			Tuple										=       pipeline(task="""fill-mask"""				,			model="""distilroberta-base"""				,			top_k=2				,			framework="""pt"""    )
										self.run_large_test(lowercase_    )







			@slow
			@require_tf
			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										lowercase_					:			Optional[Any]										=       pipeline(task="""fill-mask"""				,			model="""distilroberta-base"""				,			top_k=2				,			framework="""tf"""    )
										self.run_large_test(lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  str				,			lowercase_				:  int    ):
										lowercase_					:			int										=       unmasker("""My name is <mask>"""    )
										self.assertEqual(
										    nested_simplify(lowercase_    )				,			[
										        {"""sequence""": """My name is John""", """score""": 0.0_08, """token""": 610, """token_str""": """ John"""},
										        {"""sequence""": """My name is Chris""", """score""": 0.0_07, """token""": 1573, """token_str""": """ Chris"""},
										    ]				,			)
										lowercase_					:			Optional[Any]										=       unmasker("""The largest city in France is <mask>"""    )
										self.assertEqual(
										    nested_simplify(lowercase_    )				,			[
										        {
										            """sequence""": """The largest city in France is Paris""",
										            """score""": 0.2_51,
										            """token""": 2201,
										            """token_str""": """ Paris""",
										        },
										        {
										            """sequence""": """The largest city in France is Lyon""",
										            """score""": 0.2_14,
										            """token""": 12790,
										            """token_str""": """ Lyon""",
										        },
										    ]				,			)

										lowercase_					:			Optional[int]										=       unmasker("""My name is <mask>"""				,			targets=[""" Patrick""", """ Clara""", """ Teven"""]				,			top_k=3    )
										self.assertEqual(
										    nested_simplify(lowercase_    )				,			[
										        {"""sequence""": """My name is Patrick""", """score""": 0.0_05, """token""": 3499, """token_str""": """ Patrick"""},
										        {"""sequence""": """My name is Clara""", """score""": 0.0_00, """token""": 13606, """token_str""": """ Clara"""},
										        {"""sequence""": """My name is Te""", """score""": 0.0_00, """token""": 2941, """token_str""": """ Te"""},
										    ]				,			)







			@require_torch
			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			str										=       pipeline(task="""fill-mask"""				,			model="""sshleifer/tiny-distilroberta-base"""				,			framework="""pt"""    )
										lowercase_					:			Any										=       None
										lowercase_					:			Dict										=       None
										self.run_pipeline_test(lowercase_				,			[]    )







			@require_tf
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										lowercase_					:			Tuple										=       pipeline(task="""fill-mask"""				,			model="""sshleifer/tiny-distilroberta-base"""				,			framework="""tf"""    )
										lowercase_					:			Dict										=       None
										lowercase_					:			List[Any]										=       None
										self.run_pipeline_test(lowercase_				,			[]    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int				,			lowercase_				:  str				,			lowercase_				:  int				,			lowercase_				:  Tuple    ):
										if tokenizer is None or tokenizer.mask_token_id is None:
																	self.skipTest("""The provided tokenizer has no mask token, (probably reformer or wav2vec2)"""    )

										lowercase_					:			Union[str, Any]										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_    )
										lowercase_					:			Tuple										=       [
										    f'''This is another {tokenizer.mask_token} test''',
										]
										return fill_masker, examples







			def 	SCREAMING_SNAKE_CASE_     (	self				:  str				,			lowercase_				:  Optional[int]				,			lowercase_				:  List[Any]    ):
										lowercase_					:			List[str]										=       fill_masker.tokenizer
										lowercase_					:			List[Any]										=       fill_masker.model

										lowercase_					:			Union[str, Any]										=       fill_masker(
										    f'''This is a {tokenizer.mask_token}'''				,			)
										self.assertEqual(
										    lowercase_				,			[
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										    ]				,			)

										lowercase_					:			Dict										=       fill_masker([f'''This is a {tokenizer.mask_token}''']    )
										self.assertEqual(
										    lowercase_				,			[
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										    ]				,			)

										lowercase_					:			Any										=       fill_masker([f'''This is a {tokenizer.mask_token}''', f'''Another {tokenizer.mask_token} great test.''']    )
										self.assertEqual(
										    lowercase_				,			[
										        [
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        ],
										        [
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        ],
										    ]				,			)

										with self.assertRaises(lowercase_    ):
																	fill_masker([None]    )
										# No mask_token is not supported
										with self.assertRaises(lowercase_    ):
																	fill_masker("""This is"""    )

										self.run_test_top_k(lowercase_				,			lowercase_    )
										self.run_test_targets(lowercase_				,			lowercase_    )
										self.run_test_top_k_targets(lowercase_				,			lowercase_    )
										self.fill_mask_with_duplicate_targets_and_top_k(lowercase_				,			lowercase_    )
										self.fill_mask_with_multiple_masks(lowercase_				,			lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  str				,			lowercase_				:  Any				,			lowercase_				:  Tuple    ):
										lowercase_					:			str										=       tokenizer.get_vocab()
										lowercase_					:			Dict										=       sorted(vocab.keys()    )[:2]
										# Pipeline argument
										lowercase_					:			Any										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_				,			targets=lowercase_    )
										lowercase_					:			int										=       fill_masker(f'''This is a {tokenizer.mask_token}'''    )
										self.assertEqual(
										    lowercase_				,			[
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										    ]				,			)
										lowercase_					:			int										=       {vocab[el] for el in targets}
										self.assertEqual({el["""token"""] for el in outputs}				,			lowercase_    )
										lowercase_					:			Any										=       [tokenizer.decode([x]    ) for x in target_ids]
										self.assertEqual({el["""token_str"""] for el in outputs}				,			set(lowercase_    )    )

										# Call argument
										lowercase_					:			str										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_    )
										lowercase_					:			Dict										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			targets=lowercase_    )
										self.assertEqual(
										    lowercase_				,			[
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										    ]				,			)
										lowercase_					:			Union[str, Any]										=       {vocab[el] for el in targets}
										self.assertEqual({el["""token"""] for el in outputs}				,			lowercase_    )
										lowercase_					:			str										=       [tokenizer.decode([x]    ) for x in target_ids]
										self.assertEqual({el["""token_str"""] for el in outputs}				,			set(lowercase_    )    )

										# Score equivalence
										lowercase_					:			Optional[Any]										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			targets=lowercase_    )
										lowercase_					:			Tuple										=       [top_mask["""token_str"""] for top_mask in outputs]
										lowercase_					:			Union[str, Any]										=       [top_mask["""score"""] for top_mask in outputs]

										# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
										if set(lowercase_    ) == set(lowercase_    ):
																	lowercase_					:			Any										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			targets=lowercase_    )
																	lowercase_					:			Optional[Any]										=       [top_mask["""score"""] for top_mask in unmasked_targets]
																	self.assertEqual(nested_simplify(lowercase_    )				,			nested_simplify(lowercase_    )    )

										# Raises with invalid
										with self.assertRaises(lowercase_    ):
																	lowercase_					:			Tuple										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			targets=[]    )
										# For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
										if "" not in tokenizer.get_vocab():
																	with self.assertRaises(lowercase_    ):
																								lowercase_					:			str										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			targets=[""""""]    )
																	with self.assertRaises(lowercase_    ):
																								lowercase_					:			Any										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			targets=""""""    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]				,			lowercase_				:  Tuple				,			lowercase_				:  str    ):
										lowercase_					:			List[Any]										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_				,			top_k=2    )
										lowercase_					:			str										=       fill_masker(f'''This is a {tokenizer.mask_token}'''    )
										self.assertEqual(
										    lowercase_				,			[
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										    ]				,			)

										lowercase_					:			str										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_    )
										lowercase_					:			Any										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			top_k=2    )
										self.assertEqual(
										    lowercase_				,			[
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										    ]				,			)
										self.assertEqual(nested_simplify(lowercase_    )				,			nested_simplify(lowercase_    )    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict				,			lowercase_				:  int				,			lowercase_				:  Union[str, Any]    ):
										lowercase_					:			Dict										=       tokenizer.get_vocab()
										lowercase_					:			List[Any]										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_    )

										# top_k=2, ntargets=3
										lowercase_					:			Any										=       sorted(vocab.keys()    )[:3]
										lowercase_					:			Optional[Any]										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			top_k=2				,			targets=lowercase_    )

										# If we use the most probably targets, and filter differently, we should still
										# have the same results
										lowercase_					:			str										=       [el["""token_str"""] for el in sorted(lowercase_				,			key=lambda lowercase_    : x["score"]				,			reverse=lowercase_    )]
										# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
										if set(lowercase_    ).issubset(lowercase_    ):
																	lowercase_					:			Tuple										=       fill_masker(f'''This is a {tokenizer.mask_token}'''				,			top_k=3				,			targets=lowercase_    )
																	# They should yield exactly the same result
																	self.assertEqual(nested_simplify(lowercase_    )				,			nested_simplify(lowercase_    )    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]				,			lowercase_				:  Optional[Any]				,			lowercase_				:  Any    ):
										lowercase_					:			Optional[Any]										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_    )
										lowercase_					:			List[Any]										=       tokenizer.get_vocab()
										# String duplicates + id duplicates
										lowercase_					:			int										=       sorted(vocab.keys()    )[:3]
										lowercase_					:			Optional[int]										=       [targets[0], targets[1], targets[0], targets[2], targets[1]]
										lowercase_					:			Tuple										=       fill_masker(f'''My name is {tokenizer.mask_token}'''				,			targets=lowercase_				,			top_k=10    )

										# The target list contains duplicates, so we can't output more
										# than them
										self.assertEqual(len(lowercase_    )				,			3    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any				,			lowercase_				:  Dict				,			lowercase_				:  Tuple    ):
										lowercase_					:			List[Any]										=       FillMaskPipeline(model=lowercase_				,			tokenizer=lowercase_    )

										lowercase_					:			int										=       fill_masker(
										    f'''This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}'''				,			top_k=2    )
										self.assertEqual(
										    lowercase_				,			[
										        [
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        ],
										        [
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        ],
										        [
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										            {"""sequence""": ANY(lowercase_    ), """score""": ANY(lowercase_    ), """token""": ANY(lowercase_    ), """token_str""": ANY(lowercase_    )},
										        ],
										    ]				,			)






'''simple docstring'''




import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()





class     __magic_name__						(    _UpperCAmelCase,     unittest.TestCase):
			UpperCamelCase__				=       PriorTransformer
			UpperCamelCase__				=       '''hidden_states'''







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										lowercase_					:			Optional[Any]										=       4
										lowercase_					:			int										=       8
										lowercase_					:			Union[str, Any]										=       7

										lowercase_					:			List[str]										=       floats_tensor((batch_size, embedding_dim)    ).to(lowercase_    )

										lowercase_					:			Tuple										=       floats_tensor((batch_size, embedding_dim)    ).to(lowercase_    )
										lowercase_					:			List[str]										=       floats_tensor((batch_size, num_embeddings, embedding_dim)    ).to(lowercase_    )

										return {
										    "hidden_states": hidden_states,
										    "timestep": 2,
										    "proj_embedding": proj_embedding,
										    "encoder_hidden_states": encoder_hidden_states,
										}







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]				,			lowercase_				:  Optional[Any]=0    ):
										torch.manual_seed(lowercase_    )
										lowercase_					:			int										=       4
										lowercase_					:			Any										=       8
										lowercase_					:			Tuple										=       7

										lowercase_					:			Union[str, Any]										=       torch.randn((batch_size, embedding_dim)    ).to(lowercase_    )

										lowercase_					:			Tuple										=       torch.randn((batch_size, embedding_dim)    ).to(lowercase_    )
										lowercase_					:			Union[str, Any]										=       torch.randn((batch_size, num_embeddings, embedding_dim)    ).to(lowercase_    )

										return {
										    "hidden_states": hidden_states,
										    "timestep": 2,
										    "proj_embedding": proj_embedding,
										    "encoder_hidden_states": encoder_hidden_states,
										}







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  str    ):
										return (4, 8)







			@property
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										return (4, 8)







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										lowercase_					:			Optional[Any]										=       {
										    """num_attention_heads""": 2,
										    """attention_head_dim""": 4,
										    """num_layers""": 2,
										    """embedding_dim""": 8,
										    """num_embeddings""": 7,
										    """additional_embeddings""": 4,
										}
										lowercase_					:			List[Any]										=       self.dummy_input
										return init_dict, inputs_dict







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										lowercase_						,						lowercase_					:			Tuple										=       PriorTransformer.from_pretrained(
										    """hf-internal-testing/prior-dummy"""				,			output_loading_info=lowercase_    )
										self.assertIsNotNone(lowercase_    )
										self.assertEqual(len(loading_info["""missing_keys"""]    )				,			0    )

										model.to(lowercase_    )
										lowercase_					:			Optional[int]										=       model(**self.dummy_input    )[0]

										assert hidden_states is not None, "Make sure output is not None"







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										lowercase_						,						lowercase_					:			str										=       self.prepare_init_args_and_inputs_for_common()

										lowercase_					:			List[Any]										=       self.model_class(**lowercase_    )
										lowercase_					:			Any										=       inspect.signature(model.forward    )
										# signature.parameters is an OrderedDict => so arg_names order is deterministic
										lowercase_					:			Any										=       [*signature.parameters.keys()]

										lowercase_					:			Optional[Any]										=       ["""hidden_states""", """timestep"""]
										self.assertListEqual(arg_names[:2]				,			lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[str]    ):
										lowercase_					:			Optional[Any]										=       PriorTransformer.from_pretrained("""hf-internal-testing/prior-dummy"""    )
										lowercase_					:			Any										=       model.to(lowercase_    )

										if hasattr(lowercase_				,			"""set_default_attn_processor"""    ):
																	model.set_default_attn_processor()

										lowercase_					:			Any										=       self.get_dummy_seed_input()

										with torch.no_grad():
																	lowercase_					:			List[str]										=       model(**lowercase_    )[0]

										lowercase_					:			Tuple										=       output[0, :5].flatten().cpu()
										print(lowercase_    )

										# Since the VAE Gaussian prior's generator is seeded on the appropriate device,
										# the expected output slices are not the same for CPU and GPU.
										lowercase_					:			Dict										=       torch.tensor([-1.34_36, -0.28_70, 0.75_38, 0.43_68, -0.02_39]    )
										self.assertTrue(torch_all_close(lowercase_				,			lowercase_				,			rtol=1E-2    )    )





@slow
class     __magic_name__						(    unittest.TestCase):







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]				,			lowercase_				:  Union[str, Any]=1				,			lowercase_				:  Optional[Any]=768				,			lowercase_				:  Tuple=77				,			lowercase_				:  Optional[int]=0    ):
										torch.manual_seed(lowercase_    )
										lowercase_					:			Optional[Any]										=       batch_size
										lowercase_					:			int										=       embedding_dim
										lowercase_					:			int										=       num_embeddings

										lowercase_					:			Dict										=       torch.randn((batch_size, embedding_dim)    ).to(lowercase_    )

										lowercase_					:			Any										=       torch.randn((batch_size, embedding_dim)    ).to(lowercase_    )
										lowercase_					:			List[Any]										=       torch.randn((batch_size, num_embeddings, embedding_dim)    ).to(lowercase_    )

										return {
										    "hidden_states": hidden_states,
										    "timestep": 2,
										    "proj_embedding": proj_embedding,
										    "encoder_hidden_states": encoder_hidden_states,
										}







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Any    ):
										# clean up the VRAM after each test
										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()







			@parameterized.expand(
			    [
			        # fmt: off
			        [13, [-0.58_61,  0.12_83, -0.09_31,  0.08_82,  0.44_76,  0.13_29, -0.04_98,  0.06_40]],
			        [37, [-0.49_13,  0.01_10, -0.04_83,  0.05_41,  0.49_54, -0.01_70,  0.03_54,  0.16_51]],
			        # fmt: on
			    ]    )
			def 	SCREAMING_SNAKE_CASE_     (	self				:  int				,			lowercase_				:  Any				,			lowercase_				:  List[str]    ):
										lowercase_					:			List[Any]										=       PriorTransformer.from_pretrained("""kandinsky-community/kandinsky-2-1-prior"""				,			subfolder="""prior"""    )
										model.to(lowercase_    )
										lowercase_					:			Optional[Any]										=       self.get_dummy_seed_input(seed=lowercase_    )

										with torch.no_grad():
																	lowercase_					:			Tuple										=       model(**lowercase_    )[0]

										assert list(sample.shape    ) == [1, 768]

										lowercase_					:			Union[str, Any]										=       sample[0, :8].flatten().cpu()
										print(lowercase_    )
										lowercase_					:			Optional[Any]										=       torch.tensor(lowercase_    )

										assert torch_all_close(lowercase_				,			lowercase_				,			atol=1E-3    )





