


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 							lowercase__(						snake_case__			):






      '''simple docstring'''







      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None)					->						Any:







         """simple docstring"""







         if current_layer == 0:
            UpperCamelCase__ :			Dict 						=self.embedding_hidden_mapping_in(A_)
         else:
            UpperCamelCase__ :			Optional[int] 						=hidden_states[0]

         UpperCamelCase__ :			Tuple 						=int(self.config.num_hidden_layers / self.config.num_hidden_groups)

         # Index of the hidden group
         UpperCamelCase__ :			List[Any] 						=int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups))

         UpperCamelCase__ :			Tuple 						=self.albert_layer_groups[group_idx](
             A_   ,			A_   ,			head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group]   ,			)
         UpperCamelCase__ :			Optional[int] 						=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 							lowercase__(						snake_case__			):






      '''simple docstring'''







      def __init__(		self   ,			__SCREAMING_SNAKE_CASE)					->						Dict:







         """simple docstring"""







         super().__init__(A_)

         UpperCamelCase__ :			List[Any] 						=AlbertTransformerWithPabee(A_)

         self.init_weights()
         UpperCamelCase__ :			Tuple 						=0
         UpperCamelCase__ :			Any 						=0
         UpperCamelCase__ :			Union[str, Any] 						=0

         UpperCamelCase__ :			Union[str, Any] 						=0




      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE)					->						List[str]:







         """simple docstring"""







         UpperCamelCase__ :			List[Any] 						=threshold




      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE)					->						List[str]:







         """simple docstring"""







         UpperCamelCase__ :			Union[str, Any] 						=patience




      def    UpperCAmelCase							(		self)					->						Optional[int]:







         """simple docstring"""







         UpperCamelCase__ :			str 						=0
         UpperCamelCase__ :			List[str] 						=0




      def    UpperCAmelCase							(		self)					->						Optional[int]:







         """simple docstring"""







         UpperCamelCase__ :			Tuple 						=self.inference_layers_num / self.inference_instances_num
         UpperCamelCase__ :			Optional[int] 						=(
             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(A_)




      @add_start_docstrings_to_model_forward(A_)
      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=False   ,			)					->						str:







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

         UpperCamelCase__ :			Optional[int] 						=input_ids.device if input_ids is not None else inputs_embeds.device

         if attention_mask is None:
            UpperCamelCase__ :			Any 						=torch.ones(A_   ,			device=A_)
         if token_type_ids is None:
            UpperCamelCase__ :			int 						=torch.zeros(A_   ,			dtype=torch.long   ,			device=A_)

         UpperCamelCase__ :			Union[str, Any] 						=attention_mask.unsqueeze(1).unsqueeze(2)
         UpperCamelCase__ :			Any 						=extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
         UpperCamelCase__ :			str 						=(1.0 - extended_attention_mask) * -1_00_00.0
         UpperCamelCase__ :			str 						=self.get_head_mask(A_   ,			self.config.num_hidden_layers)

         UpperCamelCase__ :			str 						=self.embeddings(
             A_   ,			position_ids=A_   ,			token_type_ids=A_   ,			inputs_embeds=A_)
         UpperCamelCase__ :			Optional[int] 						=embedding_output

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

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

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

               UpperCamelCase__ :			int 						=logits
               if patient_counter == self.patience:
                  break
            UpperCamelCase__ :			List[Any] 						=[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 							lowercase__(						snake_case__			):






      '''simple docstring'''







      def __init__(		self   ,			__SCREAMING_SNAKE_CASE)					->						Dict:







         """simple docstring"""







         super().__init__(A_)
         UpperCamelCase__ :			str 						=config.num_labels

         UpperCamelCase__ :			str 						=AlbertModelWithPabee(A_)
         UpperCamelCase__ :			Optional[Any] 						=nn.Dropout(config.classifier_dropout_prob)
         UpperCamelCase__ :			int 						=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(A_)
      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			__SCREAMING_SNAKE_CASE=None   ,			)					->						Tuple:







         """simple docstring"""







         UpperCamelCase__ :			Any 						=self.albert(
             input_ids=A_   ,			attention_mask=A_   ,			token_type_ids=A_   ,			position_ids=A_   ,			head_mask=A_   ,			inputs_embeds=A_   ,			output_dropout=self.dropout   ,			output_layers=self.classifiers   ,			regression=self.num_labels == 1   ,			)

         UpperCamelCase__ :			List[Any] 						=(logits[-1],)

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

         return outputs




import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
  import torch

if is_vision_available():
  from PIL import Image

  from transformers import CLIPImageProcessor





class A__     (			unittest.TestCase			):



 """simple docstring"""







 def __init__(		self			:						Union[str, Any]  ,    A_			:						Optional[Any]  ,    A_			:						List[str]=7  ,    A_			:						Optional[Any]=3  ,    A_			:						int=1_8  ,    A_			:						List[Any]=3_0  ,    A_			:						str=4_0_0  ,    A_			:						Union[str, Any]=True  ,    A_			:						Union[str, Any]=None  ,    A_			:						Union[str, Any]=True  ,    A_			:						Optional[int]=None  ,    A_			:						List[str]=True  ,    A_			:						Union[str, Any]=[0.48145466, 0.4578275, 0.40821073]  ,    A_			:						int=[0.26862954, 0.26130258, 0.27577711]  ,    A_			:						Optional[int]=True  ,    ):






      '''simple docstring'''




      _lowerCAmelCase :       Optional[Any]	   =     size if size is not None else {"shortest_edge": 2_0}
      _lowerCAmelCase :       Union[str, Any]	   =     crop_size if crop_size is not None else {"height": 1_8, "width": 1_8}
      _lowerCAmelCase :       Tuple	   =     parent
      _lowerCAmelCase :       int	   =     batch_size
      _lowerCAmelCase :       str	   =     num_channels
      _lowerCAmelCase :       int	   =     image_size
      _lowerCAmelCase :       Optional[int]	   =     min_resolution
      _lowerCAmelCase :       Any	   =     max_resolution
      _lowerCAmelCase :       str	   =     do_resize
      _lowerCAmelCase :       Union[str, Any]	   =     size
      _lowerCAmelCase :       Optional[int]	   =     do_center_crop
      _lowerCAmelCase :       Union[str, Any]	   =     crop_size
      _lowerCAmelCase :       Optional[int]	   =     do_normalize
      _lowerCAmelCase :       Optional[int]	   =     image_mean
      _lowerCAmelCase :       int	   =     image_std
      _lowerCAmelCase :       Optional[Any]	   =     do_convert_rgb

 def  __magic_name__							(		self			:						Any							):






      '''simple docstring'''




      return {
          "do_resize": self.do_resize,
          "size": self.size,
          "do_center_crop": self.do_center_crop,
          "crop_size": self.crop_size,
          "do_normalize": self.do_normalize,
          "image_mean": self.image_mean,
          "image_std": self.image_std,
          "do_convert_rgb": self.do_convert_rgb,
      }


 def  __magic_name__							(		self			:						List[Any]  ,    A_			:						List[Any]=False  ,    A_			:						Optional[Any]=False  ,    A_			:						Tuple=False							):






      '''simple docstring'''




      assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

      if equal_resolution:
           _lowerCAmelCase :       Tuple	   =     []
           for i in range(self.batch_size							):
                image_inputs.append(
                    np.random.randint(
                        2_5_5  ,    size=(self.num_channels, self.max_resolution, self.max_resolution)  ,    dtype=np.uinta							)							)
      else:
           _lowerCAmelCase :       Optional[int]	   =     []
           for i in range(self.batch_size							):
                _lowerCAmelCase   ,      _lowerCAmelCase :       Union[str, Any]	   =     np.random.choice(np.arange(self.min_resolution  ,    self.max_resolution							)  ,    2							)
                image_inputs.append(np.random.randint(2_5_5  ,    size=(self.num_channels, width, height)  ,    dtype=np.uinta							)							)

      if not numpify and not torchify:
           # PIL expects the channel dimension as last dimension
           _lowerCAmelCase :       List[str]	   =     [Image.fromarray(np.moveaxis(A_  ,    0  ,    -1							)							) for x in image_inputs]

      if torchify:
           _lowerCAmelCase :       Tuple	   =     [torch.from_numpy(A_							) for x in image_inputs]

      return image_inputs





@require_torch
@require_vision
class A__     (			A	,		unittest.TestCase			):



 """simple docstring"""







 _lowercase  :       Any  			=    CLIPImageProcessor if is_vision_available() else None

 def  __magic_name__							(		self			:						Union[str, Any]							):






      '''simple docstring'''




      _lowerCAmelCase :       Tuple	   =     CLIPImageProcessingTester(self							)

 @property
 def  __magic_name__							(		self			:						List[Any]							):






      '''simple docstring'''




      return self.image_processor_tester.prepare_image_processor_dict()

 def  __magic_name__							(		self			:						Union[str, Any]							):






      '''simple docstring'''




      _lowerCAmelCase :       Optional[int]	   =     self.image_processing_class(**self.image_processor_dict							)
      self.assertTrue(hasattr(A_  ,    "do_resize"							)							)
      self.assertTrue(hasattr(A_  ,    "size"							)							)
      self.assertTrue(hasattr(A_  ,    "do_center_crop"							)							)
      self.assertTrue(hasattr(A_  ,    "center_crop"							)							)
      self.assertTrue(hasattr(A_  ,    "do_normalize"							)							)
      self.assertTrue(hasattr(A_  ,    "image_mean"							)							)
      self.assertTrue(hasattr(A_  ,    "image_std"							)							)
      self.assertTrue(hasattr(A_  ,    "do_convert_rgb"							)							)

 def  __magic_name__							(		self			:						Optional[int]							):






      '''simple docstring'''




      _lowerCAmelCase :       Union[str, Any]	   =     self.image_processing_class.from_dict(self.image_processor_dict							)
      self.assertEqual(image_processor.size  ,    {"shortest_edge": 2_0}							)
      self.assertEqual(image_processor.crop_size  ,    {"height": 1_8, "width": 1_8}							)

      _lowerCAmelCase :       List[Any]	   =     self.image_processing_class.from_dict(self.image_processor_dict  ,    size=4_2  ,    crop_size=8_4							)
      self.assertEqual(image_processor.size  ,    {"shortest_edge": 4_2}							)
      self.assertEqual(image_processor.crop_size  ,    {"height": 8_4, "width": 8_4}							)

 def  __magic_name__							(		self			:						Dict							):






      '''simple docstring'''




      pass

 def  __magic_name__							(		self			:						List[Any]							):






      '''simple docstring'''




      _lowerCAmelCase :       Optional[Any]	   =     self.image_processing_class(**self.image_processor_dict							)
      # create random PIL images
      _lowerCAmelCase :       Union[str, Any]	   =     self.image_processor_tester.prepare_inputs(equal_resolution=A_							)
      for image in image_inputs:
           self.assertIsInstance(A_  ,    Image.Image							)

      # Test not batched input
      _lowerCAmelCase :       Optional[Any]	   =     image_processing(image_inputs[0]  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              1,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )

      # Test batched
      _lowerCAmelCase :       Optional[Any]	   =     image_processing(A_  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )

 def  __magic_name__							(		self			:						List[Any]							):






      '''simple docstring'''




      _lowerCAmelCase :       List[Any]	   =     self.image_processing_class(**self.image_processor_dict							)
      # create random numpy tensors
      _lowerCAmelCase :       Dict	   =     self.image_processor_tester.prepare_inputs(equal_resolution=A_  ,    numpify=A_							)
      for image in image_inputs:
           self.assertIsInstance(A_  ,    np.ndarray							)

      # Test not batched input
      _lowerCAmelCase :       Optional[Any]	   =     image_processing(image_inputs[0]  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              1,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )

      # Test batched
      _lowerCAmelCase :       Union[str, Any]	   =     image_processing(A_  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )


 def  __magic_name__							(		self			:						Tuple							):






      '''simple docstring'''




      _lowerCAmelCase :       Any	   =     self.image_processing_class(**self.image_processor_dict							)
      # create random PyTorch tensors
      _lowerCAmelCase :       Any	   =     self.image_processor_tester.prepare_inputs(equal_resolution=A_  ,    torchify=A_							)
      for image in image_inputs:
           self.assertIsInstance(A_  ,    torch.Tensor							)

      # Test not batched input
      _lowerCAmelCase :       Tuple	   =     image_processing(image_inputs[0]  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              1,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )

      # Test batched
      _lowerCAmelCase :       List[Any]	   =     image_processing(A_  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              self.image_processor_tester.batch_size,
              self.image_processor_tester.num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )





@require_torch
@require_vision
class A__     (			A	,		unittest.TestCase			):



 """simple docstring"""







 _lowercase  :       List[str]  			=    CLIPImageProcessor if is_vision_available() else None

 def  __magic_name__							(		self			:						Tuple							):






      '''simple docstring'''




      _lowerCAmelCase :       int	   =     CLIPImageProcessingTester(self  ,    num_channels=4							)
      _lowerCAmelCase :       Tuple	   =     3

 @property
 def  __magic_name__							(		self			:						Optional[int]							):






      '''simple docstring'''




      return self.image_processor_tester.prepare_image_processor_dict()

 def  __magic_name__							(		self			:						Optional[int]							):






      '''simple docstring'''




      _lowerCAmelCase :       List[str]	   =     self.image_processing_class(**self.image_processor_dict							)
      self.assertTrue(hasattr(A_  ,    "do_resize"							)							)
      self.assertTrue(hasattr(A_  ,    "size"							)							)
      self.assertTrue(hasattr(A_  ,    "do_center_crop"							)							)
      self.assertTrue(hasattr(A_  ,    "center_crop"							)							)
      self.assertTrue(hasattr(A_  ,    "do_normalize"							)							)
      self.assertTrue(hasattr(A_  ,    "image_mean"							)							)
      self.assertTrue(hasattr(A_  ,    "image_std"							)							)
      self.assertTrue(hasattr(A_  ,    "do_convert_rgb"							)							)

 def  __magic_name__							(		self			:						Optional[Any]							):






      '''simple docstring'''




      pass


 def  __magic_name__							(		self			:						Any							):






      '''simple docstring'''




      _lowerCAmelCase :       Tuple	   =     self.image_processing_class(**self.image_processor_dict							)
      # create random PIL images
      _lowerCAmelCase :       Optional[int]	   =     self.image_processor_tester.prepare_inputs(equal_resolution=A_							)
      for image in image_inputs:
           self.assertIsInstance(A_  ,    Image.Image							)

      # Test not batched input
      _lowerCAmelCase :       Optional[Any]	   =     image_processing(image_inputs[0]  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              1,
              self.expected_encoded_image_num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )

      # Test batched
      _lowerCAmelCase :       int	   =     image_processing(A_  ,    return_tensors="pt"							).pixel_values
      self.assertEqual(
          encoded_images.shape  ,    (
              self.image_processor_tester.batch_size,
              self.expected_encoded_image_num_channels,
              self.image_processor_tester.crop_size["height"],
              self.image_processor_tester.crop_size["width"],
          )  ,    )






