




'''simple docstring'''





import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
   from sagemaker import Session, TrainingJobAnalytics
   from sagemaker.huggingface import HuggingFace

@pytest.mark.skipif(
    literal_eval(os.getenv("""TEST_SAGEMAKER"""   ,   """False"""				)				) is not True   ,   reason="""Skipping test because should only be run when releasing minor transformers version"""   ,   )
@pytest.mark.usefixtures("""sm_env"""				)
@parameterized_class(
    [
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.g4dn.xlarge""",
            """results""": {"""train_runtime""": 6_50, """eval_accuracy""": 0.6, """eval_loss""": 0.9},
        },
        {
            """framework""": """tensorflow""",
            """script""": """run_tf.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.g4dn.xlarge""",
            """results""": {"""train_runtime""": 6_00, """eval_accuracy""": 0.3, """eval_loss""": 0.9},
        },
    ]				)
class   _a							(		unittest.TestCase				):






     def 							_lowercase      (					self       )				->     int:
          if self.framework == "pytorch":
               subprocess.run(
                   f"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()   ,encoding="utf-8"   ,check=_SCREAMING_SNAKE_CASE   ,)
          assert hasattr(self   ,"env"       )






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE=1       )				->     Union[str, Any]:
          # creates estimator
          return HuggingFace(
              entry_point=self.script   ,source_dir=self.env.test_path   ,role=self.env.role   ,image_uri=self.env.image_uri   ,base_job_name=f"""{self.env.base_job_name}-single"""   ,instance_count=_SCREAMING_SNAKE_CASE   ,instance_type=self.instance_type   ,debugger_hook_config=_SCREAMING_SNAKE_CASE   ,hyperparameters={**self.env.hyperparameters, "model_name_or_path": self.model_name_or_path}   ,metric_definitions=self.env.metric_definitions   ,py_version="py36"   ,)






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE       )				->     int:
          TrainingJobAnalytics(_SCREAMING_SNAKE_CASE       ).export_csv(f"""{self.env.test_path}/{job_name}_metrics.csv"""       )






     def 							_lowercase      (					self       )				->     List[str]:
          # create estimator
          _snake_case				= self.create_estimator()

          # run training
          estimator.fit()

          # result dataframe
          _snake_case				= TrainingJobAnalytics(estimator.latest_training_job.name       ).dataframe()

          # extract kpis
          _snake_case				= list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]       )
          _snake_case				= list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]       )
          # get train time from SageMaker job, this includes starting, preprocessing, stopping
          _snake_case				= (
              Session().describe_training_job(estimator.latest_training_job.name       ).get("TrainingTimeInSeconds"   ,999_999       )
          )

          # assert kpis
          assert train_runtime <= self.results["train_runtime"]
          assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy       )
          assert all(t <= self.results["eval_loss"] for t in eval_loss       )

          # dump tests result into json file to share in PR
          with open(f"""{estimator.latest_training_job.name}.json"""   ,"w"       ) as outfile:
               json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}   ,_SCREAMING_SNAKE_CASE       )








'''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				):
     SCREAMING_SNAKE_CASE_				:       str       							=		RobertaConfig
     SCREAMING_SNAKE_CASE_				:       Dict       							=		"""roberta"""






     def __init__(					self   ,_SCREAMING_SNAKE_CASE       )				->     Optional[Any]:
          super().__init__(_SCREAMING_SNAKE_CASE       )

          _snake_case				= RobertaEmbeddings(_SCREAMING_SNAKE_CASE       )
          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				):
     SCREAMING_SNAKE_CASE_				:       List[Any]       							=		RobertaConfig
     SCREAMING_SNAKE_CASE_				:       List[str]       							=		"""roberta"""






     def __init__(					self   ,_SCREAMING_SNAKE_CASE       )				->     str:
          super().__init__(_SCREAMING_SNAKE_CASE       )
          _snake_case				= config.num_labels
          _snake_case				= config.num_hidden_layers

          _snake_case				= DeeRobertaModel(_SCREAMING_SNAKE_CASE       )
          _snake_case				= nn.Dropout(config.hidden_dropout_prob       )
          _snake_case				= nn.Linear(config.hidden_size   ,self.config.num_labels       )






     @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE       )
     def 							_lowercase      (					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=-1   ,_SCREAMING_SNAKE_CASE=False   ,)				->     List[str]:
          _snake_case				= self.num_layers
          try:
               _snake_case				= self.roberta(
                   _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   ,)

               _snake_case				= outputs[1]

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

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

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

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

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


