


"""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_5_0, """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_0_0, """eval_accuracy""": 0.3, """eval_loss""": 0.9},
        },
    ]  )
class 						_A  (  unittest.TestCase  ):







    """simple docstring"""




    def      __snake_case							( self    :   Tuple):




      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=__UpperCAmelCase    ,			)
      assert hasattr(self    ,			"env")






    def      __snake_case							( self    :   str    ,			__UpperCAmelCase    :   int=1):
      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=__UpperCAmelCase    ,			instance_type=self.instance_type    ,			debugger_hook_config=__UpperCAmelCase    ,			hyperparameters={**self.env.hyperparameters, "model_name_or_path": self.model_name_or_path}    ,			metric_definitions=self.env.metric_definitions    ,			py_version="py36"    ,			)






    def      __snake_case							( self    :   List[Any]    ,			__UpperCAmelCase    :   str):
      TrainingJobAnalytics(__UpperCAmelCase).export_csv(f'''{self.env.test_path}/{job_name}_metrics.csv''')






    def      __snake_case							( self    :   int):
      a						:		List[str]											=      self.create_estimator()

      # run training
      estimator.fit()

      # result dataframe
      a						:		List[str]											=      TrainingJobAnalytics(estimator.latest_training_job.name).dataframe()

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

      # 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}    ,			__UpperCAmelCase)





"""simple docstring"""







__lowercase    =				frozenset(
    [
        """prompt""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
        """cross_attention_kwargs""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """negative_prompt"""])

__lowercase    =				frozenset([])

__lowercase    =				frozenset(["""image"""])

__lowercase    =				frozenset(
    [
        """image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

__lowercase    =				frozenset(["""image"""])

__lowercase    =				frozenset(
    [
        """prompt""",
        """image""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """image""", """negative_prompt"""])

__lowercase    =				frozenset(
    [
        # Text guided image variation with an image mask
        """prompt""",
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """image""", """mask_image""", """negative_prompt"""])

__lowercase    =				frozenset(
    [
        # image variation with an image mask
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

__lowercase    =				frozenset(["""image""", """mask_image"""])

__lowercase    =				frozenset(
    [
        """example_image""",
        """image""",
        """mask_image""",
        """height""",
        """width""",
        """guidance_scale""",
    ]
)

__lowercase    =				frozenset(["""example_image""", """image""", """mask_image"""])

__lowercase    =				frozenset(["""class_labels"""])

__lowercase    =				frozenset(["""class_labels"""])

__lowercase    =				frozenset(["""batch_size"""])

__lowercase    =				frozenset([])

__lowercase    =				frozenset(["""batch_size"""])

__lowercase    =				frozenset([])

__lowercase    =				frozenset(
    [
        """prompt""",
        """audio_length_in_s""",
        """guidance_scale""",
        """negative_prompt""",
        """prompt_embeds""",
        """negative_prompt_embeds""",
        """cross_attention_kwargs""",
    ]
)

__lowercase    =				frozenset(["""prompt""", """negative_prompt"""])
__lowercase    =				frozenset(["""input_tokens"""])

__lowercase    =				frozenset(["""input_tokens"""])






