
"""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_		:					Tuple   =			'''roberta'''






       def __init__(  self :							Union[str, Any]       ,      UpperCAmelCase_ :							Tuple			)->  Union[str, 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_		:					Optional[Any]   =			RobertaConfig
       UpperCamelCase_		:					Any   =			'''roberta'''






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







              """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 :							int       ,      UpperCAmelCase_ :							Any=None       ,      UpperCAmelCase_ :							Optional[Any]=None       ,      UpperCAmelCase_ :							int=None       ,      UpperCAmelCase_ :							int=None       ,      UpperCAmelCase_ :							Union[str, Any]=None       ,      UpperCAmelCase_ :							Any=None       ,      UpperCAmelCase_ :							Dict=None       ,      UpperCAmelCase_ :							Optional[int]=-1       ,      UpperCAmelCase_ :							Tuple=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"""


def 	lowerCamelCase__			(							UpperCAmelCase_				)->							int:







       """simple docstring"""
       if a < 0:
              raise ValueError("Input value must be a positive integer"				)
       elif isinstance(UpperCAmelCase_		,      UpperCAmelCase_				):
              raise TypeError("Input value must be a 'int' type"				)
       return bin(UpperCAmelCase_				).count("1"				)


if __name__ == "__main__":
      import doctest

      doctest.testmod()





