




import logging

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.albert.modeling_albert import (
    ALBERT_INPUTS_DOCSTRING,
    ALBERT_START_DOCSTRING,
    AlbertModel,
    AlbertPreTrainedModel,
    AlbertTransformer,
)


UpperCamelCase__		=						logging.getLogger(__name__)


class 					__SCREAMING_SNAKE_CASE							(		snake_case_  ):

					def      lowercase_							(    self			,	__lowercase			,	__lowercase			,	__lowercase=None			,	__lowercase=None ):







						'''simple docstring'''




						if current_layer == 0:
							lowerCamelCase_							:		str							       =							self.embedding_hidden_mapping_in(__lowercase )
						else:
							lowerCamelCase_							:		List[Any]							       =							hidden_states[0]

						lowerCamelCase_							:		Optional[Any]							       =							int(self.config.num_hidden_layers / self.config.num_hidden_groups )

						# Index of the hidden group
						lowerCamelCase_							:		Any							       =							int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups) )

						lowerCamelCase_							:		Union[str, Any]							       =							self.albert_layer_groups[group_idx](
						    __lowercase			,	__lowercase			,	head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]			,	)
						lowerCamelCase_							:		Tuple							       =							layer_group_output[0]

						return (hidden_states,)


@add_start_docstrings(
    """The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top."""       ,					snake_case_       ,					)
class 					__SCREAMING_SNAKE_CASE							(		snake_case_  ):

					def __init__(    self			,	__lowercase ):







						'''simple docstring'''




						super().__init__(__lowercase )

						lowerCamelCase_							:		Optional[int]							       =							AlbertTransformerWithPabee(__lowercase )

						self.init_weights()
						lowerCamelCase_							:		List[str]							       =							0
						lowerCamelCase_							:		List[str]							       =							0
						lowerCamelCase_							:		Any							       =							0

						lowerCamelCase_							:		Union[str, Any]							       =							0

					def      lowercase_							(    self			,	__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							       =							threshold

					def      lowercase_							(    self			,	__lowercase ):







						'''simple docstring'''




						lowerCamelCase_							:		int							       =							patience

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							       =							0
						lowerCamelCase_							:		Tuple							       =							0

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		int							       =							self.inference_layers_num / self.inference_instances_num
						lowerCamelCase_							:		Any							       =							(
						    f"*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ="
						    f" {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***"
						)
						print(__lowercase )





					@add_start_docstrings_to_model_forward(__lowercase )
					def      lowercase_							(    self			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=False			,	):







						'''simple docstring'''




						if input_ids is not None and inputs_embeds is not None:
							raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time' )
						elif input_ids is not None:
							lowerCamelCase_							:		Dict							       =							input_ids.size()
						elif inputs_embeds is not None:
							lowerCamelCase_							:		Optional[int]							       =							inputs_embeds.size()[:-1]
						else:
							raise ValueError('You have to specify either input_ids or inputs_embeds' )

						lowerCamelCase_							:		Tuple							       =							input_ids.device if input_ids is not None else inputs_embeds.device

						if attention_mask is None:
							lowerCamelCase_							:		Dict							       =							torch.ones(__lowercase			,	device=__lowercase )
						if token_type_ids is None:
							lowerCamelCase_							:		List[Any]							       =							torch.zeros(__lowercase			,	dtype=torch.long			,	device=__lowercase )

						lowerCamelCase_							:		List[Any]							       =							attention_mask.unsqueeze(1 ).unsqueeze(2 )
						lowerCamelCase_							:		Tuple							       =							extended_attention_mask.to(dtype=self.dtype )  # fp16 compatibility
						lowerCamelCase_							:		Union[str, Any]							       =							(1.0 - extended_attention_mask) * -10_000.0
						lowerCamelCase_							:		Tuple							       =							self.get_head_mask(__lowercase			,	self.config.num_hidden_layers )

						lowerCamelCase_							:		Optional[int]							       =							self.embeddings(
						    __lowercase			,	position_ids=__lowercase			,	token_type_ids=__lowercase			,	inputs_embeds=__lowercase )
						lowerCamelCase_							:		Tuple							       =							embedding_output

						if self.training:
							lowerCamelCase_							:		Optional[int]							       =							[]
							for i in range(self.config.num_hidden_layers ):
								lowerCamelCase_							:		List[str]							       =							self.encoder.adaptive_forward(
								    __lowercase			,	current_layer=__lowercase			,	attention_mask=__lowercase			,	head_mask=__lowercase			,	)

								lowerCamelCase_							:		Optional[int]							       =							self.pooler_activation(self.pooler(encoder_outputs[0][:, 0] ) )
								lowerCamelCase_							:		Dict							       =							output_layers[i](output_dropout(__lowercase ) )
								res.append(__lowercase )
						elif self.patience == 0:  # Use all layers for inference
							lowerCamelCase_							:		int							       =							self.encoder(__lowercase			,	__lowercase			,	head_mask=__lowercase )
							lowerCamelCase_							:		Dict							       =							self.pooler_activation(self.pooler(encoder_outputs[0][:, 0] ) )
							lowerCamelCase_							:		Union[str, Any]							       =							[output_layers[self.config.num_hidden_layers - 1](__lowercase )]
						else:
							lowerCamelCase_							:		Optional[Any]							       =							0
							lowerCamelCase_							:		Dict							       =							None
							lowerCamelCase_							:		Optional[Any]							       =							0
							for i in range(self.config.num_hidden_layers ):
								calculated_layer_num += 1
								lowerCamelCase_							:		Optional[int]							       =							self.encoder.adaptive_forward(
								    __lowercase			,	current_layer=__lowercase			,	attention_mask=__lowercase			,	head_mask=__lowercase			,	)

								lowerCamelCase_							:		Any							       =							self.pooler_activation(self.pooler(encoder_outputs[0][:, 0] ) )
								lowerCamelCase_							:		Optional[int]							       =							output_layers[i](__lowercase )
								if regression:
									lowerCamelCase_							:		str							       =							logits.detach()
									if patient_result is not None:
										lowerCamelCase_							:		List[str]							       =							patient_result.detach()
									if (patient_result is not None) and torch.abs(patient_result - labels ) < self.regression_threshold:
										patient_counter += 1
									else:
										lowerCamelCase_							:		int							       =							0
								else:
									lowerCamelCase_							:		Dict							       =							logits.detach().argmax(dim=1 )
									if patient_result is not None:
										lowerCamelCase_							:		int							       =							patient_result.detach().argmax(dim=1 )
									if (patient_result is not None) and torch.all(labels.eq(__lowercase ) ):
										patient_counter += 1
									else:
										lowerCamelCase_							:		List[str]							       =							0

								lowerCamelCase_							:		Optional[Any]							       =							logits
								if patient_counter == self.patience:
									break
							lowerCamelCase_							:		Dict							       =							[patient_result]
							self.inference_layers_num += calculated_layer_num
							self.inference_instances_num += 1

						return res


@add_start_docstrings(
    """Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. """       ,					snake_case_       ,					)
class 					__SCREAMING_SNAKE_CASE							(		snake_case_  ):

					def __init__(    self			,	__lowercase ):







						'''simple docstring'''




						super().__init__(__lowercase )
						lowerCamelCase_							:		Any							       =							config.num_labels

						lowerCamelCase_							:		List[str]							       =							AlbertModelWithPabee(__lowercase )
						lowerCamelCase_							:		Optional[Any]							       =							nn.Dropout(config.classifier_dropout_prob )
						lowerCamelCase_							:		List[Any]							       =							nn.ModuleList(
						    [nn.Linear(config.hidden_size			,	self.config.num_labels ) for _ in range(config.num_hidden_layers )] )

						self.init_weights()





					@add_start_docstrings_to_model_forward(__lowercase )
					def      lowercase_							(    self			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	__lowercase=None			,	):







						'''simple docstring'''




						lowerCamelCase_							:		List[Any]							       =							self.albert(
						    input_ids=__lowercase			,	attention_mask=__lowercase			,	token_type_ids=__lowercase			,	position_ids=__lowercase			,	head_mask=__lowercase			,	inputs_embeds=__lowercase			,	output_dropout=self.dropout			,	output_layers=self.classifiers			,	regression=self.num_labels == 1			,	)

						lowerCamelCase_							:		List[str]							       =							(logits[-1],)

						if labels is not None:
							lowerCamelCase_							:		Optional[int]							       =							None
							lowerCamelCase_							:		Any							       =							0
							for ix, logits_item in enumerate(__lowercase ):
								if self.num_labels == 1:
									#  We are doing regression
									lowerCamelCase_							:		Any							       =							MSELoss()
									lowerCamelCase_							:		Optional[int]							       =							loss_fct(logits_item.view(-1 )			,	labels.view(-1 ) )
								else:
									lowerCamelCase_							:		List[Any]							       =							CrossEntropyLoss()
									lowerCamelCase_							:		Tuple							       =							loss_fct(logits_item.view(-1			,	self.num_labels )			,	labels.view(-1 ) )
								if total_loss is None:
									lowerCamelCase_							:		Any							       =							loss
								else:
									total_loss += loss * (ix + 1)
								total_weights += ix + 1
							lowerCamelCase_							:		str							       =							(total_loss / total_weights,) + outputs

						return outputs






import random
import unittest

import numpy as np
import torch
from PIL import Image
from torch import nn
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import KandinskyVaaPriorEmbaEmbPipeline, PriorTransformer, UnCLIPScheduler
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps

from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


class     __SCREAMING_SNAKE_CASE							(		lowerCamelCase       ,					unittest.TestCase  ):
					_lowerCamelCase		  =			KandinskyVaaPriorEmbaEmbPipeline
					_lowerCamelCase		  =			["""prompt""", """image"""]
					_lowerCamelCase		  =			["""prompt""", """image"""]
					_lowerCamelCase		  =			[
					    """num_images_per_prompt""",
					    """strength""",
					    """generator""",
					    """num_inference_steps""",
					    """latents""",
					    """negative_prompt""",
					    """guidance_scale""",
					    """output_type""",
					    """return_dict""",
					]
					_lowerCamelCase		  =			False

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						return 32

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						return 32

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						return self.time_input_dim

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						return self.time_input_dim * 4

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						return 100

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Dict							       =							CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )
						return tokenizer

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						torch.manual_seed(0 )
						lowerCamelCase_							:		Dict							       =							CLIPTextConfig(
						    bos_token_id=0			,	eos_token_id=2			,	hidden_size=self.text_embedder_hidden_size			,	projection_dim=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=1_000			,	)
						return CLIPTextModelWithProjection(__lowercase )

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						torch.manual_seed(0 )

						lowerCamelCase_							:		int							       =							{
						    'num_attention_heads': 2,
						    'attention_head_dim': 12,
						    'embedding_dim': self.text_embedder_hidden_size,
						    'num_layers': 1,
						}

						lowerCamelCase_							:		str							       =							PriorTransformer(**__lowercase )
						# clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
						lowerCamelCase_							:		Dict							       =							nn.Parameter(torch.ones(model.clip_std.shape ) )
						return model

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						torch.manual_seed(0 )
						lowerCamelCase_							:		str							       =							CLIPVisionConfig(
						    hidden_size=self.text_embedder_hidden_size			,	image_size=224			,	projection_dim=self.text_embedder_hidden_size			,	intermediate_size=37			,	num_attention_heads=4			,	num_channels=3			,	num_hidden_layers=5			,	patch_size=14			,	)

						lowerCamelCase_							:		List[str]							       =							CLIPVisionModelWithProjection(__lowercase )
						return model

					@property
					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		int							       =							CLIPImageProcessor(
						    crop_size=224			,	do_center_crop=__lowercase			,	do_normalize=__lowercase			,	do_resize=__lowercase			,	image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]			,	image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]			,	resample=3			,	size=224			,	)

						return image_processor

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		List[Any]							       =							self.dummy_prior
						lowerCamelCase_							:		str							       =							self.dummy_image_encoder
						lowerCamelCase_							:		Tuple							       =							self.dummy_text_encoder
						lowerCamelCase_							:		str							       =							self.dummy_tokenizer
						lowerCamelCase_							:		int							       =							self.dummy_image_processor

						lowerCamelCase_							:		Dict							       =							UnCLIPScheduler(
						    variance_type='fixed_small_log'			,	prediction_type='sample'			,	num_train_timesteps=1_000			,	clip_sample=__lowercase			,	clip_sample_range=10.0			,	)

						lowerCamelCase_							:		Dict							       =							{
						    'prior': prior,
						    'image_encoder': image_encoder,
						    'text_encoder': text_encoder,
						    'tokenizer': tokenizer,
						    'scheduler': scheduler,
						    'image_processor': image_processor,
						}

						return components

					def      lowercase_							(    self			,	__lowercase			,	__lowercase=0 ):







						'''simple docstring'''




						if str(__lowercase ).startswith('mps' ):
							lowerCamelCase_							:		Union[str, Any]							       =							torch.manual_seed(__lowercase )
						else:
							lowerCamelCase_							:		Optional[Any]							       =							torch.Generator(device=__lowercase ).manual_seed(__lowercase )

						lowerCamelCase_							:		Optional[int]							       =							floats_tensor((1, 3, 64, 64)			,	rng=random.Random(__lowercase ) ).to(__lowercase )
						lowerCamelCase_							:		Any							       =							image.cpu().permute(0			,	2			,	3			,	1 )[0]
						lowerCamelCase_							:		List[Any]							       =							Image.fromarray(np.uinta(__lowercase ) ).convert('RGB' ).resize((256, 256) )

						lowerCamelCase_							:		Tuple							       =							{
						    'prompt': 'horse',
						    'image': init_image,
						    'strength': 0.5,
						    'generator': generator,
						    'guidance_scale': 4.0,
						    'num_inference_steps': 2,
						    'output_type': 'np',
						}
						return inputs

					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		str							       =							'cpu'

						lowerCamelCase_							:		Union[str, Any]							       =							self.get_dummy_components()

						lowerCamelCase_							:		str							       =							self.pipeline_class(**__lowercase )
						lowerCamelCase_							:		Optional[Any]							       =							pipe.to(__lowercase )

						pipe.set_progress_bar_config(disable=__lowercase )

						lowerCamelCase_							:		List[str]							       =							pipe(**self.get_dummy_inputs(__lowercase ) )
						lowerCamelCase_							:		Union[str, Any]							       =							output.image_embeds

						lowerCamelCase_							:		List[str]							       =							pipe(
						    **self.get_dummy_inputs(__lowercase )			,	return_dict=__lowercase			,	)[0]

						lowerCamelCase_							:		List[Any]							       =							image[0, -10:]
						lowerCamelCase_							:		List[Any]							       =							image_from_tuple[0, -10:]

						assert image.shape == (1, 32)

						lowerCamelCase_							:		Any							       =							np.array(
						    [
						        0.1_07_12_84,
						        1.3_33_02_71,
						        0.61_26_02_23,
						        -0.6_69_10_65,
						        -0.3_84_68_52,
						        -1.0_30_36_61,
						        0.22_71_61_11,
						        0.03_34_89_01,
						        0.30_04_06_75,
						        -0.24_80_50_29,
						    ] )

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

					@skip_mps
					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		int							       =							torch_device == 'cpu'
						lowerCamelCase_							:		Union[str, Any]							       =							True
						lowerCamelCase_							:		str							       =							False

						self._test_inference_batch_single_identical(
						    test_max_difference=__lowercase			,	relax_max_difference=__lowercase			,	test_mean_pixel_difference=__lowercase			,	)





					@skip_mps
					def      lowercase_							(    self ):







						'''simple docstring'''




						lowerCamelCase_							:		Optional[Any]							       =							torch_device == 'cpu'
						lowerCamelCase_							:		Optional[Any]							       =							False

						self._test_attention_slicing_forward_pass(
						    test_max_difference=__lowercase			,	test_mean_pixel_difference=__lowercase			,	)
