





'''simple docstring'''






import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          DistilBertForMaskedLM,
          DistilBertForMultipleChoice,
          DistilBertForQuestionAnswering,
          DistilBertForSequenceClassification,
          DistilBertForTokenClassification,
          DistilBertModel,
      )







class 		UpperCAmelCase   (				UpperCAmelCase_):


  """simple docstring"""







  def __init__(				self    :					Any							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					List[str]=13							,	UpperCamelCase__    :					int=7							,	UpperCamelCase__    :					int=True							,	UpperCamelCase__    :					List[str]=True							,	UpperCamelCase__    :					str=False							,	UpperCamelCase__    :					str=True							,	UpperCamelCase__    :					Tuple=99							,	UpperCamelCase__    :					List[str]=32							,	UpperCamelCase__    :					Optional[int]=5							,	UpperCamelCase__    :					Union[str, Any]=4							,	UpperCamelCase__    :					int=37							,	UpperCamelCase__    :					Union[str, Any]="gelu"							,	UpperCamelCase__    :					Tuple=0.1							,	UpperCamelCase__    :					Optional[int]=0.1							,	UpperCamelCase__    :					Union[str, Any]=512							,	UpperCamelCase__    :					Optional[Any]=16							,	UpperCamelCase__    :					List[str]=2							,	UpperCamelCase__    :					int=0.02							,	UpperCamelCase__    :					str=3							,	UpperCamelCase__    :					Optional[int]=4							,	UpperCamelCase__    :					Optional[int]=None							,	)     ->						List[Any]:
    _UpperCamelCase       			=parent
    _UpperCamelCase       			=batch_size
    _UpperCamelCase       			=seq_length
    _UpperCamelCase       			=is_training
    _UpperCamelCase       			=use_input_mask
    _UpperCamelCase       			=use_token_type_ids
    _UpperCamelCase       			=use_labels
    _UpperCamelCase       			=vocab_size
    _UpperCamelCase       			=hidden_size
    _UpperCamelCase       			=num_hidden_layers
    _UpperCamelCase       			=num_attention_heads
    _UpperCamelCase       			=intermediate_size
    _UpperCamelCase       			=hidden_act
    _UpperCamelCase       			=hidden_dropout_prob
    _UpperCamelCase       			=attention_probs_dropout_prob
    _UpperCamelCase       			=max_position_embeddings
    _UpperCamelCase       			=type_vocab_size
    _UpperCamelCase       			=type_sequence_label_size
    _UpperCamelCase       			=initializer_range
    _UpperCamelCase       			=num_labels
    _UpperCamelCase       			=num_choices
    _UpperCamelCase       			=scope


  def   UpperCamelCase__      (				self    :					Dict  )     ->						List[Any]:
    _UpperCamelCase       			=ids_tensor([self.batch_size, self.seq_length]							,	self.vocab_size  )

    _UpperCamelCase       			=None
    if self.use_input_mask:
      _UpperCamelCase       			=random_attention_mask([self.batch_size, self.seq_length]  )

    _UpperCamelCase       			=None
    _UpperCamelCase       			=None
    _UpperCamelCase       			=None
    if self.use_labels:
      _UpperCamelCase       			=ids_tensor([self.batch_size]							,	self.type_sequence_label_size  )
      _UpperCamelCase       			=ids_tensor([self.batch_size, self.seq_length]							,	self.num_labels  )
      _UpperCamelCase       			=ids_tensor([self.batch_size]							,	self.num_choices  )

    _UpperCamelCase       			=self.get_config()

    return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


  def   UpperCamelCase__      (				self    :					Any  )     ->						Optional[int]:
    return DistilBertConfig(
        vocab_size=self.vocab_size							,	dim=self.hidden_size							,	n_layers=self.num_hidden_layers							,	n_heads=self.num_attention_heads							,	hidden_dim=self.intermediate_size							,	hidden_act=self.hidden_act							,	dropout=self.hidden_dropout_prob							,	attention_dropout=self.attention_probs_dropout_prob							,	max_position_embeddings=self.max_position_embeddings							,	initializer_range=self.initializer_range							,	)


  def   UpperCamelCase__      (				self    :					Dict							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Optional[int]  )     ->						Dict:
    _UpperCamelCase       			=DistilBertModel(config=_lowercase  )
    model.to(_lowercase  )
    model.eval()
    _UpperCamelCase       			=model(_lowercase							,	_lowercase  )
    _UpperCamelCase       			=model(_lowercase  )
    self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, self.seq_length, self.hidden_size)  )


  def   UpperCamelCase__      (				self    :					Tuple							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Optional[int]  )     ->						Dict:
    _UpperCamelCase       			=DistilBertForMaskedLM(config=_lowercase  )
    model.to(_lowercase  )
    model.eval()
    _UpperCamelCase       			=model(_lowercase							,	attention_mask=_lowercase							,	labels=_lowercase  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.vocab_size)  )


  def   UpperCamelCase__      (				self    :					Any							,	UpperCamelCase__    :					str							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Optional[Any]  )     ->						List[Any]:
    _UpperCamelCase       			=DistilBertForQuestionAnswering(config=_lowercase  )
    model.to(_lowercase  )
    model.eval()
    _UpperCamelCase       			=model(
        _lowercase							,	attention_mask=_lowercase							,	start_positions=_lowercase							,	end_positions=_lowercase  )
    self.parent.assertEqual(result.start_logits.shape							,	(self.batch_size, self.seq_length)  )
    self.parent.assertEqual(result.end_logits.shape							,	(self.batch_size, self.seq_length)  )


  def   UpperCamelCase__      (				self    :					Any							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					Optional[int]							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Dict  )     ->						Optional[int]:
    _UpperCamelCase       			=self.num_labels
    _UpperCamelCase       			=DistilBertForSequenceClassification(_lowercase  )
    model.to(_lowercase  )
    model.eval()
    _UpperCamelCase       			=model(_lowercase							,	attention_mask=_lowercase							,	labels=_lowercase  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.num_labels)  )


  def   UpperCamelCase__      (				self    :					Optional[int]							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Optional[Any]							,	UpperCamelCase__    :					int							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					List[str]  )     ->						int:
    _UpperCamelCase       			=self.num_labels
    _UpperCamelCase       			=DistilBertForTokenClassification(config=_lowercase  )
    model.to(_lowercase  )
    model.eval()

    _UpperCamelCase       			=model(_lowercase							,	attention_mask=_lowercase							,	labels=_lowercase  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.num_labels)  )


  def   UpperCamelCase__      (				self    :					List[Any]							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					Union[str, Any]							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Union[str, Any]  )     ->						Union[str, Any]:
    _UpperCamelCase       			=self.num_choices
    _UpperCamelCase       			=DistilBertForMultipleChoice(config=_lowercase  )
    model.to(_lowercase  )
    model.eval()
    _UpperCamelCase       			=input_ids.unsqueeze(1  ).expand(-1							,	self.num_choices							,	-1  ).contiguous()
    _UpperCamelCase       			=input_mask.unsqueeze(1  ).expand(-1							,	self.num_choices							,	-1  ).contiguous()
    _UpperCamelCase       			=model(
        _lowercase							,	attention_mask=_lowercase							,	labels=_lowercase							,	)
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.num_choices)  )




  def   UpperCamelCase__      (				self    :					int  )     ->						Dict:
    _UpperCamelCase       			=self.prepare_config_and_inputs()
    ((_UpperCamelCase)			,     (_UpperCamelCase)			,     (_UpperCamelCase)			,     (_UpperCamelCase)			,     (_UpperCamelCase)			,     (_UpperCamelCase))       			=config_and_inputs
    _UpperCamelCase       			={'''input_ids''': input_ids, '''attention_mask''': input_mask}
    return config, inputs_dict







@require_torch
class 		UpperCAmelCase   (				UpperCAmelCase_   ,     UpperCAmelCase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				(
      (
          DistilBertModel,
          DistilBertForMaskedLM,
          DistilBertForMultipleChoice,
          DistilBertForQuestionAnswering,
          DistilBertForSequenceClassification,
          DistilBertForTokenClassification,
      )
      if is_torch_available()
      else None
  )
  lowerCAmelCase_					=				(
      {
          """feature-extraction""": DistilBertModel,
          """fill-mask""": DistilBertForMaskedLM,
          """question-answering""": DistilBertForQuestionAnswering,
          """text-classification""": DistilBertForSequenceClassification,
          """token-classification""": DistilBertForTokenClassification,
          """zero-shot""": DistilBertForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  lowerCAmelCase_					=				True
  lowerCAmelCase_					=				True
  lowerCAmelCase_					=				True
  lowerCAmelCase_					=				True


  def   UpperCamelCase__      (				self    :					Dict  )     ->						Optional[Any]:
    _UpperCamelCase       			=DistilBertModelTester(self  )
    _UpperCamelCase       			=ConfigTester(self							,	config_class=_lowercase							,	dim=37  )


  def   UpperCamelCase__      (				self    :					int  )     ->						Union[str, Any]:
    self.config_tester.run_common_tests()


  def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Any:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_model(*_lowercase  )


  def   UpperCamelCase__      (				self    :					str  )     ->						str:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase  )


  def   UpperCamelCase__      (				self    :					Any  )     ->						Optional[Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase  )


  def   UpperCamelCase__      (				self    :					int  )     ->						List[Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase  )


  def   UpperCamelCase__      (				self    :					str  )     ->						Optional[int]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase  )


  def   UpperCamelCase__      (				self    :					str  )     ->						List[Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase  )


  @slow
  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Tuple:
    for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      _UpperCamelCase       			=DistilBertModel.from_pretrained(_lowercase  )
      self.assertIsNotNone(_lowercase  )




  @slow
  @require_torch_gpu
  def   UpperCamelCase__      (				self    :					Optional[int]  )     ->						Union[str, Any]:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()
    for model_class in self.all_model_classes:
      # BertForMultipleChoice behaves incorrectly in JIT environments.
      if model_class == DistilBertForMultipleChoice:
        return

      _UpperCamelCase       			=True
      _UpperCamelCase       			=model_class(config=_lowercase  )

      _UpperCamelCase       			=self._prepare_for_class(_lowercase							,	_lowercase  )
      _UpperCamelCase       			=torch.jit.trace(
          _lowercase							,	(inputs_dict['''input_ids'''].to('''cpu'''  ), inputs_dict['''attention_mask'''].to('''cpu'''  ))  )

      with tempfile.TemporaryDirectory() as tmp:
        torch.jit.save(_lowercase							,	os.path.join(_lowercase							,	'''traced_model.pt'''  )  )
        _UpperCamelCase       			=torch.jit.load(os.path.join(_lowercase							,	'''traced_model.pt'''  )							,	map_location=_lowercase  )
        loaded(inputs_dict['''input_ids'''].to(_lowercase  )							,	inputs_dict['''attention_mask'''].to(_lowercase  )  )







@require_torch
class 		UpperCAmelCase   (				unittest.TestCase):


  """simple docstring"""







  @slow
  def   UpperCamelCase__      (				self    :					int  )     ->						int:
    _UpperCamelCase       			=DistilBertModel.from_pretrained('''distilbert-base-uncased'''  )
    _UpperCamelCase       			=torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]  )
    _UpperCamelCase       			=torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
    with torch.no_grad():
      _UpperCamelCase       			=model(_lowercase							,	attention_mask=_lowercase  )[0]
    _UpperCamelCase       			=torch.Size((1, 11, 768)  )
    self.assertEqual(output.shape							,	_lowercase  )
    _UpperCamelCase       			=torch.tensor(
        [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]  )

    self.assertTrue(torch.allclose(output[:, 1:4, 1:4]							,	_lowercase							,	atol=1E-4  )  )











'''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). """   ,     lowercase_   ,     )
class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				RobertaConfig
  lowerCAmelCase_					=				"""roberta"""


  def __init__(				self    :					Dict							,	UpperCamelCase__    :					List[str]  )     ->						List[Any]:
    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. """   ,     lowercase_   ,     )
class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				RobertaConfig
  lowerCAmelCase_					=				"""roberta"""


  def __init__(				self    :					Tuple							,	UpperCamelCase__    :					Any  )     ->						Dict:
    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   UpperCamelCase__      (				self    :					List[Any]							,	UpperCamelCase__    :					Optional[int]=None							,	UpperCamelCase__    :					Union[str, Any]=None							,	UpperCamelCase__    :					List[str]=None							,	UpperCamelCase__    :					Optional[Any]=None							,	UpperCamelCase__    :					int=None							,	UpperCamelCase__    :					Union[str, Any]=None							,	UpperCamelCase__    :					Tuple=None							,	UpperCamelCase__    :					Union[str, Any]=-1							,	UpperCamelCase__    :					Dict=False							,	)     ->						Union[str, Any]:
    _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






