





import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest



class 							_UpperCamelCase					(	lowerCAmelCase_     ):
       _UpperCamelCase				:					Optional[int]         =   (PNDMScheduler,)
       _UpperCamelCase				:					List[Any]         =   (('''num_inference_steps''', 5_0),)

       def lowercase (		self: Optional[Any]					,   **_SCREAMING_SNAKE_CASE: Tuple		)       -> List[Any]:




              """simple docstring"""


              UpperCamelCase_          =  {
                  "num_train_timesteps": 1000,
                  "beta_start": 0.00_01,
                  "beta_end": 0.02,
                  "beta_schedule": "linear",
              }

              config.update(**_SCREAMING_SNAKE_CASE		)
              return config

       def lowercase (		self: List[Any]					,   _SCREAMING_SNAKE_CASE: Any=0					,   **_SCREAMING_SNAKE_CASE: Tuple		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  dict(self.forward_default_kwargs		)
              UpperCamelCase_          =  kwargs.pop("num_inference_steps"					,   _SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  self.dummy_sample
              UpperCamelCase_          =  0.1 * sample
              UpperCamelCase_          =  [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

              for scheduler_class in self.scheduler_classes:
                     UpperCamelCase_          =  self.get_scheduler_config(**_SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  scheduler_class(**_SCREAMING_SNAKE_CASE		)
                     scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)
                     # copy over dummy past residuals
                     UpperCamelCase_          =  dummy_past_residuals[:]

                     with tempfile.TemporaryDirectory() as tmpdirname:
                            scheduler.save_config(_SCREAMING_SNAKE_CASE		)
                            UpperCamelCase_          =  scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE		)
                            new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)
                            # copy over dummy past residuals
                            UpperCamelCase_          =  dummy_past_residuals[:]

                     UpperCamelCase_          =  scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample
                     UpperCamelCase_          =  new_scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample

                     assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

                     UpperCamelCase_          =  scheduler.step_plms(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample
                     UpperCamelCase_          =  new_scheduler.step_plms(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample

                     assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

       def lowercase (		self: List[str]		)       -> Optional[Any]:




              """simple docstring"""


              pass

       def lowercase (		self: Dict					,   _SCREAMING_SNAKE_CASE: Optional[Any]=0					,   **_SCREAMING_SNAKE_CASE: Dict		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  dict(self.forward_default_kwargs		)
              UpperCamelCase_          =  kwargs.pop("num_inference_steps"					,   _SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  self.dummy_sample
              UpperCamelCase_          =  0.1 * sample
              UpperCamelCase_          =  [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

              for scheduler_class in self.scheduler_classes:
                     UpperCamelCase_          =  self.get_scheduler_config()
                     UpperCamelCase_          =  scheduler_class(**_SCREAMING_SNAKE_CASE		)
                     scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)

                     # copy over dummy past residuals (must be after setting timesteps)
                     UpperCamelCase_          =  dummy_past_residuals[:]

                     with tempfile.TemporaryDirectory() as tmpdirname:
                            scheduler.save_config(_SCREAMING_SNAKE_CASE		)
                            UpperCamelCase_          =  scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE		)
                            # copy over dummy past residuals
                            new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)

                            # copy over dummy past residual (must be after setting timesteps)
                            UpperCamelCase_          =  dummy_past_residuals[:]

                     UpperCamelCase_          =  scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample
                     UpperCamelCase_          =  new_scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample

                     assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

                     UpperCamelCase_          =  scheduler.step_plms(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample
                     UpperCamelCase_          =  new_scheduler.step_plms(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample

                     assert torch.sum(torch.abs(output - new_output		)		) < 1e-5, "Scheduler outputs are not identical"

       def lowercase (		self: List[str]					,   **_SCREAMING_SNAKE_CASE: Optional[int]		)       -> Union[str, Any]:




              """simple docstring"""


              UpperCamelCase_          =  self.scheduler_classes[0]
              UpperCamelCase_          =  self.get_scheduler_config(**_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_          =  scheduler_class(**_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  10
              UpperCamelCase_          =  self.dummy_model()
              UpperCamelCase_          =  self.dummy_sample_deter
              scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)

              for i, t in enumerate(scheduler.prk_timesteps		):
                     UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		).prev_sample

              for i, t in enumerate(scheduler.plms_timesteps		):
                     UpperCamelCase_          =  model(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		)
                     UpperCamelCase_          =  scheduler.step_plms(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		).prev_sample

              return sample

       def lowercase (		self: int		)       -> Optional[int]:




              """simple docstring"""


              UpperCamelCase_          =  dict(self.forward_default_kwargs		)

              UpperCamelCase_          =  kwargs.pop("num_inference_steps"					,   _SCREAMING_SNAKE_CASE		)

              for scheduler_class in self.scheduler_classes:
                     UpperCamelCase_          =  self.get_scheduler_config()
                     UpperCamelCase_          =  scheduler_class(**_SCREAMING_SNAKE_CASE		)

                     UpperCamelCase_          =  self.dummy_sample
                     UpperCamelCase_          =  0.1 * sample

                     if num_inference_steps is not None and hasattr(_SCREAMING_SNAKE_CASE					,   "set_timesteps"		):
                            scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)
                     elif num_inference_steps is not None and not hasattr(_SCREAMING_SNAKE_CASE					,   "set_timesteps"		):
                            UpperCamelCase_          =  num_inference_steps

                     # copy over dummy past residuals (must be done after set_timesteps)
                     UpperCamelCase_          =  [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
                     UpperCamelCase_          =  dummy_past_residuals[:]

                     UpperCamelCase_          =  scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   0					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample
                     UpperCamelCase_          =  scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   1					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample

                     self.assertEqual(output_a.shape					,   sample.shape		)
                     self.assertEqual(output_a.shape					,   output_a.shape		)

                     UpperCamelCase_          =  scheduler.step_plms(_SCREAMING_SNAKE_CASE					,   0					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample
                     UpperCamelCase_          =  scheduler.step_plms(_SCREAMING_SNAKE_CASE					,   1					,   _SCREAMING_SNAKE_CASE					,   **_SCREAMING_SNAKE_CASE		).prev_sample

                     self.assertEqual(output_a.shape					,   sample.shape		)
                     self.assertEqual(output_a.shape					,   output_a.shape		)

       def lowercase (		self: int		)       -> str:




              """simple docstring"""


              for timesteps in [100, 1000]:
                     self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Any		)       -> int:




              """simple docstring"""


              for steps_offset in [0, 1]:
                     self.check_over_configs(steps_offset=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_          =  self.scheduler_classes[0]
              UpperCamelCase_          =  self.get_scheduler_config(steps_offset=1		)
              UpperCamelCase_          =  scheduler_class(**_SCREAMING_SNAKE_CASE		)
              scheduler.set_timesteps(10		)
              assert torch.equal(
                  scheduler.timesteps					,   torch.LongTensor(
                      [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]		)					,   )

       def lowercase (		self: Optional[Any]		)       -> Dict:




              """simple docstring"""


              for beta_start, beta_end in zip([0.00_01, 0.0_01]					,   [0.0_02, 0.02]		):
                     self.check_over_configs(beta_start=_SCREAMING_SNAKE_CASE					,   beta_end=_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Tuple		)       -> List[Any]:




              """simple docstring"""


              for schedule in ["linear", "squaredcos_cap_v2"]:
                     self.check_over_configs(beta_schedule=_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Tuple		)       -> Any:




              """simple docstring"""


              for prediction_type in ["epsilon", "v_prediction"]:
                     self.check_over_configs(prediction_type=_SCREAMING_SNAKE_CASE		)

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




              """simple docstring"""


              for t in [1, 5, 10]:
                     self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Tuple		)       -> List[Any]:




              """simple docstring"""


              for t, num_inference_steps in zip([1, 5, 10]					,   [10, 50, 100]		):
                     self.check_over_forward(num_inference_steps=_SCREAMING_SNAKE_CASE		)

       def lowercase (		self: Tuple		)       -> List[str]:




              """simple docstring"""


              UpperCamelCase_          =  27

              for scheduler_class in self.scheduler_classes:
                     UpperCamelCase_          =  self.dummy_sample
                     UpperCamelCase_          =  0.1 * sample

                     UpperCamelCase_          =  self.get_scheduler_config()
                     UpperCamelCase_          =  scheduler_class(**_SCREAMING_SNAKE_CASE		)

                     scheduler.set_timesteps(_SCREAMING_SNAKE_CASE		)

                     # before power of 3 fix, would error on first step, so we only need to do two
                     for i, t in enumerate(scheduler.prk_timesteps[:2]		):
                            UpperCamelCase_          =  scheduler.step_prk(_SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE					,   _SCREAMING_SNAKE_CASE		).prev_sample

       def lowercase (		self: int		)       -> List[str]:




              """simple docstring"""


              with self.assertRaises(_SCREAMING_SNAKE_CASE		):
                     UpperCamelCase_          =  self.scheduler_classes[0]
                     UpperCamelCase_          =  self.get_scheduler_config()
                     UpperCamelCase_          =  scheduler_class(**_SCREAMING_SNAKE_CASE		)

                     scheduler.step_plms(self.dummy_sample					,   1					,   self.dummy_sample		).prev_sample

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




              """simple docstring"""


              UpperCamelCase_          =  self.full_loop()
              UpperCamelCase_          =  torch.sum(torch.abs(_SCREAMING_SNAKE_CASE		)		)
              UpperCamelCase_          =  torch.mean(torch.abs(_SCREAMING_SNAKE_CASE		)		)

              assert abs(result_sum.item() - 1_98.13_18		) < 1e-2
              assert abs(result_mean.item() - 0.25_80		) < 1e-3

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




              """simple docstring"""


              UpperCamelCase_          =  self.full_loop(prediction_type="v_prediction"		)
              UpperCamelCase_          =  torch.sum(torch.abs(_SCREAMING_SNAKE_CASE		)		)
              UpperCamelCase_          =  torch.mean(torch.abs(_SCREAMING_SNAKE_CASE		)		)

              assert abs(result_sum.item() - 67.39_86		) < 1e-2
              assert abs(result_mean.item() - 0.08_78		) < 1e-3

       def lowercase (		self: Optional[int]		)       -> List[Any]:




              """simple docstring"""


              UpperCamelCase_          =  self.full_loop(set_alpha_to_one=_SCREAMING_SNAKE_CASE					,   beta_start=0.01		)
              UpperCamelCase_          =  torch.sum(torch.abs(_SCREAMING_SNAKE_CASE		)		)
              UpperCamelCase_          =  torch.mean(torch.abs(_SCREAMING_SNAKE_CASE		)		)

              assert abs(result_sum.item() - 2_30.03_99		) < 1e-2
              assert abs(result_mean.item() - 0.29_95		) < 1e-3



       def lowercase (		self: Optional[Any]		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  self.full_loop(set_alpha_to_one=_SCREAMING_SNAKE_CASE					,   beta_start=0.01		)
              UpperCamelCase_          =  torch.sum(torch.abs(_SCREAMING_SNAKE_CASE		)		)
              UpperCamelCase_          =  torch.mean(torch.abs(_SCREAMING_SNAKE_CASE		)		)

              assert abs(result_sum.item() - 1_86.94_82		) < 1e-2
              assert abs(result_mean.item() - 0.24_34		) < 1e-3








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

import pytest
from parameterized import parameterized, 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.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_5_0, '''eval_accuracy''': 0.7, '''eval_loss''': 0.6},
        },
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_ddp.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.7, '''eval_loss''': 0.6},
        },
        {
            '''framework''': '''tensorflow''',
            '''script''': '''run_tf_dist.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.p3.16xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.6, '''eval_loss''': 0.7},
        },
    ]     )
class 							_UpperCamelCase					(	unittest.TestCase     ):

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




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

       def lowercase (		self: List[str]					,   _SCREAMING_SNAKE_CASE: Dict		)       -> str:




              """simple docstring"""


              UpperCamelCase_          =  f'''{self.env.base_job_name}-{instance_count}-{"ddp" if "ddp" in self.script else "smd"}'''
              # distributed data settings
              UpperCamelCase_          =  {"smdistributed": {"dataparallel": {"enabled": True}}} if self.script != "run_ddp.py" else None

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

       def lowercase (		self: List[str]					,   _SCREAMING_SNAKE_CASE: Any		)       -> List[Any]:




              """simple docstring"""


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



       @parameterized.expand([(2,)]		)
       def lowercase (		self: str					,   _SCREAMING_SNAKE_CASE: Tuple		)       -> Optional[Any]:




              """simple docstring"""


              UpperCamelCase_          =  self.create_estimator(_SCREAMING_SNAKE_CASE		)

              # run training
              estimator.fit()

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

              # extract kpis
              UpperCamelCase_          =  list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]		)
              UpperCamelCase_          =  list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]		)
              # get train time from SageMaker job, this includes starting, preprocessing, stopping
              UpperCamelCase_          =  (
                  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}					,   _SCREAMING_SNAKE_CASE		)

