





import argparse
import json
import logging
import os
import shutil
import sys
import tempfile
import unittest
from unittest import mock

import torch
from accelerate.utils import write_basic_config

from transformers.testing_utils import TestCasePlus, get_gpu_count, run_command, slow, torch_device
from transformers.utils import is_apex_available


logging.basicConfig(level=logging.DEBUG)

lowercase    : Any    	=  logging.getLogger()


def     SCREAMING_SNAKE_CASE			(      ):
       _UpperCamelCase              =						argparse.ArgumentParser()
       parser.add_argument('''-f'''					)
       _UpperCamelCase              =						parser.parse_args()
       return args.f


def     SCREAMING_SNAKE_CASE			(      lowerCAmelCase					):
       _UpperCamelCase              =						{}
       _UpperCamelCase              =						os.path.join(__UpperCamelCase    ,       '''all_results.json'''					)
       if os.path.exists(__UpperCamelCase					):
              with open(__UpperCamelCase    ,       '''r'''					) as f:
                     _UpperCamelCase              =						json.load(__UpperCamelCase					)
       else:
              raise ValueError(f'''can\'t find {path}'''					)
       return results


def     SCREAMING_SNAKE_CASE			(      ):
       _UpperCamelCase              =						torch.cuda.is_available() and torch_device == '''cuda'''
       return is_using_cuda and is_apex_available()


lowercase    : Union[str, Any]    	=  logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)

class  __A(						__UpperCAmelCase ):

     @classmethod
     def 	_UpperCamelCase     (      cls      ):





            """simple docstring"""

            _UpperCamelCase              =						tempfile.mkdtemp()
            _UpperCamelCase              =						os.path.join(cls.tmpdir,       '''default_config.yml'''      )
            write_basic_config(save_location=cls.configPath      )
            _UpperCamelCase              =						['''accelerate''', '''launch''', '''--config_file''', cls.configPath]

     @classmethod
     def 	_UpperCamelCase     (      cls      ):





            """simple docstring"""

            shutil.rmtree(cls.tmpdir      )

     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/text-classification/run_glue_no_trainer.py\n            --model_name_or_path distilbert-base-uncased\n            --output_dir {tmp_dir}\n            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv\n            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --learning_rate=1e-4\n            --seed=42\n            --checkpointing_steps epoch\n            --with_tracking\n        '''.split()

            if is_cuda_and_apex_available():
                   testargs.append('''--fp16'''      )

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertGreaterEqual(result['''eval_accuracy'''],       0.75      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''epoch_0'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''glue_no_trainer'''      )      )      )

     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/language-modeling/run_clm_no_trainer.py\n            --model_name_or_path distilgpt2\n            --train_file ./tests/fixtures/sample_text.txt\n            --validation_file ./tests/fixtures/sample_text.txt\n            --block_size 128\n            --per_device_train_batch_size 5\n            --per_device_eval_batch_size 5\n            --num_train_epochs 2\n            --output_dir {tmp_dir}\n            --checkpointing_steps epoch\n            --with_tracking\n        '''.split()

            if torch.cuda.device_count() > 1:
                   # Skipping because there are not enough batches to train the model + would need a drop_last to work.
                   return

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertLess(result['''perplexity'''],       100      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''epoch_0'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''clm_no_trainer'''      )      )      )

     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/language-modeling/run_mlm_no_trainer.py\n            --model_name_or_path distilroberta-base\n            --train_file ./tests/fixtures/sample_text.txt\n            --validation_file ./tests/fixtures/sample_text.txt\n            --output_dir {tmp_dir}\n            --num_train_epochs=1\n            --checkpointing_steps epoch\n            --with_tracking\n        '''.split()

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertLess(result['''perplexity'''],       42      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''epoch_0'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''mlm_no_trainer'''      )      )      )

     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						7 if get_gpu_count() > 1 else 2

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/token-classification/run_ner_no_trainer.py\n            --model_name_or_path bert-base-uncased\n            --train_file tests/fixtures/tests_samples/conll/sample.json\n            --validation_file tests/fixtures/tests_samples/conll/sample.json\n            --output_dir {tmp_dir}\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=2\n            --num_train_epochs={epochs}\n            --seed 7\n            --checkpointing_steps epoch\n            --with_tracking\n        '''.split()

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertGreaterEqual(result['''eval_accuracy'''],       0.75      )
            self.assertLess(result['''train_loss'''],       0.5      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''epoch_0'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''ner_no_trainer'''      )      )      )

     @unittest.skip(reason='''Fix me @muellerzr'''      )
     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/question-answering/run_qa_no_trainer.py\n            --model_name_or_path bert-base-uncased\n            --version_2_with_negative\n            --train_file tests/fixtures/tests_samples/SQUAD/sample.json\n            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json\n            --output_dir {tmp_dir}\n            --seed=42\n            --max_train_steps=10\n            --num_warmup_steps=2\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --checkpointing_steps epoch\n            --with_tracking\n        '''.split()

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            # Because we use --version_2_with_negative the testing script uses SQuAD v2 metrics.
            self.assertGreaterEqual(result['''eval_f1'''],       28      )
            self.assertGreaterEqual(result['''eval_exact'''],       28      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''epoch_0'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''qa_no_trainer'''      )      )      )

     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/multiple-choice/run_swag_no_trainer.py\n            --model_name_or_path bert-base-uncased\n            --train_file tests/fixtures/tests_samples/swag/sample.json\n            --validation_file tests/fixtures/tests_samples/swag/sample.json\n            --output_dir {tmp_dir}\n            --max_train_steps=20\n            --num_warmup_steps=2\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --with_tracking\n        '''.split()

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertGreaterEqual(result['''eval_accuracy'''],       0.8      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''swag_no_trainer'''      )      )      )

     @slow
     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/summarization/run_summarization_no_trainer.py\n            --model_name_or_path t5-small\n            --train_file tests/fixtures/tests_samples/xsum/sample.json\n            --validation_file tests/fixtures/tests_samples/xsum/sample.json\n            --output_dir {tmp_dir}\n            --max_train_steps=50\n            --num_warmup_steps=8\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --checkpointing_steps epoch\n            --with_tracking\n        '''.split()

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertGreaterEqual(result['''eval_rouge1'''],       10      )
            self.assertGreaterEqual(result['''eval_rouge2'''],       2      )
            self.assertGreaterEqual(result['''eval_rougeL'''],       7      )
            self.assertGreaterEqual(result['''eval_rougeLsum'''],       7      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''epoch_0'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''summarization_no_trainer'''      )      )      )

     @slow
     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/translation/run_translation_no_trainer.py\n            --model_name_or_path sshleifer/student_marian_en_ro_6_1\n            --source_lang en\n            --target_lang ro\n            --train_file tests/fixtures/tests_samples/wmt16/sample.json\n            --validation_file tests/fixtures/tests_samples/wmt16/sample.json\n            --output_dir {tmp_dir}\n            --max_train_steps=50\n            --num_warmup_steps=8\n            --num_beams=6\n            --learning_rate=3e-3\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --source_lang en_XX\n            --target_lang ro_RO\n            --checkpointing_steps epoch\n            --with_tracking\n        '''.split()

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertGreaterEqual(result['''eval_bleu'''],       30      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''epoch_0'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''translation_no_trainer'''      )      )      )

     @slow
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						logging.StreamHandler(sys.stdout      )
            logger.addHandler(_SCREAMING_SNAKE_CASE      )

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/semantic-segmentation/run_semantic_segmentation_no_trainer.py\n            --dataset_name huggingface/semantic-segmentation-test-sample\n            --output_dir {tmp_dir}\n            --max_train_steps=10\n            --num_warmup_steps=2\n            --learning_rate=2e-4\n            --per_device_train_batch_size=2\n            --per_device_eval_batch_size=1\n            --checkpointing_steps epoch\n        '''.split()

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            self.assertGreaterEqual(result['''eval_overall_accuracy'''],       0.10      )



     @mock.patch.dict(os.environ,       {'''WANDB_MODE''': '''offline'''}      )
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						self.get_auto_remove_tmp_dir()
            _UpperCamelCase              =						F'''\n            {self.examples_dir}/pytorch/image-classification/run_image_classification_no_trainer.py\n            --model_name_or_path google/vit-base-patch16-224-in21k\n            --dataset_name hf-internal-testing/cats_vs_dogs_sample\n            --learning_rate 1e-4\n            --per_device_train_batch_size 2\n            --per_device_eval_batch_size 1\n            --max_train_steps 2\n            --train_val_split 0.1\n            --seed 42\n            --output_dir {tmp_dir}\n            --with_tracking\n            --checkpointing_steps 1\n        '''.split()

            if is_cuda_and_apex_available():
                   testargs.append('''--fp16'''      )

            run_command(self._launch_args + testargs      )
            _UpperCamelCase              =						get_results(_SCREAMING_SNAKE_CASE      )
            # The base model scores a 25%
            self.assertGreaterEqual(result['''eval_accuracy'''],       0.6      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''step_1'''      )      )      )
            self.assertTrue(os.path.exists(os.path.join(_SCREAMING_SNAKE_CASE,       '''image_classification_no_trainer'''      )      )      )













import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
       import jax
       import jax.numpy as jnp
       from flax.jax_utils import replicate
       from flax.training.common_utils import shard

@slow
@require_flax
class  __A(						unittest.TestCase ):

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            super().tearDown()
            gc.collect()

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxStableDiffusionPipeline.from_pretrained(
                '''stabilityai/stable-diffusion-2''',       revision='''bf16''',       dtype=jnp.bfloataa,       )

            _UpperCamelCase              =						'''A painting of a squirrel eating a burger'''
            _UpperCamelCase              =						jax.device_count()
            _UpperCamelCase              =						num_samples * [prompt]
            _UpperCamelCase              =						sd_pipe.prepare_inputs(A      )

            _UpperCamelCase              =						replicate(A      )
            _UpperCamelCase              =						shard(A      )

            _UpperCamelCase              =						jax.random.PRNGKey(0      )
            _UpperCamelCase              =						jax.random.split(A,       jax.device_count()      )

            _UpperCamelCase              =						sd_pipe(A,       A,       A,       num_inference_steps=25,       jit=A      )[0]
            assert images.shape == (jax.device_count(), 1, 768, 768, 3)

            _UpperCamelCase              =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
            _UpperCamelCase              =						images[0, 253:256, 253:256, -1]

            _UpperCamelCase              =						jnp.asarray(jax.device_get(image_slice.flatten()      )      )
            _UpperCamelCase              =						jnp.array([0.4_238, 0.4_414, 0.4_395, 0.4_453, 0.4_629, 0.4_590, 0.4_531, 0.45_508, 0.4_512]      )
            print(F'''output_slice: {output_slice}'''      )
            assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2



     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						'''stabilityai/stable-diffusion-2'''
            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxDPMSolverMultistepScheduler.from_pretrained(A,       subfolder='''scheduler'''      )
            _UpperCamelCase       ,		_UpperCamelCase              =						FlaxStableDiffusionPipeline.from_pretrained(
                A,       scheduler=A,       revision='''bf16''',       dtype=jnp.bfloataa,       )
            _UpperCamelCase              =						scheduler_params

            _UpperCamelCase              =						'''A painting of a squirrel eating a burger'''
            _UpperCamelCase              =						jax.device_count()
            _UpperCamelCase              =						num_samples * [prompt]
            _UpperCamelCase              =						sd_pipe.prepare_inputs(A      )

            _UpperCamelCase              =						replicate(A      )
            _UpperCamelCase              =						shard(A      )

            _UpperCamelCase              =						jax.random.PRNGKey(0      )
            _UpperCamelCase              =						jax.random.split(A,       jax.device_count()      )

            _UpperCamelCase              =						sd_pipe(A,       A,       A,       num_inference_steps=25,       jit=A      )[0]
            assert images.shape == (jax.device_count(), 1, 768, 768, 3)

            _UpperCamelCase              =						images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]      )
            _UpperCamelCase              =						images[0, 253:256, 253:256, -1]

            _UpperCamelCase              =						jnp.asarray(jax.device_get(image_slice.flatten()      )      )
            _UpperCamelCase              =						jnp.array([0.4_336, 0.42_969, 0.4_453, 0.4_199, 0.4_297, 0.4_531, 0.4_434, 0.4_434, 0.4_297]      )
            print(F'''output_slice: {output_slice}'''      )
            assert jnp.abs(output_slice - expected_slice      ).max() < 1E-2






