"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
	import tensorflow as tf

	from transformers.models.distilbert.modeling_tf_distilbert import (
	    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TFDistilBertForMaskedLM,
	    TFDistilBertForMultipleChoice,
	    TFDistilBertForQuestionAnswering,
	    TFDistilBertForSequenceClassification,
	    TFDistilBertForTokenClassification,
	    TFDistilBertModel,
	)

class 		UpperCamelCase      :



				def __init__(self    :				List[str]     ,  _A    :				Any     ,  )   ->     Optional[Any]:
					__snake_case      :							Any						=			parent
					__snake_case      :							Union[str, Any]						=			13
					__snake_case      :							Union[str, Any]						=			7
					__snake_case      :							List[str]						=			True
					__snake_case      :							str						=			True
					__snake_case      :							str						=			False
					__snake_case      :							Union[str, Any]						=			True
					__snake_case      :							Any						=			99
					__snake_case      :							Any						=			32
					__snake_case      :							int						=			2
					__snake_case      :							Tuple						=			4
					__snake_case      :							Tuple						=			37
					__snake_case      :							Optional[Any]						=			'gelu'
					__snake_case      :							Optional[int]						=			0.1
					__snake_case      :							Optional[int]						=			0.1
					__snake_case      :							Optional[Any]						=			5_12
					__snake_case      :							str						=			16
					__snake_case      :							Optional[Any]						=			2
					__snake_case      :							Any						=			0.02
					__snake_case      :							int						=			3
					__snake_case      :							Dict						=			4
					__snake_case      :							List[Any]						=			None



				def  _lowercase					(self    :				int)   ->     List[str]:
					__snake_case      :							Any						=			ids_tensor([self.batch_size, self.seq_length]     ,  self.vocab_size)

					__snake_case      :							List[str]						=			None
					if self.use_input_mask:
						__snake_case      :							int						=			random_attention_mask([self.batch_size, self.seq_length])

					__snake_case      :							List[str]						=			None
					__snake_case      :							Dict						=			None
					__snake_case      :							Dict						=			None
					if self.use_labels:
						__snake_case      :							List[Any]						=			ids_tensor([self.batch_size]     ,  self.type_sequence_label_size)
						__snake_case      :							int						=			ids_tensor([self.batch_size, self.seq_length]     ,  self.num_labels)
						__snake_case      :							List[str]						=			ids_tensor([self.batch_size]     ,  self.num_choices)

					__snake_case      :							Dict						=			DistilBertConfig(
					    vocab_size=self.vocab_size     ,  dim=self.hidden_size     ,  n_layers=self.num_hidden_layers     ,  n_heads=self.num_attention_heads     ,  hidden_dim=self.intermediate_size     ,  hidden_act=self.hidden_act     ,  dropout=self.hidden_dropout_prob     ,  attention_dropout=self.attention_probs_dropout_prob     ,  max_position_embeddings=self.max_position_embeddings     ,  initializer_range=self.initializer_range     ,  )

					return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



				def  _lowercase					(self    :				List[Any]     ,  _A    :				Tuple     ,  _A    :				int     ,  _A    :				Any     ,  _A    :				int     ,  _A    :				List[str]     ,  _A    :				str)   ->     int:
					__snake_case      :							Tuple						=			TFDistilBertModel(config=_A)
					__snake_case      :							Union[str, Any]						=			{'input_ids': input_ids, 'attention_mask': input_mask}

					__snake_case      :							Optional[int]						=			model(_A)

					__snake_case      :							Union[str, Any]						=			[input_ids, input_mask]

					__snake_case      :							Any						=			model(_A)

					self.parent.assertEqual(result.last_hidden_state.shape     ,  (self.batch_size, self.seq_length, self.hidden_size))



				def  _lowercase					(self    :				List[str]     ,  _A    :				List[Any]     ,  _A    :				Dict     ,  _A    :				str     ,  _A    :				List[Any]     ,  _A    :				List[str]     ,  _A    :				Optional[int])   ->     int:
					__snake_case      :							Optional[Any]						=			TFDistilBertForMaskedLM(config=_A)
					__snake_case      :							Union[str, Any]						=			{'input_ids': input_ids, 'attention_mask': input_mask}
					__snake_case      :							Dict						=			model(_A)
					self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.seq_length, self.vocab_size))



				def  _lowercase					(self    :				str     ,  _A    :				int     ,  _A    :				int     ,  _A    :				Dict     ,  _A    :				Union[str, Any]     ,  _A    :				Optional[int]     ,  _A    :				List[Any])   ->     Optional[Any]:
					__snake_case      :							int						=			TFDistilBertForQuestionAnswering(config=_A)
					__snake_case      :							str						=			{
					    'input_ids': input_ids,
					    'attention_mask': input_mask,
					}
					__snake_case      :							Dict						=			model(_A)
					self.parent.assertEqual(result.start_logits.shape     ,  (self.batch_size, self.seq_length))
					self.parent.assertEqual(result.end_logits.shape     ,  (self.batch_size, self.seq_length))



				def  _lowercase					(self    :				List[Any]     ,  _A    :				Any     ,  _A    :				Tuple     ,  _A    :				Dict     ,  _A    :				Union[str, Any]     ,  _A    :				Any     ,  _A    :				List[str])   ->     int:
					__snake_case      :							int						=			self.num_labels
					__snake_case      :							str						=			TFDistilBertForSequenceClassification(_A)
					__snake_case      :							Any						=			{'input_ids': input_ids, 'attention_mask': input_mask}
					__snake_case      :							Optional[int]						=			model(_A)
					self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.num_labels))



				def  _lowercase					(self    :				Optional[int]     ,  _A    :				List[Any]     ,  _A    :				List[str]     ,  _A    :				int     ,  _A    :				Optional[Any]     ,  _A    :				Optional[Any]     ,  _A    :				int)   ->     Tuple:
					__snake_case      :							str						=			self.num_choices
					__snake_case      :							List[Any]						=			TFDistilBertForMultipleChoice(_A)
					__snake_case      :							int						=			tf.tile(tf.expand_dims(_A     ,  1)     ,  (1, self.num_choices, 1))
					__snake_case      :							Union[str, Any]						=			tf.tile(tf.expand_dims(_A     ,  1)     ,  (1, self.num_choices, 1))
					__snake_case      :							int						=			{
					    'input_ids': multiple_choice_inputs_ids,
					    'attention_mask': multiple_choice_input_mask,
					}
					__snake_case      :							str						=			model(_A)
					self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.num_choices))



				def  _lowercase					(self    :				List[Any]     ,  _A    :				str     ,  _A    :				Any     ,  _A    :				Optional[int]     ,  _A    :				Tuple     ,  _A    :				Dict     ,  _A    :				Union[str, Any])   ->     Optional[Any]:
					__snake_case      :							Any						=			self.num_labels
					__snake_case      :							Optional[int]						=			TFDistilBertForTokenClassification(_A)
					__snake_case      :							int						=			{'input_ids': input_ids, 'attention_mask': input_mask}
					__snake_case      :							Optional[Any]						=			model(_A)
					self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.seq_length, self.num_labels))






				def  _lowercase					(self    :				Tuple)   ->     Any:
					__snake_case      :							Union[str, Any]						=			self.prepare_config_and_inputs()
					((__snake_case)		,    (__snake_case)		,    (__snake_case)		,    (__snake_case)		,    (__snake_case)		,    (__snake_case))      :							Optional[Any]						=			config_and_inputs
					__snake_case      :							Tuple						=			{'input_ids': input_ids, 'attention_mask': input_mask}
					return config, inputs_dict







@require_tf
class 		UpperCamelCase      (   lowercase				,     lowercase				,     unittest.TestCase						):
				UpperCAmelCase    :				Optional[Any]   =		(
				    (
				        TFDistilBertModel,
				        TFDistilBertForMaskedLM,
				        TFDistilBertForQuestionAnswering,
				        TFDistilBertForSequenceClassification,
				        TFDistilBertForTokenClassification,
				        TFDistilBertForMultipleChoice,
				    )
				    if is_tf_available()
				    else None
				)
				UpperCAmelCase    :				str   =		(
				    {
				        """feature-extraction""": TFDistilBertModel,
				        """fill-mask""": TFDistilBertForMaskedLM,
				        """question-answering""": TFDistilBertForQuestionAnswering,
				        """text-classification""": TFDistilBertForSequenceClassification,
				        """token-classification""": TFDistilBertForTokenClassification,
				        """zero-shot""": TFDistilBertForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				UpperCAmelCase    :				Tuple   =		False
				UpperCAmelCase    :				List[str]   =		False



				def  _lowercase					(self    :				Optional[int])   ->     Optional[int]:
					__snake_case      :							Dict						=			TFDistilBertModelTester(self)
					__snake_case      :							List[Any]						=			ConfigTester(self     ,  config_class=_A     ,  dim=37)



				def  _lowercase					(self    :				Tuple)   ->     Union[str, Any]:
					self.config_tester.run_common_tests()



				def  _lowercase					(self    :				Optional[int])   ->     List[str]:
					__snake_case      :							List[str]						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_model(*_A)



				def  _lowercase					(self    :				List[str])   ->     Optional[Any]:
					__snake_case      :							Optional[int]						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_masked_lm(*_A)



				def  _lowercase					(self    :				Union[str, Any])   ->     Dict:
					__snake_case      :							Dict						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_question_answering(*_A)



				def  _lowercase					(self    :				List[str])   ->     str:
					__snake_case      :							Optional[Any]						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_sequence_classification(*_A)



				def  _lowercase					(self    :				Dict)   ->     Dict:
					__snake_case      :							Union[str, Any]						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_multiple_choice(*_A)



				def  _lowercase					(self    :				Dict)   ->     Dict:
					__snake_case      :							Union[str, Any]						=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_distilbert_for_token_classification(*_A)






				@slow
				def  _lowercase					(self    :				Optional[int])   ->     List[str]:
					for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]):
						__snake_case      :							int						=			TFDistilBertModel.from_pretrained(_A)
						self.assertIsNotNone(_A)







@require_tf
class 		UpperCamelCase      (   unittest.TestCase						):



				@slow
				def  _lowercase					(self    :				Optional[Any])   ->     List[str]:
					__snake_case      :							Dict						=			TFDistilBertModel.from_pretrained('distilbert-base-uncased')
					__snake_case      :							List[Any]						=			tf.constant([[0, 1, 2, 3, 4, 5]])
					__snake_case      :							Dict						=			model(_A)[0]

					__snake_case      :							Optional[Any]						=			[1, 6, 7_68]
					self.assertEqual(output.shape     ,  _A)

					__snake_case      :							str						=			tf.constant(
					    [
					        [
					            [0.19_261_885, -0.13_732_955, 0.4_119_799],
					            [0.22_150_156, -0.07_422_661, 0.39_037_204],
					            [0.22_756_018, -0.0_896_414, 0.3_701_467],
					        ]
					    ])
					tf.debugging.assert_near(output[:, :3, :3]     ,  _A     ,  atol=1E-4)






"""simple docstring"""







import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


_a					:					List[str]=			False

class 		UpperCamelCase      (   unittest.TestCase						):



				def  _lowercase					(self    :				List[str])   ->     List[Any]:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()



				@property
				def  _lowercase					(self    :				Optional[Any])   ->     Union[str, Any]:
					return 12



				@property
				def  _lowercase					(self    :				Dict)   ->     Union[str, Any]:
					return 12



				@property
				def  _lowercase					(self    :				int)   ->     Tuple:
					return 32



				@property
				def  _lowercase					(self    :				Optional[int])   ->     Dict:
					torch.manual_seed(0)
					__snake_case      :							Any						=			VQModel(
					    block_out_channels=[32, 64]     ,  in_channels=3     ,  out_channels=3     ,  down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']     ,  up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']     ,  latent_channels=3     ,  num_vq_embeddings=self.num_embed     ,  vq_embed_dim=3     ,  )
					return model



				@property
				def  _lowercase					(self    :				List[Any])   ->     Optional[int]:
					__snake_case      :							Dict						=			CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')
					return tokenizer



				@property
				def  _lowercase					(self    :				Union[str, Any])   ->     Optional[int]:
					torch.manual_seed(0)
					__snake_case      :							Any						=			CLIPTextConfig(
					    bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=self.text_embedder_hidden_size     ,  intermediate_size=37     ,  layer_norm_eps=1E-05     ,  num_attention_heads=4     ,  num_hidden_layers=5     ,  pad_token_id=1     ,  vocab_size=10_00     ,  )
					return CLIPTextModel(_A)



				@property
				def  _lowercase					(self    :				Union[str, Any])   ->     Dict:
					torch.manual_seed(0)

					__snake_case      :							Any						=			12
					__snake_case      :							int						=			12

					__snake_case      :							List[Any]						=			{
					    'attention_bias': True,
					    'cross_attention_dim': 32,
					    'attention_head_dim': height * width,
					    'num_attention_heads': 1,
					    'num_vector_embeds': self.num_embed,
					    'num_embeds_ada_norm': self.num_embeds_ada_norm,
					    'norm_num_groups': 32,
					    'sample_size': width,
					    'activation_fn': 'geglu-approximate',
					}

					__snake_case      :							Union[str, Any]						=			TransformeraDModel(**_A)
					return model



				def  _lowercase					(self    :				Union[str, Any])   ->     Dict:
					__snake_case      :							Tuple						=			'cpu'

					__snake_case      :							List[str]						=			self.dummy_vqvae
					__snake_case      :							str						=			self.dummy_text_encoder
					__snake_case      :							Optional[Any]						=			self.dummy_tokenizer
					__snake_case      :							Dict						=			self.dummy_transformer
					__snake_case      :							Optional[int]						=			VQDiffusionScheduler(self.num_embed)
					__snake_case      :							List[Any]						=			LearnedClassifierFreeSamplingEmbeddings(learnable=_A)

					__snake_case      :							List[Any]						=			VQDiffusionPipeline(
					    vqvae=_A     ,  text_encoder=_A     ,  tokenizer=_A     ,  transformer=_A     ,  scheduler=_A     ,  learned_classifier_free_sampling_embeddings=_A     ,  )
					__snake_case      :							List[Any]						=			pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case      :							Optional[Any]						=			'teddy bear playing in the pool'

					__snake_case      :							str						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Union[str, Any]						=			pipe([prompt]     ,  generator=_A     ,  num_inference_steps=2     ,  output_type='np')
					__snake_case      :							Optional[int]						=			output.images

					__snake_case      :							int						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Tuple						=			pipe(
					    [prompt]     ,  generator=_A     ,  output_type='np'     ,  return_dict=_A     ,  num_inference_steps=2)[0]

					__snake_case      :							str						=			image[0, -3:, -3:, -1]
					__snake_case      :							Union[str, Any]						=			image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 24, 24, 3)

					__snake_case      :							str						=			np.array([0.6_551, 0.6_168, 0.5_008, 0.5_676, 0.5_659, 0.4_295, 0.6_073, 0.5_599, 0.4_992])

					assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2






				def  _lowercase					(self    :				Tuple)   ->     Optional[int]:
					__snake_case      :							Optional[Any]						=			'cpu'

					__snake_case      :							Optional[int]						=			self.dummy_vqvae
					__snake_case      :							List[str]						=			self.dummy_text_encoder
					__snake_case      :							Optional[int]						=			self.dummy_tokenizer
					__snake_case      :							Optional[Any]						=			self.dummy_transformer
					__snake_case      :							Union[str, Any]						=			VQDiffusionScheduler(self.num_embed)
					__snake_case      :							Optional[int]						=			LearnedClassifierFreeSamplingEmbeddings(
					    learnable=_A     ,  hidden_size=self.text_embedder_hidden_size     ,  length=tokenizer.model_max_length)

					__snake_case      :							Union[str, Any]						=			VQDiffusionPipeline(
					    vqvae=_A     ,  text_encoder=_A     ,  tokenizer=_A     ,  transformer=_A     ,  scheduler=_A     ,  learned_classifier_free_sampling_embeddings=_A     ,  )
					__snake_case      :							Union[str, Any]						=			pipe.to(_A)
					pipe.set_progress_bar_config(disable=_A)

					__snake_case      :							Union[str, Any]						=			'teddy bear playing in the pool'

					__snake_case      :							Optional[int]						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Tuple						=			pipe([prompt]     ,  generator=_A     ,  num_inference_steps=2     ,  output_type='np')
					__snake_case      :							Optional[Any]						=			output.images

					__snake_case      :							str						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							Dict						=			pipe(
					    [prompt]     ,  generator=_A     ,  output_type='np'     ,  return_dict=_A     ,  num_inference_steps=2)[0]

					__snake_case      :							int						=			image[0, -3:, -3:, -1]
					__snake_case      :							Tuple						=			image_from_tuple[0, -3:, -3:, -1]

					assert image.shape == (1, 24, 24, 3)

					__snake_case      :							Optional[Any]						=			np.array([0.6_693, 0.6_075, 0.4_959, 0.5_701, 0.5_583, 0.4_333, 0.6_171, 0.5_684, 0.4_988])

					assert np.abs(image_slice.flatten() - expected_slice).max() < 2.0
					assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1E-2







@slow
@require_torch_gpu
class 		UpperCamelCase      (   unittest.TestCase						):



				def  _lowercase					(self    :				Any)   ->     Union[str, Any]:
					# clean up the VRAM after each test
					super().tearDown()
					gc.collect()
					torch.cuda.empty_cache()






				def  _lowercase					(self    :				Tuple)   ->     Optional[int]:
					__snake_case      :							List[Any]						=			load_numpy(
					    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
					    '/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy')

					__snake_case      :							Union[str, Any]						=			VQDiffusionPipeline.from_pretrained('microsoft/vq-diffusion-ithq')
					__snake_case      :							Tuple						=			pipeline.to(_A)
					pipeline.set_progress_bar_config(disable=_A)

					# requires GPU generator for gumbel softmax
					# don't use GPU generator in tests though
					__snake_case      :							Optional[int]						=			torch.Generator(device=_A).manual_seed(0)
					__snake_case      :							int						=			pipeline(
					    'teddy bear playing in the pool'     ,  num_images_per_prompt=1     ,  generator=_A     ,  output_type='np'     ,  )

					__snake_case      :							int						=			output.images[0]

					assert image.shape == (2_56, 2_56, 3)
					assert np.abs(expected_image - image).max() < 2.0





