code: 
  
  
  
  
  
  
  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)
  
  SCREAMING_SNAKE_CASE__		:       List[Any]   							=							get_tests_dir("""fixtures""")
  SCREAMING_SNAKE_CASE__		:       List[Any]   							=							os.path.join(dirname(get_tests_dir()))
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							dirname(os.path.abspath(__file__))
  
  # default torch.distributed port
  SCREAMING_SNAKE_CASE__		:       Dict   							=							"""10999"""
  
  SCREAMING_SNAKE_CASE__		:       str   							=							"""t5-small"""
  
  # *** Working Models ***
  SCREAMING_SNAKE_CASE__		:       List[str]   							=							"""hf-internal-testing/tiny-albert"""
  SCREAMING_SNAKE_CASE__		:       Dict   							=							"""sshleifer/bart-tiny-random"""
  SCREAMING_SNAKE_CASE__		:       Any   							=							"""hf-internal-testing/tiny-bert"""
  SCREAMING_SNAKE_CASE__		:       List[Any]   							=							"""hf-internal-testing/tiny-random-bigbird_pegasus"""
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							"""hf-internal-testing/tiny-random-big_bird"""
  SCREAMING_SNAKE_CASE__		:       Optional[Any]   							=							"""hf-internal-testing/tiny-random-blenderbot"""
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							"""bigscience/bigscience-small-testing"""
  SCREAMING_SNAKE_CASE__		:       List[Any]   							=							"""hf-internal-testing/tiny-random-deberta"""
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							"""hf-internal-testing/tiny-random-deberta-v2"""
  SCREAMING_SNAKE_CASE__		:       Any   							=							"""sshleifer/tiny-distilbert-base-cased"""
  SCREAMING_SNAKE_CASE__		:       str   							=							"""hf-internal-testing/tiny-electra"""
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							"""hf-internal-testing/tiny-random-flaubert"""
  SCREAMING_SNAKE_CASE__		:       Tuple   							=							"""stas/tiny-wmt19-en-de"""
  SCREAMING_SNAKE_CASE__		:       Any   							=							"""hf-internal-testing/tiny-random-funnel"""
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							"""sshleifer/tiny-gpt2"""
  SCREAMING_SNAKE_CASE__		:       Optional[Any]   							=							"""hf-internal-testing/tiny-random-gptj"""
  SCREAMING_SNAKE_CASE__		:       Union[str, Any]   							=							"""hf-internal-testing/tiny-random-gpt_neo"""
  SCREAMING_SNAKE_CASE__		:       str   							=							"""hf-internal-testing/tiny-layoutlm"""
  SCREAMING_SNAKE_CASE__		:       str   							=							"""hf-internal-testing/tiny-random-led"""
  SCREAMING_SNAKE_CASE__		:       Tuple   							=							"""hf-internal-testing/tiny-random-longformer"""
  SCREAMING_SNAKE_CASE__		:       Tuple   							=							"""stas/tiny-m2m_100"""  # hf tiny model is unsuitable
  SCREAMING_SNAKE_CASE__		:       Optional[Any]   							=							"""sshleifer/tiny-marian-en-de"""
  SCREAMING_SNAKE_CASE__		:       int   							=							"""sshleifer/tiny-mbart"""
  SCREAMING_SNAKE_CASE__		:       int   							=							"""hf-internal-testing/tiny-random-mobilebert"""
  SCREAMING_SNAKE_CASE__		:       Any   							=							"""hf-internal-testing/tiny-random-mpnet"""
  SCREAMING_SNAKE_CASE__		:       Optional[Any]   							=							"""stas/pegasus-cnn_dailymail-tiny-random"""
  SCREAMING_SNAKE_CASE__		:       List[Any]   							=							"""hf-internal-testing/tiny-random-prophetnet"""
  SCREAMING_SNAKE_CASE__		:       Any   							=							"""sshleifer/tiny-distilroberta-base"""
  SCREAMING_SNAKE_CASE__		:       str   							=							"""hf-internal-testing/tiny-random-squeezebert"""
  SCREAMING_SNAKE_CASE__		:       Union[str, Any]   							=							"""patrickvonplaten/t5-tiny-random"""
  SCREAMING_SNAKE_CASE__		:       Dict   							=							"""hf-internal-testing/tiny-random-t5-v1.1"""
  SCREAMING_SNAKE_CASE__		:       int   							=							"""hf-internal-testing/tiny-random-vit"""
  SCREAMING_SNAKE_CASE__		:       List[str]   							=							"""hf-internal-testing/tiny-xlm-roberta"""
  SCREAMING_SNAKE_CASE__		:       Tuple   							=							"""sshleifer/tiny-xlnet-base-cased"""
  
  
  # *** To Fix ***
  
  
  # *** tiny model issues ***
  # missing model files:
  SCREAMING_SNAKE_CASE__		:       List[str]   							=							"""hf-internal-testing/tiny-random-mt5"""
  SCREAMING_SNAKE_CASE__		:       str   							=							"""hf-internal-testing/tiny-random-camembert"""
  SCREAMING_SNAKE_CASE__		:       Any   							=							"""hf-internal-testing/tiny-random-openai-gpt"""
  
  # missing tokenizer files
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							"""hf-internal-testing/tiny-random-convbert"""
  SCREAMING_SNAKE_CASE__		:       int   							=							"""hf-internal-testing/tiny-random-layoutlmv2"""
  SCREAMING_SNAKE_CASE__		:       List[Any]   							=							"""hf-internal-testing/tiny-random-hubert"""
  
  # issues with tokenizer
  SCREAMING_SNAKE_CASE__		:       List[str]   							=							"""hf-internal-testing/tiny-random-ctrl"""
  SCREAMING_SNAKE_CASE__		:       Union[str, Any]   							=							"""hf-internal-testing/tiny-random-transfo-xl"""  # same as ctrl
  
  # other issues with tiny models
  SCREAMING_SNAKE_CASE__		:       Optional[int]   							=							"""hf-internal-testing/tiny-random-ibert"""  # multiple issues with either mlm/qa/clas
  SCREAMING_SNAKE_CASE__		:       Dict   							=							"""hf-internal-testing/tiny-random-reformer"""  # multiple issues with either mlm/qa/clas
  
  # *** Lacking official examples to test with ***
  # or not working with examples
  SCREAMING_SNAKE_CASE__		:       List[str]   							=							"""hf-internal-testing/tiny-random-dpr"""
  # - "dpr"  examples/research_projects/rag-end2end-retriever/
  SCREAMING_SNAKE_CASE__		:       str   							=							"""hf-internal-testing/tiny-random-rag"""
  # - "rag" research_projects
  SCREAMING_SNAKE_CASE__		:       Any   							=							""""""
  # - "luke" Entities classes - no plan to make such example
  SCREAMING_SNAKE_CASE__		:       int   							=							"""hf-internal-testing/tiny-random-lxmert"""
  # - "lxmert" doesn't work with run_qa.py
  SCREAMING_SNAKE_CASE__		:       int   							=							"""hf-internal-testing/tiny-random-clip"""
  # - "clip" nothing under pytorch examples - XXX: Suraj is working on adding some - check by end of Sep
  SCREAMING_SNAKE_CASE__		:       List[Any]   							=							"""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
  SCREAMING_SNAKE_CASE__		:       int   							=							"""hf-internal-testing/tiny-random-tapas"""
  
  
  
  
  
  def    _lowerCamelCase	(	__lowerCamelCase=False					)   -> int:
  
       '''simple docstring'''
       # 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)
       UpperCAmelCase__			:     int										= min(2   ,       get_gpu_count()					) if distributed else 1
       UpperCAmelCase__			:     List[str]										= os.environ.get("""DS_TEST_PORT"""   ,       __lowerCamelCase					)
       return F"deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}".split()
  
  
  
  
  
  def    _lowerCamelCase	(	)   -> Any:
  
       '''simple docstring'''
       UpperCAmelCase__			:     str										= F"{FIXTURE_DIRECTORY}/tests_samples"
       UpperCAmelCase__			:     List[str]										= F"{data_dir_samples}/wmt_en_ro"
       UpperCAmelCase__			:     Any										= F"{data_dir_samples}/xsum"
       UpperCAmelCase__			:     str										= """
          --do_train
          --max_train_samples 4
          --per_device_train_batch_size 2
          --num_train_epochs 1
          --fp16
          --report_to none
          --overwrite_output_dir
          """.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.
       UpperCAmelCase__			:     Tuple										= {
           """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""",
           ],
       }
  
       UpperCAmelCase__			:     List[Any]										= F"{ROOT_DIRECTORY}/examples/pytorch"
  
       UpperCAmelCase__			:     str										= {
           """trans""": F"\n        {scripts_dir}/translation/run_translation.py\n        --train_file {data_dir_wmt}/train.json\n        --source_lang en\n        --target_lang ro\n        ",
           """sum""": F"\n        {scripts_dir}/summarization/run_summarization.py\n        --train_file {data_dir_xsum}/sample.json\n        --max_source_length 12\n        --max_target_length 12\n        --lang en\n        ",
           """clm""": F"\n        {scripts_dir}/language-modeling/run_clm.py\n        --train_file {FIXTURE_DIRECTORY}/sample_text.txt\n        --block_size 8\n        ",
           """mlm""": F"\n        {scripts_dir}/language-modeling/run_mlm.py\n        --train_file {FIXTURE_DIRECTORY}/sample_text.txt\n        ",
           """qa""": F"\n        {scripts_dir}/question-answering/run_qa.py\n        --train_file {data_dir_samples}/SQUAD/sample.json\n        ",
           """clas""": F"\n        {scripts_dir}/text-classification/run_glue.py\n        --train_file {data_dir_samples}/MRPC/train.csv\n        --max_seq_length 12\n        --task_name MRPC\n        ",
           """img_clas""": F"\n        {scripts_dir}/image-classification/run_image_classification.py\n            --dataset_name hf-internal-testing/cats_vs_dogs_sample\n            --remove_unused_columns False\n            --max_steps 10\n            --image_processor_name {DS_TESTS_DIRECTORY}/vit_feature_extractor.json\n        ",
       }
  
       UpperCAmelCase__			:     Any										= get_launcher(distributed=__lowerCamelCase					)
  
       UpperCAmelCase__			:     Tuple										= {}
       for task, args in tasks.items():
            UpperCAmelCase__			:     Any										= args.split()
            for model in tasksamodels[task]:
                 UpperCAmelCase__			:     List[str]										= globals()[F"{model.upper().replace('-'   ,       '_'					)}_TINY"]
                 UpperCAmelCase__			:     Union[str, Any]										= F"--model_name_or_path {model_name}".split()
                 UpperCAmelCase__			:     Tuple										= 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
  
  
  SCREAMING_SNAKE_CASE__		:       Any   							=							make_task_cmds()
  
  SCREAMING_SNAKE_CASE__		:       Any   							=							"""zero2"""
  SCREAMING_SNAKE_CASE__		:       Optional[Any]   							=							"""zero3"""
  
  SCREAMING_SNAKE_CASE__		:       Optional[Any]   							=							[ZEROa, ZEROa]
  
  
  
  
  
  def    _lowerCamelCase	(	__lowerCamelCase   ,       __lowerCamelCase   ,       __lowerCamelCase					)   -> Dict:
  
       '''simple docstring'''
       # 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
       UpperCAmelCase__			:     Dict										= parameterized.to_safe_name("""_""".join(str(__lowerCamelCase					) for x in param.args					)					)
       return F"{func.__name__}_{param_based_name}"
  
  
  # Cartesian-product of zero stages with models to test
  SCREAMING_SNAKE_CASE__		:       Tuple   							=							list(itertools.product(stages, task_cmds.keys()))
  
  
  
  
  
  
  
  @slow
  @require_deepspeed
  @require_torch_gpu
  class 							UpperCAmelCase_     (		__lowerCamelCase		):
  
  
  
  
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase		):
            # 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()}"		)
  
            UpperCAmelCase__			:     Union[str, Any]										= task_cmds[task]
            UpperCAmelCase__			:     List[str]										= f"--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json".split()
  
            UpperCAmelCase__			:     List[str]										= self.get_auto_remove_tmp_dir()
            UpperCAmelCase__			:     List[str]										= f"--output_dir {output_dir}".split()
  
            cmd += args_ds + args_out
  
            return cmd, output_dir
  
  
  
  
       @parameterized.expand(_lowerCAmelCase  ,	name_func=_lowerCAmelCase		)
       def      __UpperCAmelCase     (							self  ,	_lowerCAmelCase  ,	_lowerCAmelCase		):
            # testing the ability to do a run followed by recovery of full fp32 weights
  
            UpperCAmelCase__			,		UpperCAmelCase__			:     Union[str, Any]										= 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
            UpperCAmelCase__			:     Optional[Any]										= f"{output_dir}/checkpoint-1"
            UpperCAmelCase__			:     Dict										= f"{chkpt_dir}/out.bin"
            UpperCAmelCase__			:     List[str]										= 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
  
  
  
  
  
  
  
style_context: 
  
  
  
  
  
  
  from typing import TYPE_CHECKING
  
  from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available
  
  
  SCREAMING_SNAKE_CASE__		:       Any   							=							{"""configuration_dinat""": ["""DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DinatConfig"""]}
  
  
  try:
        if not is_torch_available():
              raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
        pass
  else:
        SCREAMING_SNAKE_CASE__		:       str   							=							[
            """DINAT_PRETRAINED_MODEL_ARCHIVE_LIST""",
            """DinatForImageClassification""",
            """DinatModel""",
            """DinatPreTrainedModel""",
            """DinatBackbone""",
        ]
  
  if TYPE_CHECKING:
        from .configuration_dinat import DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP, DinatConfig
  
        try:
              if not is_torch_available():
                    raise OptionalDependencyNotAvailable()
        except OptionalDependencyNotAvailable:
              pass
        else:
              from .modeling_dinat import (
                  DINAT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  DinatBackbone,
                  DinatForImageClassification,
                  DinatModel,
                  DinatPreTrainedModel,
              )
  
  else:
        import sys
  
        SCREAMING_SNAKE_CASE__		:       Dict   							=							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
  
  
  
  
  
  
  