





"""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 				__lowercase			(			_UpperCamelCase   ,      unittest.TestCase							):







					'''simple docstring'''




					__lowerCAmelCase  =  PriorTransformer
					__lowerCAmelCase  =  '''hidden_states'''







					@property
					def  _lowerCamelCase		(						self     ):
									__a :						Union[str, Any]			       =  4
									__a :						List[Any]			       =  8
									__a :						int			       =  7

									__a :						Dict			       =  floats_tensor((batch_size, embedding_dim)     ).to(_UpperCAmelCase     )

									__a :						Dict			       =  floats_tensor((batch_size, embedding_dim)     ).to(_UpperCAmelCase     )
									__a :						List[str]			       =  floats_tensor((batch_size, num_embeddings, embedding_dim)     ).to(_UpperCAmelCase     )

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







					def  _lowerCamelCase		(						self	,			_UpperCAmelCase=0     ):
									torch.manual_seed(_UpperCAmelCase     )
									__a :						List[str]			       =  4
									__a :						Any			       =  8
									__a :						List[str]			       =  7

									__a :						str			       =  torch.randn((batch_size, embedding_dim)     ).to(_UpperCAmelCase     )

									__a :						List[str]			       =  torch.randn((batch_size, embedding_dim)     ).to(_UpperCAmelCase     )
									__a :						Optional[int]			       =  torch.randn((batch_size, num_embeddings, embedding_dim)     ).to(_UpperCAmelCase     )

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







					@property
					def  _lowerCamelCase		(						self     ):
									return (4, 8)







					@property
					def  _lowerCamelCase		(						self     ):
									return (4, 8)







					def  _lowerCamelCase		(						self     ):
									__a :						Dict			       =  {
									    '''num_attention_heads''': 2,
									    '''attention_head_dim''': 4,
									    '''num_layers''': 2,
									    '''embedding_dim''': 8,
									    '''num_embeddings''': 7,
									    '''additional_embeddings''': 4,
									}
									__a :						Dict			       =  self.dummy_input
									return init_dict, inputs_dict







					def  _lowerCamelCase		(						self     ):
									__a :						Dict			       =  PriorTransformer.from_pretrained(
									    '''hf-internal-testing/prior-dummy'''	,			output_loading_info=_UpperCAmelCase     )
									self.assertIsNotNone(_UpperCAmelCase     )
									self.assertEqual(len(loading_info['''missing_keys''']     )	,			0     )

									model.to(_UpperCAmelCase     )
									__a :						List[str]			       =  model(**self.dummy_input     )[0]

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







					def  _lowerCamelCase		(						self     ):
									__a :						Optional[int]			       =  self.prepare_init_args_and_inputs_for_common()

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

									__a :						Tuple			       =  ['''hidden_states''', '''timestep''']
									self.assertListEqual(arg_names[:2]	,			_UpperCAmelCase     )







					def  _lowerCamelCase		(						self     ):
									__a :						str			       =  PriorTransformer.from_pretrained('''hf-internal-testing/prior-dummy'''     )
									__a :						Optional[Any]			       =  model.to(_UpperCAmelCase     )

									if hasattr(_UpperCAmelCase	,			'''set_default_attn_processor'''     ):
													model.set_default_attn_processor()

									__a :						Union[str, Any]			       =  self.get_dummy_seed_input()

									with torch.no_grad():
													__a :						List[Any]			       =  model(**_UpperCAmelCase     )[0]

									__a :						Any			       =  output[0, :5].flatten().cpu()
									print(_UpperCAmelCase     )

									# 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.
									__a :						Optional[int]			       =  torch.tensor([-1.3_4_3_6, -0.2_8_7_0, 0.7_5_3_8, 0.4_3_6_8, -0.0_2_3_9]     )
									self.assertTrue(torch_all_close(_UpperCAmelCase	,			_UpperCAmelCase	,			rtol=1e-2     )     )






@slow
class 				__lowercase			(			unittest.TestCase							):







					'''simple docstring'''




					def  _lowerCamelCase		(						self	,			_UpperCAmelCase=1	,			_UpperCAmelCase=768	,			_UpperCAmelCase=77	,			_UpperCAmelCase=0     ):
									torch.manual_seed(_UpperCAmelCase     )
									__a :						List[Any]			       =  batch_size
									__a :						Optional[int]			       =  embedding_dim
									__a :						Optional[int]			       =  num_embeddings

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

									__a :						Dict			       =  torch.randn((batch_size, embedding_dim)     ).to(_UpperCAmelCase     )
									__a :						str			       =  torch.randn((batch_size, num_embeddings, embedding_dim)     ).to(_UpperCAmelCase     )

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







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







					@parameterized.expand(
					    [
					        # fmt: off
					        [13, [-0.5_8_6_1,  0.1_2_8_3, -0.0_9_3_1,  0.0_8_8_2,  0.4_4_7_6,  0.1_3_2_9, -0.0_4_9_8,  0.0_6_4_0]],
					        [37, [-0.4_9_1_3,  0.0_1_1_0, -0.0_4_8_3,  0.0_5_4_1,  0.4_9_5_4, -0.0_1_7_0,  0.0_3_5_4,  0.1_6_5_1]],
					        # fmt: on
					    ]     )
					def  _lowerCamelCase		(						self	,			_UpperCAmelCase	,			_UpperCAmelCase     ):
									__a :						Optional[int]			       =  PriorTransformer.from_pretrained('''kandinsky-community/kandinsky-2-1-prior'''	,			subfolder='''prior'''     )
									model.to(_UpperCAmelCase     )
									__a :						int			       =  self.get_dummy_seed_input(seed=_UpperCAmelCase     )

									with torch.no_grad():
													__a :						Optional[Any]			       =  model(**_UpperCAmelCase     )[0]

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

									__a :						Tuple			       =  sample[0, :8].flatten().cpu()
									print(_UpperCAmelCase     )
									__a :						Dict			       =  torch.tensor(_UpperCAmelCase     )

									assert torch_all_close(_UpperCAmelCase	,			_UpperCAmelCase	,			atol=1e-3     )






"""simple docstring"""






import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFAutoModel, is_tensorflow_text_available, is_tf_available
from transformers.models.bert.tokenization_bert import BertTokenizer
from transformers.testing_utils import require_tensorflow_text, require_tf, slow


if is_tf_available():
		import tensorflow as tf

if is_tensorflow_text_available():
		from transformers.models.bert import TFBertTokenizer


A 			=  ['''bert-base-uncased''', '''bert-base-cased''']
A 			=  '''hf-internal-testing/tiny-bert-tf-only'''

if is_tf_available():




		class 				__lowercase			(			tf.keras.Model							):







							'''simple docstring'''




							def __init__(						self	,			_UpperCAmelCase     ):
											super().__init__()
											__a :						Any			       =  tokenizer
											__a :						Optional[Any]			       =  AutoConfig.from_pretrained(_UpperCAmelCase     )
											__a :						str			       =  TFAutoModel.from_config(_UpperCAmelCase     )







							def  _lowerCamelCase		(						self	,			_UpperCAmelCase     ):
											__a :						Any			       =  self.tokenizer(_UpperCAmelCase     )
											__a :						int			       =  self.bert(**_UpperCAmelCase     )
											return out["pooler_output"]




@require_tf
@require_tensorflow_text
class 				__lowercase			(			unittest.TestCase							):







					'''simple docstring'''




					def  _lowerCamelCase		(						self     ):
									super().setUp()

									__a :						Any			       =  [
									    BertTokenizer.from_pretrained(_UpperCAmelCase     ) for checkpoint in (TOKENIZER_CHECKPOINTS * 2)
									]  # repeat for when fast_bert_tokenizer=false
									__a :						Union[str, Any]			       =  [TFBertTokenizer.from_pretrained(_UpperCAmelCase     ) for checkpoint in TOKENIZER_CHECKPOINTS] + [
									    TFBertTokenizer.from_pretrained(_UpperCAmelCase	,			use_fast_bert_tokenizer=_UpperCAmelCase     )
									    for checkpoint in TOKENIZER_CHECKPOINTS
									]
									assert len(self.tokenizers     ) == len(self.tf_tokenizers     )

									__a :						Tuple			       =  [
									    '''This is a straightforward English test sentence.''',
									    '''This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.''',
									    '''Now we\'re going to add some Chinese: 一 二 三 一二三''',
									    '''And some much more rare Chinese: 齉 堃 齉堃''',
									    '''Je vais aussi écrire en français pour tester les accents''',
									    '''Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ''',
									]
									__a :						Any			       =  list(zip(self.test_sentences	,			self.test_sentences[::-1]     )     )







					def  _lowerCamelCase		(						self     ):
									for tokenizer, tf_tokenizer in zip(self.tokenizers	,			self.tf_tokenizers     ):
													for test_inputs in (self.test_sentences, self.paired_sentences):
																	__a :						List[Any]			       =  tokenizer(_UpperCAmelCase	,			return_tensors='''tf'''	,			padding='''longest'''     )
																	__a :						List[str]			       =  tf_tokenizer(_UpperCAmelCase     )

																	for key in python_outputs.keys():
																					self.assertTrue(tf.reduce_all(python_outputs[key].shape == tf_outputs[key].shape     )     )
																					self.assertTrue(tf.reduce_all(tf.cast(python_outputs[key]	,			tf.intaa     ) == tf_outputs[key]     )     )







					@slow
					def  _lowerCamelCase		(						self     ):
									for tf_tokenizer in self.tf_tokenizers:
													__a :						Dict			       =  tf_tokenizer(self.paired_sentences     )
													__a :						str			       =  tf_tokenizer(
													    text=[sentence[0] for sentence in self.paired_sentences]	,			text_pair=[sentence[1] for sentence in self.paired_sentences]	,			)
													for key in merged_outputs.keys():
																	self.assertTrue(tf.reduce_all(tf.cast(merged_outputs[key]	,			tf.intaa     ) == separated_outputs[key]     )     )







					@slow
					def  _lowerCamelCase		(						self     ):
									for tf_tokenizer in self.tf_tokenizers:
													__a :						Tuple			       =  tf.function(_UpperCAmelCase     )
													for test_inputs in (self.test_sentences, self.paired_sentences):
																	__a :						List[str]			       =  tf.constant(_UpperCAmelCase     )
																	__a :						Tuple			       =  compiled_tokenizer(_UpperCAmelCase     )
																	__a :						Union[str, Any]			       =  tf_tokenizer(_UpperCAmelCase     )

																	for key in eager_outputs.keys():
																					self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]     )     )







					@slow
					def  _lowerCamelCase		(						self     ):
									for tf_tokenizer in self.tf_tokenizers:
													__a :						Dict			       =  ModelToSave(tokenizer=_UpperCAmelCase     )
													__a :						Tuple			       =  tf.convert_to_tensor(self.test_sentences     )
													__a :						List[str]			       =  model(_UpperCAmelCase     )  # Build model with some sample inputs
													with TemporaryDirectory() as tempdir:
																	__a :						Tuple			       =  Path(_UpperCAmelCase     ) / '''saved.model'''
																	model.save(_UpperCAmelCase     )
																	__a :						Tuple			       =  tf.keras.models.load_model(_UpperCAmelCase     )
													__a :						int			       =  loaded_model(_UpperCAmelCase     )
													# We may see small differences because the loaded model is compiled, so we need an epsilon for the test
													self.assertLessEqual(tf.reduce_max(tf.abs(out - loaded_output     )     )	,			1e-5     )