






'''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,
)


__UpperCamelCase   :							Dict         =   logging.getLogger(__name__)
class 				__SCREAMING_SNAKE_CASE							(				UpperCAmelCase_  ):




      def 			__UpperCamelCase				(	self			, lowerCamelCase			, lowerCamelCase			, lowerCamelCase=None			, lowerCamelCase=None   )				->Union[str, Any]:



         '''simple docstring'''
         __a					  =     self.layer[current_layer](SCREAMING_SNAKE_CASE_			, SCREAMING_SNAKE_CASE_			, head_mask[current_layer]   )

         __a					  =     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."   ,   UpperCAmelCase_   ,   )
class 				__SCREAMING_SNAKE_CASE							(				UpperCAmelCase_  ):




      def __init__(	self			, lowerCamelCase   )				->Dict:



         '''simple docstring'''
         super().__init__(SCREAMING_SNAKE_CASE_   )

         __a					  =     BertEncoderWithPabee(SCREAMING_SNAKE_CASE_   )

         self.init_weights()
         __a					  =     0
         __a					  =     0
         __a					  =     0

         __a					  =     0




      def 			__UpperCamelCase				(	self			, lowerCamelCase   )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     threshold




      def 			__UpperCamelCase				(	self			, lowerCamelCase   )				->Optional[Any]:



         '''simple docstring'''
         __a					  =     patience




      def 			__UpperCamelCase				(	self   )				->Dict:



         '''simple docstring'''
         __a					  =     0
         __a					  =     0




      def 			__UpperCamelCase				(	self   )				->str:



         '''simple docstring'''
         __a					  =     self.inference_layers_num / self.inference_instances_num
         __a					  =     (
             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(SCREAMING_SNAKE_CASE_   )






      @add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_   )
      def 			__UpperCamelCase				(	self			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=False			, )				->Dict:



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

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

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

         # 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.
         __a					  =     self.get_extended_attention_mask(SCREAMING_SNAKE_CASE_			, SCREAMING_SNAKE_CASE_			, SCREAMING_SNAKE_CASE_   )

         # 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:
            __a					  =     encoder_hidden_states.size()
            __a					  =     (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
               __a					  =     torch.ones(SCREAMING_SNAKE_CASE_			, device=SCREAMING_SNAKE_CASE_   )
            __a					  =     self.invert_attention_mask(SCREAMING_SNAKE_CASE_   )
         else:
            __a					  =     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]
         __a					  =     self.get_head_mask(SCREAMING_SNAKE_CASE_			, self.config.num_hidden_layers   )

         __a					  =     self.embeddings(
             input_ids=SCREAMING_SNAKE_CASE_			, position_ids=SCREAMING_SNAKE_CASE_			, token_type_ids=SCREAMING_SNAKE_CASE_			, inputs_embeds=SCREAMING_SNAKE_CASE_   )
         __a					  =     embedding_output

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

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

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

               __a					  =     logits
               if patient_counter == self.patience:
                  break
            __a					  =     [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\n    the pooled output) e.g. for GLUE tasks. "   ,   UpperCAmelCase_   ,   )
class 				__SCREAMING_SNAKE_CASE							(				UpperCAmelCase_  ):




      def __init__(	self			, lowerCamelCase   )				->int:



         '''simple docstring'''
         super().__init__(SCREAMING_SNAKE_CASE_   )
         __a					  =     config.num_labels

         __a					  =     BertModelWithPabee(SCREAMING_SNAKE_CASE_   )
         __a					  =     nn.Dropout(config.hidden_dropout_prob   )
         __a					  =     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(SCREAMING_SNAKE_CASE_   )
      def 			__UpperCamelCase				(	self			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, lowerCamelCase=None			, )				->Tuple:



         '''simple docstring'''
         __a					  =     self.bert(
             input_ids=SCREAMING_SNAKE_CASE_			, attention_mask=SCREAMING_SNAKE_CASE_			, token_type_ids=SCREAMING_SNAKE_CASE_			, position_ids=SCREAMING_SNAKE_CASE_			, head_mask=SCREAMING_SNAKE_CASE_			, inputs_embeds=SCREAMING_SNAKE_CASE_			, output_dropout=self.dropout			, output_layers=self.classifiers			, regression=self.num_labels == 1			, )

         __a					  =     (logits[-1],)

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

         return outputs

from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
     from PIL import Image

     from ..image_utils import load_image

if is_torch_available():
     import torch

     from transformers.modeling_outputs import BaseModelOutput

     from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

lowerCamelCase__     :		Optional[Any]        = logging.get_logger(__name__)



@add_end_docstrings(UpperCAmelCase_				)
class    _snake_case (						UpperCAmelCase_				):


 def __init__(       self       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   super().__init__(**SCREAMING_SNAKE_CASE_)

   if self.framework == "tf":
     raise ValueError(f'The {self.__class__} is only available in PyTorch.')

   requires_backends(self       ,	"""vision""")
   self.check_model_type(SCREAMING_SNAKE_CASE_)


 def __call__(       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_ = None       ,	**SCREAMING_SNAKE_CASE_       ,	):


   '''simple docstring'''




   if "text_queries" in kwargs:
     lowercase__	: Any     		=   kwargs.pop("""text_queries""")

   if isinstance(SCREAMING_SNAKE_CASE_       ,	(str, Image.Image)):
     lowercase__	: Optional[Any]     		=   {"""image""": image, """candidate_labels""": candidate_labels}
   else:
     lowercase__	: int     		=   image
   lowercase__	: List[str]     		=   super().__call__(SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_)
   return results


 def   lowercase__  (       self       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: Tuple     		=   {}
   if "threshold" in kwargs:
     lowercase__	: List[Any]     		=   kwargs["""threshold"""]
   if "top_k" in kwargs:
     lowercase__	: int     		=   kwargs["""top_k"""]
   return {}, {}, postprocess_params


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: str     		=   load_image(inputs["""image"""])
   lowercase__	: Any     		=   inputs["""candidate_labels"""]
   if isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):
     lowercase__	: List[str]     		=   candidate_labels.split(""",""")

   lowercase__	: Tuple     		=   torch.tensor([[image.height, image.width]]       ,	dtype=torch.intaa)
   for i, candidate_label in enumerate(SCREAMING_SNAKE_CASE_):
     lowercase__	: Optional[Any]     		=   self.tokenizer(SCREAMING_SNAKE_CASE_       ,	return_tensors=self.framework)
     lowercase__	: Union[str, Any]     		=   self.image_processor(SCREAMING_SNAKE_CASE_       ,	return_tensors=self.framework)
     yield {
         "is_last": i == len(SCREAMING_SNAKE_CASE_) - 1,
         "target_size": target_size,
         "candidate_label": candidate_label,
         **text_inputs,
         **image_features,
     }


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: str     		=   model_inputs.pop("""target_size""")
   lowercase__	: Optional[int]     		=   model_inputs.pop("""candidate_label""")
   lowercase__	: Dict     		=   model_inputs.pop("""is_last""")

   lowercase__	: Union[str, Any]     		=   self.model(**SCREAMING_SNAKE_CASE_)

   lowercase__	: Union[str, Any]     		=   {"""target_size""": target_size, """candidate_label""": candidate_label, """is_last""": is_last, **outputs}
   return model_outputs


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_=0.1       ,	SCREAMING_SNAKE_CASE_=None):


   '''simple docstring'''




   lowercase__	: Union[str, Any]     		=   []
   for model_output in model_outputs:
     lowercase__	: Optional[int]     		=   model_output["""candidate_label"""]
     lowercase__	: Tuple     		=   BaseModelOutput(SCREAMING_SNAKE_CASE_)
     lowercase__	: List[str]     		=   self.image_processor.post_process_object_detection(
         outputs=SCREAMING_SNAKE_CASE_       ,	threshold=SCREAMING_SNAKE_CASE_       ,	target_sizes=model_output["""target_size"""])[0]

     for index in outputs["scores"].nonzero():
       lowercase__	: Optional[Any]     		=   outputs["""scores"""][index].item()
       lowercase__	: Optional[Any]     		=   self._get_bounding_box(outputs["""boxes"""][index][0])

       lowercase__	: Tuple     		=   {"""score""": score, """label""": label, """box""": box}
       results.append(SCREAMING_SNAKE_CASE_)

   lowercase__	: int     		=   sorted(SCREAMING_SNAKE_CASE_       ,	key=lambda SCREAMING_SNAKE_CASE_: x["score"]       ,	reverse=SCREAMING_SNAKE_CASE_)
   if top_k:
     lowercase__	: Any     		=   results[:top_k]

   return results



 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   if self.framework != "pt":
     raise ValueError("""The ZeroShotObjectDetectionPipeline is only available in PyTorch.""")
   lowercase__		,     lowercase__		,     lowercase__		,     lowercase__	: List[Any]     		=   box.int().tolist()
   lowercase__	: Optional[int]     		=   {
       """xmin""": xmin,
       """ymin""": ymin,
       """xmax""": xmax,
       """ymax""": ymax,
   }
   return bbox






