



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(4_2)

UpperCAmelCase_          =				get_tests_dir("fixtures")
UpperCAmelCase_          =				os.path.join(dirname(get_tests_dir()))
UpperCAmelCase_          =				dirname(os.path.abspath(__file__))

# default torch.distributed port
UpperCAmelCase_          =				"10999"

UpperCAmelCase_          =				"t5-small"

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


# *** To Fix ***


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

# missing tokenizer files
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-convbert"
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-layoutlmv2"
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-hubert"

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

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

# *** Lacking official examples to test with ***
# or not working with examples
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-dpr"
# - "dpr"  examples/research_projects/rag-end2end-retriever/
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-rag"
# - "rag" research_projects
UpperCAmelCase_          =				""
# - "luke" Entities classes - no plan to make such example
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-lxmert"
# - "lxmert" doesn't work with run_qa.py
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-clip"
# - "clip" nothing under pytorch examples - XXX: Suraj is working on adding some - check by end of Sep
UpperCAmelCase_          =				"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
UpperCAmelCase_          =				"hf-internal-testing/tiny-random-tapas"
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						str=False					)->List[str]:
      # 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)
      _lowerCAmelCase			      =	min(2  ,					get_gpu_count()					) if distributed else 1
      _lowerCAmelCase			      =	os.environ.get('''DS_TEST_PORT'''  ,					_SCREAMING_SNAKE_CASE					)
      return f'''deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}'''.split()
def 				UpperCAmelCase__		(					)->Optional[Any]:
      _lowerCAmelCase			      =	f'''{FIXTURE_DIRECTORY}/tests_samples'''
      _lowerCAmelCase			      =	f'''{data_dir_samples}/wmt_en_ro'''
      _lowerCAmelCase			      =	f'''{data_dir_samples}/xsum'''
      _lowerCAmelCase			      =	'''
        --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.
      _lowerCAmelCase			      =	{
          '''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''',
          ],
      }

      _lowerCAmelCase			      =	f'''{ROOT_DIRECTORY}/examples/pytorch'''

      _lowerCAmelCase			      =	{
          '''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
        ''',
      }

      _lowerCAmelCase			      =	get_launcher(distributed=_SCREAMING_SNAKE_CASE					)

      _lowerCAmelCase			      =	{}
      for task, args in tasks.items():
            _lowerCAmelCase			      =	args.split()
            for model in tasksamodels[task]:
                  _lowerCAmelCase			      =	globals()[f'''{model.upper().replace("-"  ,					"_"					)}_TINY''']
                  _lowerCAmelCase			      =	f'''--model_name_or_path {model_name}'''.split()
                  _lowerCAmelCase			      =	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


UpperCAmelCase_          =				make_task_cmds()

UpperCAmelCase_          =				"zero2"
UpperCAmelCase_          =				"zero3"

UpperCAmelCase_          =				[ZEROa, ZEROa]
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Union[str, Any]  ,					_SCREAMING_SNAKE_CASE :						Union[str, Any]  ,					_SCREAMING_SNAKE_CASE :						List[Any]					)->str:
      # 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
      _lowerCAmelCase			      =	parameterized.to_safe_name('''_'''.join(str(_SCREAMING_SNAKE_CASE					) for x in param.args					)					)
      return f'''{func.__name__}_{param_based_name}'''


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





@slow
@require_deepspeed
@require_torch_gpu
class 		UpperCAmelCase    (    snake_case_					):





   def 						__lowerCAmelCase				(						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()}'''      )

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

         _lowerCAmelCase			      =	self.get_auto_remove_tmp_dir()
         _lowerCAmelCase			      =	F'''--output_dir {output_dir}'''.split()

         cmd += args_ds + args_out

         return cmd, output_dir





   @parameterized.expand(_lowerCAmelCase							,  name_func=_lowerCAmelCase      )
   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase							,  _lowerCAmelCase      ):
         # testing the ability to do a run followed by recovery of full fp32 weights

         _lowerCAmelCase      , _lowerCAmelCase			      =	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
         _lowerCAmelCase			      =	F'''{output_dir}/checkpoint-1'''
         _lowerCAmelCase			      =	F'''{chkpt_dir}/out.bin'''
         _lowerCAmelCase			      =	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




from typing import List, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase_          =				logging.get_logger(__name__)

UpperCAmelCase_          =				{
    "allenai/led-base-16384": "https://huggingface.co/allenai/led-base-16384/resolve/main/config.json",
    # See all LED models at https://huggingface.co/models?filter=led
}





class 		UpperCAmelCase    (    snake_case_					):
   SCREAMING_SNAKE_CASE__					      =  '''led'''
   SCREAMING_SNAKE_CASE__					      =  {
       '''num_attention_heads''': '''encoder_attention_heads''',
       '''hidden_size''': '''d_model''',
       '''attention_probs_dropout_prob''': '''attention_dropout''',
       '''initializer_range''': '''init_std''',
   }





   def __init__(						self							,  _lowerCAmelCase=50_265							,  _lowerCAmelCase=16_384							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=12							,  _lowerCAmelCase=4_096							,  _lowerCAmelCase=16							,  _lowerCAmelCase=12							,  _lowerCAmelCase=4_096							,  _lowerCAmelCase=16							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=True							,  _lowerCAmelCase=True							,  _lowerCAmelCase="gelu"							,  _lowerCAmelCase=1_024							,  _lowerCAmelCase=0.1							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=0.02							,  _lowerCAmelCase=2							,  _lowerCAmelCase=0.0							,  _lowerCAmelCase=1							,  _lowerCAmelCase=0							,  _lowerCAmelCase=2							,  _lowerCAmelCase = 512							,  **_lowerCAmelCase							,  ):
         _lowerCAmelCase			      =	vocab_size
         _lowerCAmelCase			      =	max_encoder_position_embeddings
         _lowerCAmelCase			      =	max_decoder_position_embeddings
         _lowerCAmelCase			      =	d_model
         _lowerCAmelCase			      =	encoder_ffn_dim
         _lowerCAmelCase			      =	encoder_layers
         _lowerCAmelCase			      =	encoder_attention_heads
         _lowerCAmelCase			      =	decoder_ffn_dim
         _lowerCAmelCase			      =	decoder_layers
         _lowerCAmelCase			      =	decoder_attention_heads
         _lowerCAmelCase			      =	dropout
         _lowerCAmelCase			      =	attention_dropout
         _lowerCAmelCase			      =	activation_dropout
         _lowerCAmelCase			      =	activation_function
         _lowerCAmelCase			      =	init_std
         _lowerCAmelCase			      =	encoder_layerdrop
         _lowerCAmelCase			      =	decoder_layerdrop
         _lowerCAmelCase			      =	classifier_dropout
         _lowerCAmelCase			      =	use_cache
         _lowerCAmelCase			      =	encoder_layers
         _lowerCAmelCase			      =	attention_window

         super().__init__(
             pad_token_id=_lowerCAmelCase							,  bos_token_id=_lowerCAmelCase							,  eos_token_id=_lowerCAmelCase							,  is_encoder_decoder=_lowerCAmelCase							,  decoder_start_token_id=_lowerCAmelCase							,  **_lowerCAmelCase							,  )