
"""simple docstring"""


from __future__ import absolute_import, division, print_function, unicode_literals

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

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy



@add_start_docstrings(
    '''The RoBERTa Model transformer with early exiting (DeeRoBERTa). '''     ,  _lowerCAmelCase     ,  )
class 	__a	(   _lowerCAmelCase						):
       UpperCamelCase_		:       List[Any]				=      RobertaConfig
       UpperCamelCase_		:       List[Any]				=      '''roberta'''






       def __init__(	self			:				Dict							,   UpperCAmelCase_			:				List[Any]      )->	List[Any]:







              """simple docstring"""
              super().__init__(UpperCAmelCase_      )

              UpperCamelCase       			=		RobertaEmbeddings(UpperCAmelCase_      )
              self.init_weights()





@add_start_docstrings(
    '''RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,
    also takes care of multi-layer training. '''     ,  _lowerCAmelCase     ,  )
class 	__a	(   _lowerCAmelCase						):
       UpperCamelCase_		:       List[Any]				=      RobertaConfig
       UpperCamelCase_		:       Any				=      '''roberta'''






       def __init__(	self			:				Union[str, Any]							,   UpperCAmelCase_			:				Tuple      )->	List[str]:







              """simple docstring"""
              super().__init__(UpperCAmelCase_      )
              UpperCamelCase       			=		config.num_labels
              UpperCamelCase       			=		config.num_hidden_layers

              UpperCamelCase       			=		DeeRobertaModel(UpperCAmelCase_      )
              UpperCamelCase       			=		nn.Dropout(config.hidden_dropout_prob      )
              UpperCamelCase       			=		nn.Linear(config.hidden_size							,   self.config.num_labels      )






       @add_start_docstrings_to_model_forward(UpperCAmelCase_      )
       def 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple							,   UpperCAmelCase_			:				List[Any]=None							,   UpperCAmelCase_			:				Any=None							,   UpperCAmelCase_			:				int=None							,   UpperCAmelCase_			:				List[Any]=None							,   UpperCAmelCase_			:				Optional[int]=None							,   UpperCAmelCase_			:				List[Any]=None							,   UpperCAmelCase_			:				str=None							,   UpperCAmelCase_			:				int=-1							,   UpperCAmelCase_			:				Union[str, Any]=False							,   )->	Optional[int]:







              """simple docstring"""
              UpperCamelCase       			=		self.num_layers
              try:
                     UpperCamelCase       			=		self.roberta(
                         UpperCAmelCase_							,   attention_mask=UpperCAmelCase_							,   token_type_ids=UpperCAmelCase_							,   position_ids=UpperCAmelCase_							,   head_mask=UpperCAmelCase_							,   inputs_embeds=UpperCAmelCase_							,   )

                     UpperCamelCase       			=		outputs[1]

                     UpperCamelCase       			=		self.dropout(UpperCAmelCase_      )
                     UpperCamelCase       			=		self.classifier(UpperCAmelCase_      )
                     UpperCamelCase       			=		(logits,) + outputs[2:]  # add hidden states and attention if they are here
              except HighwayException as e:
                     UpperCamelCase       			=		e.message
                     UpperCamelCase       			=		e.exit_layer
                     UpperCamelCase       			=		outputs[0]

              if not self.training:
                     UpperCamelCase       			=		entropy(UpperCAmelCase_      )
                     UpperCamelCase       			=		[]
                     UpperCamelCase       			=		[]
              if labels is not None:
                     if self.num_labels == 1:
                            #  We are doing regression
                            UpperCamelCase       			=		MSELoss()
                            UpperCamelCase       			=		loss_fct(logits.view(-1      )							,   labels.view(-1      )      )
                     else:
                            UpperCamelCase       			=		CrossEntropyLoss()
                            UpperCamelCase       			=		loss_fct(logits.view(-1							,   self.num_labels      )							,   labels.view(-1      )      )

                     # work with highway exits
                     UpperCamelCase       			=		[]
                     for highway_exit in outputs[-1]:
                            UpperCamelCase       			=		highway_exit[0]
                            if not self.training:
                                   highway_logits_all.append(UpperCAmelCase_      )
                                   highway_entropy.append(highway_exit[2]      )
                            if self.num_labels == 1:
                                   #  We are doing regression
                                   UpperCamelCase       			=		MSELoss()
                                   UpperCamelCase       			=		loss_fct(highway_logits.view(-1      )							,   labels.view(-1      )      )
                            else:
                                   UpperCamelCase       			=		CrossEntropyLoss()
                                   UpperCamelCase       			=		loss_fct(highway_logits.view(-1							,   self.num_labels      )							,   labels.view(-1      )      )
                            highway_losses.append(UpperCAmelCase_      )

                     if train_highway:
                            UpperCamelCase       			=		(sum(highway_losses[:-1]      ),) + outputs
                            # exclude the final highway, of course
                     else:
                            UpperCamelCase       			=		(loss,) + outputs
              if not self.training:
                     UpperCamelCase       			=		outputs + ((original_entropy, highway_entropy), exit_layer)
                     if output_layer >= 0:
                            UpperCamelCase       			=		(
                                (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
                            )  # use the highway of the last layer

              return outputs  # (loss), logits, (hidden_states), (attentions), entropy







"""simple docstring"""


import os
from typing import List, Optional, Union

from ...image_processing_utils import BatchFeature
from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType
from ..auto import AutoTokenizer



class 	__a	(   _lowerCAmelCase						):
       UpperCamelCase_		:       List[str]				=      ['''image_processor''', '''tokenizer''']
       UpperCamelCase_		:       Dict				=      '''BlipImageProcessor'''
       UpperCamelCase_		:       Optional[int]				=      '''AutoTokenizer'''






       def __init__(	self			:				List[str]							,   UpperCAmelCase_			:				Optional[int]							,   UpperCAmelCase_			:				List[str]							,   UpperCAmelCase_			:				Optional[int]      )->	Union[str, Any]:







              """simple docstring"""
              super().__init__(UpperCAmelCase_							,   UpperCAmelCase_      )

              # add QFormer tokenizer
              UpperCamelCase       			=		qformer_tokenizer






       def __call__(	self			:				Tuple							,   UpperCAmelCase_			:				ImageInput = None							,   UpperCAmelCase_			:				Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None							,   UpperCAmelCase_			:				bool = True							,   UpperCAmelCase_			:				Union[bool, str, PaddingStrategy] = False							,   UpperCAmelCase_			:				Union[bool, str, TruncationStrategy] = None							,   UpperCAmelCase_			:				Optional[int] = None							,   UpperCAmelCase_			:				int = 0							,   UpperCAmelCase_			:				Optional[int] = None							,   UpperCAmelCase_			:				Optional[bool] = None							,   UpperCAmelCase_			:				bool = False							,   UpperCAmelCase_			:				bool = False							,   UpperCAmelCase_			:				bool = False							,   UpperCAmelCase_			:				bool = False							,   UpperCAmelCase_			:				bool = False							,   UpperCAmelCase_			:				bool = True							,   UpperCAmelCase_			:				Optional[Union[str, TensorType]] = None							,   **UpperCAmelCase_			:				str							,   )->	BatchFeature:







              """simple docstring"""
              if images is None and text is None:
                     raise ValueError("You have to specify at least images or text."      )

              UpperCamelCase       			=		BatchFeature()

              if text is not None:
                     UpperCamelCase       			=		self.tokenizer(
                         text=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_							,   padding=UpperCAmelCase_							,   truncation=UpperCAmelCase_							,   max_length=UpperCAmelCase_							,   stride=UpperCAmelCase_							,   pad_to_multiple_of=UpperCAmelCase_							,   return_attention_mask=UpperCAmelCase_							,   return_overflowing_tokens=UpperCAmelCase_							,   return_special_tokens_mask=UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   return_token_type_ids=UpperCAmelCase_							,   return_length=UpperCAmelCase_							,   verbose=UpperCAmelCase_							,   return_tensors=UpperCAmelCase_							,   **UpperCAmelCase_							,   )
                     encoding.update(UpperCAmelCase_      )
                     UpperCamelCase       			=		self.qformer_tokenizer(
                         text=UpperCAmelCase_							,   add_special_tokens=UpperCAmelCase_							,   padding=UpperCAmelCase_							,   truncation=UpperCAmelCase_							,   max_length=UpperCAmelCase_							,   stride=UpperCAmelCase_							,   pad_to_multiple_of=UpperCAmelCase_							,   return_attention_mask=UpperCAmelCase_							,   return_overflowing_tokens=UpperCAmelCase_							,   return_special_tokens_mask=UpperCAmelCase_							,   return_offsets_mapping=UpperCAmelCase_							,   return_token_type_ids=UpperCAmelCase_							,   return_length=UpperCAmelCase_							,   verbose=UpperCAmelCase_							,   return_tensors=UpperCAmelCase_							,   **UpperCAmelCase_							,   )
                     UpperCamelCase       			=		qformer_text_encoding.pop("input_ids"      )
                     UpperCamelCase       			=		qformer_text_encoding.pop("attention_mask"      )

              if images is not None:
                     UpperCamelCase       			=		self.image_processor(UpperCAmelCase_							,   return_tensors=UpperCAmelCase_      )
                     encoding.update(UpperCAmelCase_      )

              return encoding






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple							,   *UpperCAmelCase_			:				Optional[int]							,   **UpperCAmelCase_			:				List[str]      )->	Dict:







              """simple docstring"""
              return self.tokenizer.batch_decode(*UpperCAmelCase_							,   **UpperCAmelCase_      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				int							,   *UpperCAmelCase_			:				Dict							,   **UpperCAmelCase_			:				Optional[int]      )->	List[Any]:







              """simple docstring"""
              return self.tokenizer.decode(*UpperCAmelCase_							,   **UpperCAmelCase_      )






       @property
       # Copied from transformers.models.blip.processing_blip.BlipProcessor.model_input_names
       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any      )->	Tuple:







              """simple docstring"""
              UpperCamelCase       			=		self.tokenizer.model_input_names
              UpperCamelCase       			=		self.image_processor.model_input_names
              return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Tuple							,   UpperCAmelCase_			:				Tuple							,   **UpperCAmelCase_			:				Tuple      )->	Union[str, Any]:







              """simple docstring"""
              if os.path.isfile(UpperCAmelCase_      ):
                     raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file"      )
              os.makedirs(UpperCAmelCase_							,   exist_ok=UpperCAmelCase_      )
              UpperCamelCase       			=		os.path.join(UpperCAmelCase_							,   "qformer_tokenizer"      )
              self.qformer_tokenizer.save_pretrained(UpperCAmelCase_      )
              return super().save_pretrained(UpperCAmelCase_							,   **UpperCAmelCase_      )






       @classmethod
       def 						_SCREAMING_SNAKE_CASE							(	cls			:				Dict							,   UpperCAmelCase_			:				int							,   **UpperCAmelCase_			:				int      )->	int:







              """simple docstring"""
              UpperCamelCase       			=		AutoTokenizer.from_pretrained(UpperCAmelCase_							,   subfolder="qformer_tokenizer"      )
              UpperCamelCase       			=		cls._get_arguments_from_pretrained(UpperCAmelCase_							,   **UpperCAmelCase_      )
              args.append(UpperCAmelCase_      )
              return cls(*UpperCAmelCase_      )





