






"""simple docstring"""




import unittest

import numpy as np
import torch

from diffusers import PNDMPipeline, PNDMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()




class    lowercase_    (unittest.TestCase     ):






      @property
      def 			lowerCamelCase__    (    self	)	->Dict:







        '''simple docstring'''


        torch.manual_seed(0	)
        _a							=	UNetaDModel(
            block_out_channels=(3_2, 6_4)     ,     layers_per_block=2     ,     sample_size=3_2     ,     in_channels=3     ,     out_channels=3     ,     down_block_types=("DownBlock2D", "AttnDownBlock2D")     ,     up_block_types=("AttnUpBlock2D", "UpBlock2D")     ,     )
        return model






      def 			lowerCamelCase__    (    self	)	->List[Any]:







        '''simple docstring'''


        _a							=	self.dummy_uncond_unet
        _a							=	PNDMScheduler()

        _a							=	PNDMPipeline(unet=a_     ,     scheduler=a_	)
        pndm.to(a_	)
        pndm.set_progress_bar_config(disable=a_	)

        _a							=	torch.manual_seed(0	)
        _a							=	pndm(generator=a_     ,     num_inference_steps=2_0     ,     output_type="numpy"	).images

        _a							=	torch.manual_seed(0	)
        _a							=	pndm(generator=a_     ,     num_inference_steps=2_0     ,     output_type="numpy"     ,     return_dict=a_	)[0]

        _a							=	image[0, -3:, -3:, -1]
        _a							=	image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 3_2, 3_2, 3)
        _a							=	np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]	)

        assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1E-2







@slow
@require_torch
class    lowercase_    (unittest.TestCase     ):






      def 			lowerCamelCase__    (    self	)	->List[str]:







        '''simple docstring'''


        _a							=	"google/ddpm-cifar10-32"

        _a							=	UNetaDModel.from_pretrained(a_	)
        _a							=	PNDMScheduler()

        _a							=	PNDMPipeline(unet=a_     ,     scheduler=a_	)
        pndm.to(a_	)
        pndm.set_progress_bar_config(disable=a_	)
        _a							=	torch.manual_seed(0	)
        _a							=	pndm(generator=a_     ,     output_type="numpy"	).images

        _a							=	image[0, -3:, -3:, -1]

        assert image.shape == (1, 3_2, 3_2, 3)
        _a							=	np.array([0.1_564, 0.14_645, 0.1_406, 0.14_715, 0.12_425, 0.14_045, 0.13_115, 0.12_175, 0.125]	)

        assert np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2










"""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''': 650, '''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''': 600, '''eval_accuracy''': 0.3, '''eval_loss''': 0.9},
        },
    ]     )
class    lowercase_    (unittest.TestCase     ):






      def 			lowerCamelCase__    (    self	)	->Union[str, Any]:







        '''simple docstring'''


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






      def 			lowerCamelCase__    (    self     ,     a_=1	)	->int:







        '''simple docstring'''


        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=a_     ,     instance_type=self.instance_type     ,     debugger_hook_config=a_     ,     hyperparameters={**self.env.hyperparameters, "model_name_or_path": self.model_name_or_path}     ,     metric_definitions=self.env.metric_definitions     ,     py_version="py36"     ,     )






      def 			lowerCamelCase__    (    self     ,     a_	)	->str:







        '''simple docstring'''


        TrainingJobAnalytics(a_	).export_csv(f'''{self.env.test_path}/{job_name}_metrics.csv'''	)






      def 			lowerCamelCase__    (    self	)	->Any:







        '''simple docstring'''


        _a							=	self.create_estimator()

        # run training
        estimator.fit()

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

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

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


