


'''simple docstring'''




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.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


__snake_case    						=logging.getLogger(__name__)

class 		UpperCAmelCase_      (  __lowercase		):





     def 							__UpperCAmelCase      ( self       :      str    ,      UpperCAmelCase__       :      Optional[int]    ,      UpperCAmelCase__       :      Optional[Any]    ,      UpperCAmelCase__       :      List[str]=None    ,      UpperCAmelCase__       :      Optional[Any]=None    )				->  str:
        lowerCAmelCase     =			self.layer[current_layer](UpperCAmelCase__    ,      UpperCAmelCase__    ,      head_mask[current_layer]    )

        lowerCAmelCase     =			layer_outputs[0]

        return hidden_states




@add_start_docstrings(
    '''The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''       ,	__lowercase       ,	)
class 		UpperCAmelCase_      (  __lowercase		):





     def __init__( self       :      int    ,      UpperCAmelCase__       :      Tuple    )				->  Dict:
        super().__init__(UpperCAmelCase__    )

        lowerCAmelCase     =			BertEncoderWithPabee(UpperCAmelCase__    )

        self.init_weights()
        lowerCAmelCase     =			0
        lowerCAmelCase     =			0
        lowerCAmelCase     =			0

        lowerCAmelCase     =			0





     def 							__UpperCAmelCase      ( self       :      Tuple    ,      UpperCAmelCase__       :      Union[str, Any]    )				->  Union[str, Any]:
        lowerCAmelCase     =			threshold





     def 							__UpperCAmelCase      ( self       :      Union[str, Any]    ,      UpperCAmelCase__       :      Dict    )				->  Dict:
        lowerCAmelCase     =			patience





     def 							__UpperCAmelCase      ( self       :      Union[str, Any]    )				->  List[str]:
        lowerCAmelCase     =			0
        lowerCAmelCase     =			0





     def 							__UpperCAmelCase      ( self       :      int    )				->  str:
        lowerCAmelCase     =			self.inference_layers_num / self.inference_instances_num
        lowerCAmelCase     =			(
            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(UpperCAmelCase__    )





     @add_start_docstrings_to_model_forward(UpperCAmelCase__    )
     def 							__UpperCAmelCase      ( self       :      Tuple    ,      UpperCAmelCase__       :      List[Any]=None    ,      UpperCAmelCase__       :      Dict=None    ,      UpperCAmelCase__       :      Any=None    ,      UpperCAmelCase__       :      List[str]=None    ,      UpperCAmelCase__       :      Any=None    ,      UpperCAmelCase__       :      Any=None    ,      UpperCAmelCase__       :      List[Any]=None    ,      UpperCAmelCase__       :      Optional[int]=None    ,      UpperCAmelCase__       :      Any=None    ,      UpperCAmelCase__       :      Union[str, Any]=None    ,      UpperCAmelCase__       :      Dict=False    ,      )				->  List[Any]:

        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     =			input_ids.size()
        elif inputs_embeds is not None:
           lowerCAmelCase     =			inputs_embeds.size()[:-1]
        else:
           raise ValueError('You have to specify either input_ids or inputs_embeds'    )

        lowerCAmelCase     =			input_ids.device if input_ids is not None else inputs_embeds.device

        if attention_mask is None:
           lowerCAmelCase     =			torch.ones(UpperCAmelCase__    ,      device=UpperCAmelCase__    )
        if token_type_ids is None:
           lowerCAmelCase     =			torch.zeros(UpperCAmelCase__    ,      dtype=torch.long    ,      device=UpperCAmelCase__    )

        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        lowerCAmelCase     =			self.get_extended_attention_mask(UpperCAmelCase__    ,      UpperCAmelCase__    ,      UpperCAmelCase__    )

        # If a 2D ou 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.is_decoder and encoder_hidden_states is not None:
           lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     =			encoder_hidden_states.size()
           lowerCAmelCase     =			(encoder_batch_size, encoder_sequence_length)
           if encoder_attention_mask is None:
              lowerCAmelCase     =			torch.ones(UpperCAmelCase__    ,      device=UpperCAmelCase__    )
           lowerCAmelCase     =			self.invert_attention_mask(UpperCAmelCase__    )
        else:
           lowerCAmelCase     =			None

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        lowerCAmelCase     =			self.get_head_mask(UpperCAmelCase__    ,      self.config.num_hidden_layers    )

        lowerCAmelCase     =			self.embeddings(
            input_ids=UpperCAmelCase__    ,      position_ids=UpperCAmelCase__    ,      token_type_ids=UpperCAmelCase__    ,      inputs_embeds=UpperCAmelCase__    )
        lowerCAmelCase     =			embedding_output

        if self.training:
           lowerCAmelCase     =			[]
           for i in range(self.config.num_hidden_layers    ):
              lowerCAmelCase     =			self.encoder.adaptive_forward(
                  UpperCAmelCase__    ,      current_layer=UpperCAmelCase__    ,      attention_mask=UpperCAmelCase__    ,      head_mask=UpperCAmelCase__    )

              lowerCAmelCase     =			self.pooler(UpperCAmelCase__    )
              lowerCAmelCase     =			output_layers[i](output_dropout(UpperCAmelCase__    )    )
              res.append(UpperCAmelCase__    )
        elif self.patience == 0:  # Use all layers for inference
           lowerCAmelCase     =			self.encoder(
               UpperCAmelCase__    ,      attention_mask=UpperCAmelCase__    ,      head_mask=UpperCAmelCase__    ,      encoder_hidden_states=UpperCAmelCase__    ,      encoder_attention_mask=UpperCAmelCase__    ,      )
           lowerCAmelCase     =			self.pooler(encoder_outputs[0]    )
           lowerCAmelCase     =			[output_layers[self.config.num_hidden_layers - 1](UpperCAmelCase__    )]
        else:
           lowerCAmelCase     =			0
           lowerCAmelCase     =			None
           lowerCAmelCase     =			0
           for i in range(self.config.num_hidden_layers    ):
              calculated_layer_num += 1
              lowerCAmelCase     =			self.encoder.adaptive_forward(
                  UpperCAmelCase__    ,      current_layer=UpperCAmelCase__    ,      attention_mask=UpperCAmelCase__    ,      head_mask=UpperCAmelCase__    )

              lowerCAmelCase     =			self.pooler(UpperCAmelCase__    )
              lowerCAmelCase     =			output_layers[i](UpperCAmelCase__    )
              if regression:
                 lowerCAmelCase     =			logits.detach()
                 if patient_result is not None:
                    lowerCAmelCase     =			patient_result.detach()
                 if (patient_result is not None) and torch.abs(patient_result - labels    ) < self.regression_threshold:
                    patient_counter += 1
                 else:
                    lowerCAmelCase     =			0
              else:
                 lowerCAmelCase     =			logits.detach().argmax(dim=1    )
                 if patient_result is not None:
                    lowerCAmelCase     =			patient_result.detach().argmax(dim=1    )
                 if (patient_result is not None) and torch.all(labels.eq(UpperCAmelCase__    )    ):
                    patient_counter += 1
                 else:
                    lowerCAmelCase     =			0

              lowerCAmelCase     =			logits
              if patient_counter == self.patience:
                 break
           lowerCAmelCase     =			[patient_result]
           self.inference_layers_num += calculated_layer_num
           self.inference_instances_num += 1

        return res




@add_start_docstrings(
    '''Bert 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. '''       ,	__lowercase       ,	)
class 		UpperCAmelCase_      (  __lowercase		):





     def __init__( self       :      List[Any]    ,      UpperCAmelCase__       :      List[Any]    )				->  List[str]:
        super().__init__(UpperCAmelCase__    )
        lowerCAmelCase     =			config.num_labels

        lowerCAmelCase     =			BertModelWithPabee(UpperCAmelCase__    )
        lowerCAmelCase     =			nn.Dropout(config.hidden_dropout_prob    )
        lowerCAmelCase     =			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(UpperCAmelCase__    )
     def 							__UpperCAmelCase      ( self       :      List[str]    ,      UpperCAmelCase__       :      List[Any]=None    ,      UpperCAmelCase__       :      List[str]=None    ,      UpperCAmelCase__       :      Optional[int]=None    ,      UpperCAmelCase__       :      Tuple=None    ,      UpperCAmelCase__       :      List[str]=None    ,      UpperCAmelCase__       :      int=None    ,      UpperCAmelCase__       :      Any=None    ,      )				->  Any:
        lowerCAmelCase     =			self.bert(
            input_ids=UpperCAmelCase__    ,      attention_mask=UpperCAmelCase__    ,      token_type_ids=UpperCAmelCase__    ,      position_ids=UpperCAmelCase__    ,      head_mask=UpperCAmelCase__    ,      inputs_embeds=UpperCAmelCase__    ,      output_dropout=self.dropout    ,      output_layers=self.classifiers    ,      regression=self.num_labels == 1    ,      )

        lowerCAmelCase     =			(logits[-1],)

        if labels is not None:
           lowerCAmelCase     =			None
           lowerCAmelCase     =			0
           for ix, logits_item in enumerate(UpperCAmelCase__    ):
              if self.num_labels == 1:
                 #  We are doing regression
                 lowerCAmelCase     =			MSELoss()
                 lowerCAmelCase     =			loss_fct(logits_item.view(-1    )    ,      labels.view(-1    )    )
              else:
                 lowerCAmelCase     =			CrossEntropyLoss()
                 lowerCAmelCase     =			loss_fct(logits_item.view(-1    ,      self.num_labels    )    ,      labels.view(-1    )    )
              if total_loss is None:
                 lowerCAmelCase     =			loss
              else:
                 total_loss += loss * (ix + 1)
              total_weights += ix + 1
           lowerCAmelCase     =			(total_loss / total_weights,) + outputs

        return outputs









'''simple docstring'''




import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

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





     def 							__UpperCAmelCase      ( self       :      Optional[int]    )				->  int:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()





     def 							__UpperCAmelCase      ( self       :      Tuple    )				->  Any:
        lowerCAmelCase     =			StableDiffusionKDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'    )
        lowerCAmelCase     =			sd_pipe.to(UpperCAmelCase__    )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__    )

        sd_pipe.set_scheduler('sample_euler'    )

        lowerCAmelCase     =			'A painting of a squirrel eating a burger'
        lowerCAmelCase     =			torch.manual_seed(0    )
        lowerCAmelCase     =			sd_pipe([prompt]    ,      generator=UpperCAmelCase__    ,      guidance_scale=9.0    ,      num_inference_steps=2_0    ,      output_type='np'    )

        lowerCAmelCase     =			output.images

        lowerCAmelCase     =			image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase     =			np.array([0.0_447, 0.0_492, 0.0_468, 0.0_408, 0.0_383, 0.0_408, 0.0_354, 0.0_380, 0.0_339]    )

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





     def 							__UpperCAmelCase      ( self       :      List[str]    )				->  Dict:
        lowerCAmelCase     =			StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
        lowerCAmelCase     =			sd_pipe.to(UpperCAmelCase__    )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__    )

        sd_pipe.set_scheduler('sample_euler'    )

        lowerCAmelCase     =			'A painting of a squirrel eating a burger'
        lowerCAmelCase     =			torch.manual_seed(0    )
        lowerCAmelCase     =			sd_pipe([prompt]    ,      generator=UpperCAmelCase__    ,      guidance_scale=9.0    ,      num_inference_steps=2_0    ,      output_type='np'    )

        lowerCAmelCase     =			output.images

        lowerCAmelCase     =			image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase     =			np.array([0.1_237, 0.1_320, 0.1_438, 0.1_359, 0.1_390, 0.1_132, 0.1_277, 0.1_175, 0.1_112]    )

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





     def 							__UpperCAmelCase      ( self       :      Optional[Any]    )				->  List[str]:
        lowerCAmelCase     =			StableDiffusionKDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1-base'    )
        lowerCAmelCase     =			sd_pipe.to(UpperCAmelCase__    )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__    )

        sd_pipe.set_scheduler('sample_dpmpp_2m'    )

        lowerCAmelCase     =			'A painting of a squirrel eating a burger'
        lowerCAmelCase     =			torch.manual_seed(0    )
        lowerCAmelCase     =			sd_pipe(
            [prompt]    ,      generator=UpperCAmelCase__    ,      guidance_scale=7.5    ,      num_inference_steps=1_5    ,      output_type='np'    ,      use_karras_sigmas=UpperCAmelCase__    ,      )

        lowerCAmelCase     =			output.images

        lowerCAmelCase     =			image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_1_2, 5_1_2, 3)
        lowerCAmelCase     =			np.array(
            [0.11_381_689, 0.12_112_921, 0.1_389_457, 0.12_549_606, 0.1_244_964, 0.10_831_517, 0.11_562_866, 0.10_867_816, 0.10_499_048]    )

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





