

"""simple docstring"""






import itertools
import os
import subprocess
from os.path import dirname

from parameterized import parameterized

from tests.trainer.test_trainer import TrainerIntegrationCommon  # noqa
from transformers import is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_tests_dir,
    require_deepspeed,
    require_torch_gpu,
    slow,
)
from transformers.trainer_utils import set_seed


if is_torch_available():
					from tests.trainer.test_trainer import (  # noqa
					    RegressionModelConfig,
					    RegressionPreTrainedModel,
					    get_regression_trainer,
					)


set_seed(42)

_lowercase							 =							get_tests_dir('''fixtures''')
_lowercase							 =							os.path.join(dirname(get_tests_dir()))
_lowercase							 =							dirname(os.path.abspath(__file__))

# default torch.distributed port
_lowercase							 =							'''10999'''

_lowercase							 =							'''t5-small'''

# *** Working Models ***
_lowercase							 =							'''hf-internal-testing/tiny-albert'''
_lowercase							 =							'''sshleifer/bart-tiny-random'''
_lowercase							 =							'''hf-internal-testing/tiny-bert'''
_lowercase							 =							'''hf-internal-testing/tiny-random-bigbird_pegasus'''
_lowercase							 =							'''hf-internal-testing/tiny-random-big_bird'''
_lowercase							 =							'''hf-internal-testing/tiny-random-blenderbot'''
_lowercase							 =							'''bigscience/bigscience-small-testing'''
_lowercase							 =							'''hf-internal-testing/tiny-random-deberta'''
_lowercase							 =							'''hf-internal-testing/tiny-random-deberta-v2'''
_lowercase							 =							'''sshleifer/tiny-distilbert-base-cased'''
_lowercase							 =							'''hf-internal-testing/tiny-electra'''
_lowercase							 =							'''hf-internal-testing/tiny-random-flaubert'''
_lowercase							 =							'''stas/tiny-wmt19-en-de'''
_lowercase							 =							'''hf-internal-testing/tiny-random-funnel'''
_lowercase							 =							'''sshleifer/tiny-gpt2'''
_lowercase							 =							'''hf-internal-testing/tiny-random-gptj'''
_lowercase							 =							'''hf-internal-testing/tiny-random-gpt_neo'''
_lowercase							 =							'''hf-internal-testing/tiny-layoutlm'''
_lowercase							 =							'''hf-internal-testing/tiny-random-led'''
_lowercase							 =							'''hf-internal-testing/tiny-random-longformer'''
_lowercase							 =							'''stas/tiny-m2m_100'''  # hf tiny model is unsuitable
_lowercase							 =							'''sshleifer/tiny-marian-en-de'''
_lowercase							 =							'''sshleifer/tiny-mbart'''
_lowercase							 =							'''hf-internal-testing/tiny-random-mobilebert'''
_lowercase							 =							'''hf-internal-testing/tiny-random-mpnet'''
_lowercase							 =							'''stas/pegasus-cnn_dailymail-tiny-random'''
_lowercase							 =							'''hf-internal-testing/tiny-random-prophetnet'''
_lowercase							 =							'''sshleifer/tiny-distilroberta-base'''
_lowercase							 =							'''hf-internal-testing/tiny-random-squeezebert'''
_lowercase							 =							'''patrickvonplaten/t5-tiny-random'''
_lowercase							 =							'''hf-internal-testing/tiny-random-t5-v1.1'''
_lowercase							 =							'''hf-internal-testing/tiny-random-vit'''
_lowercase							 =							'''hf-internal-testing/tiny-xlm-roberta'''
_lowercase							 =							'''sshleifer/tiny-xlnet-base-cased'''


# *** To Fix ***


# *** tiny model issues ***
# missing model files:
_lowercase							 =							'''hf-internal-testing/tiny-random-mt5'''
_lowercase							 =							'''hf-internal-testing/tiny-random-camembert'''
_lowercase							 =							'''hf-internal-testing/tiny-random-openai-gpt'''

# missing tokenizer files
_lowercase							 =							'''hf-internal-testing/tiny-random-convbert'''
_lowercase							 =							'''hf-internal-testing/tiny-random-layoutlmv2'''
_lowercase							 =							'''hf-internal-testing/tiny-random-hubert'''

# issues with tokenizer
_lowercase							 =							'''hf-internal-testing/tiny-random-ctrl'''
_lowercase							 =							'''hf-internal-testing/tiny-random-transfo-xl'''  # same as ctrl

# other issues with tiny models
_lowercase							 =							'''hf-internal-testing/tiny-random-ibert'''  # multiple issues with either mlm/qa/clas
_lowercase							 =							'''hf-internal-testing/tiny-random-reformer'''  # multiple issues with either mlm/qa/clas

# *** Lacking official examples to test with ***
# or not working with examples
_lowercase							 =							'''hf-internal-testing/tiny-random-dpr'''
# - "dpr"  examples/research_projects/rag-end2end-retriever/
_lowercase							 =							'''hf-internal-testing/tiny-random-rag'''
# - "rag" research_projects
_lowercase							 =							''''''
# - "luke" Entities classes - no plan to make such example
_lowercase							 =							'''hf-internal-testing/tiny-random-lxmert'''
# - "lxmert" doesn't work with run_qa.py
_lowercase							 =							'''hf-internal-testing/tiny-random-clip'''
# - "clip" nothing under pytorch examples - XXX: Suraj is working on adding some - check by end of Sep
_lowercase							 =							'''hf-internal-testing/tiny-random-speech_to_text'''
# - "speech_to_text", nothing under pytorch examples


# *** Reactive mode ***
# models with low usage, unstable API, things about to change - do nothing about the following until someone runs into a problem
_lowercase							 =							'''hf-internal-testing/tiny-random-tapas'''



def     lowerCAmelCase__ (		__magic_name__=False			)	->int:
			# 1. explicitly set --num_nodes=1 just in case these tests end up run on a multi-node setup
			# - it won't be able to handle that
			# 2. for now testing with just 2 gpus max (since some quality tests may give different
			# results with mode gpus because we use very little data)
			__lowercase						      =					min(2 ,      get_gpu_count()			) if distributed else 1
			__lowercase						      =					os.environ.get("DS_TEST_PORT" ,      __magic_name__			)
			return F'''deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}'''.split()



def     lowerCAmelCase__ (		)	->Optional[Any]:
			__lowercase						      =					F'''{FIXTURE_DIRECTORY}/tests_samples'''
			__lowercase						      =					F'''{data_dir_samples}/wmt_en_ro'''
			__lowercase						      =					F'''{data_dir_samples}/xsum'''
			__lowercase						      =					"\n        --do_train\n        --max_train_samples 4\n        --per_device_train_batch_size 2\n        --num_train_epochs 1\n        --fp16\n        --report_to none\n        --overwrite_output_dir\n        ".split()

			# try to cover as many models as possible once (it's enough to run on one task per model)
			# but need a tiny model for each
			#
			# should have "{model_type.upper()}_TINY" corresponding vars defined, e.g., T5_TINY, etc.
			__lowercase						      =					{
			    "trans": [
			        "bart",
			        "fsmt",
			        "m2m_100",
			        "marian",
			        "mbart",
			        "t5",
			        "t5_v1",
			        # "mt5", missing model files
			    ],
			    "sum": [
			        "pegasus",
			    ],
			    "clm": [
			        "big_bird",
			        "bigbird_pegasus",
			        "blenderbot",
			        "bloom",
			        "gpt2",
			        "gpt_neo",
			        "gptj",
			        "xlm-roberta",
			        "prophetnet",
			        # "camembert", missing model files
			    ],
			    "mlm": [
			        "albert",
			        "deberta",
			        "deberta-v2",
			        "distilbert",
			        "electra",
			        "flaubert",
			        "funnel",
			        "layoutlm",
			        # "reformer", # multiple issues with either mlm/qa/clas
			    ],
			    "qa": [
			        "led",
			        "longformer",
			        "mobilebert",
			        "mpnet",
			        "roberta",
			        "squeezebert",
			        # "convbert", # missing tokenizer files
			        # "layoutlmv2", missing model files
			    ],
			    "clas": [
			        "bert",
			        "xlnet",
			        # "hubert", # missing tokenizer files
			        # "ibert", # multiple issues with either mlm/qa/clas
			        # "transfo-xl", # tokenizer issues as ctrl
			        # "ctrl", # tokenizer issues
			        # "openai-gpt", missing model files
			        # "tapas", multiple issues
			    ],
			    "img_clas": [
			        "vit",
			    ],
			}

			__lowercase						      =					F'''{ROOT_DIRECTORY}/examples/pytorch'''

			__lowercase						      =					{
			    "trans": F'''
        {scripts_dir}/translation/run_translation.py
        --train_file {data_dir_wmt}/train.json
        --source_lang en
        --target_lang ro
        ''',
			    "sum": F'''
        {scripts_dir}/summarization/run_summarization.py
        --train_file {data_dir_xsum}/sample.json
        --max_source_length 12
        --max_target_length 12
        --lang en
        ''',
			    "clm": F'''
        {scripts_dir}/language-modeling/run_clm.py
        --train_file {FIXTURE_DIRECTORY}/sample_text.txt
        --block_size 8
        ''',
			    "mlm": F'''
        {scripts_dir}/language-modeling/run_mlm.py
        --train_file {FIXTURE_DIRECTORY}/sample_text.txt
        ''',
			    "qa": F'''
        {scripts_dir}/question-answering/run_qa.py
        --train_file {data_dir_samples}/SQUAD/sample.json
        ''',
			    "clas": F'''
        {scripts_dir}/text-classification/run_glue.py
        --train_file {data_dir_samples}/MRPC/train.csv
        --max_seq_length 12
        --task_name MRPC
        ''',
			    "img_clas": F'''
        {scripts_dir}/image-classification/run_image_classification.py
            --dataset_name hf-internal-testing/cats_vs_dogs_sample
            --remove_unused_columns False
            --max_steps 10
            --image_processor_name {DS_TESTS_DIRECTORY}/vit_feature_extractor.json
        ''',
			}

			__lowercase						      =					get_launcher(distributed=__magic_name__			)

			__lowercase						      =					{}
			for task, args in tasks.items():
						__lowercase						      =					args.split()
						for model in tasksamodels[task]:
									__lowercase						      =					globals()[F'''{model.upper().replace('-' ,      '_'			)}_TINY''']
									__lowercase						      =					F'''--model_name_or_path {model_name}'''.split()
									__lowercase						      =					launcher + args + args_model + args_main

									# # generation special case
									# if task == "gen":
									#     launcher = f"deepspeed --num_nodes 1 --num_gpus 1".split()
									#     args_model += f"--model_type {model}".split()
									#     cmds[f"{task}_{model}"] = launcher + args + args_model
									# else:

			return cmds


_lowercase							 =							make_task_cmds()

_lowercase							 =							'''zero2'''
_lowercase							 =							'''zero3'''

_lowercase							 =							[ZEROa, ZEROa]



def     lowerCAmelCase__ (		__magic_name__ ,      __magic_name__ ,      __magic_name__			)	->Any:
			# customize the test name generator function as we want both params to appear in the sub-test
			# name, as by default it shows only the first param
			__lowercase						      =					parameterized.to_safe_name("_".join(str(__magic_name__			) for x in param.args			)			)
			return F'''{func.__name__}_{param_based_name}'''


# Cartesian-product of zero stages with models to test
_lowercase							 =							list(itertools.product(stages, task_cmds.keys()))







@slow
@require_deepspeed
@require_torch_gpu
class 				__a    (   __a					):



							'''simple docstring'''







							def      SCREAMING_SNAKE_CASE	(    self     ,						_lowerCamelCase     ,						_lowerCamelCase					) ->       Any:




										'''simple docstring'''
										# return a ready to run train cmd
										if task not in task_cmds:
													raise ValueError(f'''don\'t know of task {task}, have {task_cmds.keys()}'''					)

										__lowercase						      =					task_cmds[task]
										__lowercase						      =					f'''--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json'''.split()

										__lowercase						      =					self.get_auto_remove_tmp_dir()
										__lowercase						      =					f'''--output_dir {output_dir}'''.split()

										cmd += args_ds + args_out

										return cmd, output_dir





							@parameterized.expand(_lowerCamelCase     ,						name_func=_lowerCamelCase					)
							def      SCREAMING_SNAKE_CASE	(    self     ,						_lowerCamelCase     ,						_lowerCamelCase					) ->       Dict:




										'''simple docstring'''
										__lowercase      ,				__lowercase						      =					self.get_task_cmd(_lowerCamelCase     ,						_lowerCamelCase					)

										# 1. generate the checkpoint
										cmd += "--save_steps 1".split()
										# keep for quick debug
										# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] + cmd)); die
										execute_subprocess_async(_lowerCamelCase     ,						env=self.get_env()					)

										# 2. test that the fp32 weights get reconsolidated
										__lowercase						      =					f'''{output_dir}/checkpoint-1'''
										__lowercase						      =					f'''{chkpt_dir}/out.bin'''
										__lowercase						      =					f'''{chkpt_dir}/zero_to_fp32.py {chkpt_dir} {recovered_model_path}'''
										# keep for quick debug
										# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
										subprocess.check_call(_lowerCamelCase     ,						shell=_lowerCamelCase					)
										assert os.path.exists(_lowerCamelCase					), f'''{recovered_model_path} was not found'''

										# possibly could also test that the resulting saved model is usable but given that we use
										# random models we won't know if it's any good







"""simple docstring"""






import os

from packaging import version

from .. import __version__
from .accelerate_utils import apply_forward_hook
from .constants import (
    CONFIG_NAME,
    DEPRECATED_REVISION_ARGS,
    DIFFUSERS_CACHE,
    DIFFUSERS_DYNAMIC_MODULE_NAME,
    FLAX_WEIGHTS_NAME,
    HF_MODULES_CACHE,
    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
    ONNX_EXTERNAL_WEIGHTS_NAME,
    ONNX_WEIGHTS_NAME,
    SAFETENSORS_WEIGHTS_NAME,
    TEXT_ENCODER_ATTN_MODULE,
    WEIGHTS_NAME,
)
from .deprecation_utils import deprecate
from .doc_utils import replace_example_docstring
from .dynamic_modules_utils import get_class_from_dynamic_module
from .hub_utils import (
    HF_HUB_OFFLINE,
    _add_variant,
    _get_model_file,
    extract_commit_hash,
    http_user_agent,
)
from .import_utils import (
    BACKENDS_MAPPING,
    ENV_VARS_TRUE_AND_AUTO_VALUES,
    ENV_VARS_TRUE_VALUES,
    USE_JAX,
    USE_TF,
    USE_TORCH,
    DummyObject,
    OptionalDependencyNotAvailable,
    is_accelerate_available,
    is_accelerate_version,
    is_bsa_available,
    is_flax_available,
    is_ftfy_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_omegaconf_available,
    is_onnx_available,
    is_safetensors_available,
    is_scipy_available,
    is_tensorboard_available,
    is_tf_available,
    is_torch_available,
    is_torch_version,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    is_wandb_available,
    is_xformers_available,
    requires_backends,
)
from .logging import get_logger
from .outputs import BaseOutput
from .pil_utils import PIL_INTERPOLATION, numpy_to_pil, pt_to_pil
from .torch_utils import is_compiled_module, randn_tensor


if is_torch_available():
					from .testing_utils import (
					    floats_tensor,
					    load_hf_numpy,
					    load_image,
					    load_numpy,
					    load_pt,
					    nightly,
					    parse_flag_from_env,
					    print_tensor_test,
					    require_torch_a,
					    require_torch_gpu,
					    skip_mps,
					    slow,
					    torch_all_close,
					    torch_device,
					)
					from .torch_utils import maybe_allow_in_graph

from .testing_utils import export_to_gif, export_to_video


_lowercase							 =							get_logger(__name__)



def     lowerCAmelCase__ (		__magic_name__			)	->Optional[int]:
			if version.parse(__magic_name__			) < version.parse(__magic_name__			):
						if "dev" in min_version:
									__lowercase						      =					(
									    "This example requires a source install from HuggingFace diffusers (see "
									    "`https://huggingface.co/docs/diffusers/installation#install-from-source`),"
									)
						else:
									__lowercase						      =					F'''This example requires a minimum version of {min_version},'''
						error_message += F''' but the version found is {__version__}.\n'''
						raise ImportError(__magic_name__			)





