



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   SCREAMING_SNAKE_CASE_       (  _snake_case  :Tuple=False							)     ->    Tuple:
   # 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)
   _A						     =					min(2  ,      get_gpu_count()							) if distributed else 1
   _A						     =					os.environ.get('''DS_TEST_PORT'''  ,      _snake_case							)
   return F'''deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}'''.split()

def   SCREAMING_SNAKE_CASE_       (  )     ->    Optional[Any]:
   _A						     =					F'''{FIXTURE_DIRECTORY}/tests_samples'''
   _A						     =					F'''{data_dir_samples}/wmt_en_ro'''
   _A						     =					F'''{data_dir_samples}/xsum'''
   _A						     =					'''
        --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.
   _A						     =					{
       '''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''',
       ],
   }

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

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

   _A						     =					get_launcher(distributed=_snake_case							)

   _A						     =					{}
   for task, args in tasks.items():
      _A						     =					args.split()
      for model in tasksamodels[task]:
         _A						     =					globals()[F'''{model.upper().replace('-'  ,      '_'							)}_TINY''']
         _A						     =					F'''--model_name_or_path {model_name}'''.split()
         _A						     =					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   SCREAMING_SNAKE_CASE_       (  _snake_case  :str  ,      _snake_case  :Union[str, Any]  ,      _snake_case  :Optional[Any]							)     ->    List[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
   _A						     =					parameterized.to_safe_name('''_'''.join(str(_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   lowerCamelCase__  (					_A):



 """simple docstring"""

 def 						snake_case_			(						self				:				str ,	__lowerCAmelCase				:				int ,	__lowerCAmelCase				:				List[Any]			)      ->  List[Any]:
    # 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()}'''			)

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

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

    cmd += args_ds + args_out

    return cmd, output_dir

 @parameterized.expand(__lowerCAmelCase ,	name_func=__lowerCAmelCase			)
 def 						snake_case_			(						self				:				Any ,	__lowerCAmelCase				:				List[Any] ,	__lowerCAmelCase				:				Union[str, Any]			)      ->  List[str]:
    # testing the ability to do a run followed by recovery of full fp32 weights

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








import argparse
import logging
import os
import sys
import time

import tensorflow as tf
from datasets import load_dataset
from tqdm import tqdm

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
from transformers.utils import is_sagemaker_dp_enabled


if os.environ.get("""SDP_ENABLED""") or is_sagemaker_dp_enabled():
   UpperCAmelCase_		   =				True
   UpperCAmelCase_		   =				"""p3dn.24xlarge"""
   import smdistributed.dataparallel.tensorflow as sdp
else:
   UpperCAmelCase_		   =				False

def   SCREAMING_SNAKE_CASE_       (  _snake_case  :Optional[Any]  ,      _snake_case  :Union[str, Any]  ,      _snake_case  :List[str]  ,      _snake_case  :int  ,      _snake_case  :Any  ,      _snake_case  :List[Any]  ,      _snake_case  :Optional[int]=None							)     ->    Union[str, Any]:
   _A						     =					tqdm(_snake_case							)
   for i, batch in enumerate(_snake_case							):
      with tf.GradientTape() as tape:
         _A     ,						_A						     =					batch
         _A						     =					model(_snake_case							)
         _A						     =					loss(_snake_case  ,      outputs.logits							)

      if SDP_ENABLED:
         _A						     =					sdp.DistributedGradientTape(_snake_case  ,      sparse_as_dense=_snake_case							)

      _A						     =					tape.gradient(_snake_case  ,      model.trainable_variables							)
      opt.apply_gradients(zip(_snake_case  ,      model.trainable_variables							)							)

      pbar.set_description(F'''Loss: {loss_value:.4f}'''							)

      if SDP_ENABLED and i == 0:
         sdp.broadcast_variables(model.variables  ,      root_rank=0							)
         sdp.broadcast_variables(opt.variables()  ,      root_rank=0							)

      if max_steps and i >= max_steps:
         break

   _A						     =					{'''loss''': loss_value.numpy()}
   return train_results

def   SCREAMING_SNAKE_CASE_       (  _snake_case  :List[Any]  ,      _snake_case  :Optional[Any]  ,      _snake_case  :List[str]							)     ->    Dict:
   # Load dataset
   _A     ,						_A						     =					load_dataset('''imdb'''  ,      split=['''train''', '''test''']							)

   # Preprocess train dataset
   _A						     =					train_dataset.map(
       lambda _snake_case							: tokenizer(e['''text''']  ,      truncation=_snake_case  ,      padding='''max_length'''							)  ,      batched=_snake_case							)
   train_dataset.set_format(type='''tensorflow'''  ,      columns=['''input_ids''', '''attention_mask''', '''label''']							)

   _A						     =					{
       x: train_dataset[x].to_tensor(default_value=0  ,      shape=[None, tokenizer.model_max_length]							)
       for x in ['''input_ids''', '''attention_mask''']
   }
   _A						     =					tf.data.Dataset.from_tensor_slices((train_features, train_dataset['''label'''])							)

   # Preprocess test dataset
   _A						     =					test_dataset.map(
       lambda _snake_case							: tokenizer(e['''text''']  ,      truncation=_snake_case  ,      padding='''max_length'''							)  ,      batched=_snake_case							)
   test_dataset.set_format(type='''tensorflow'''  ,      columns=['''input_ids''', '''attention_mask''', '''label''']							)

   _A						     =					{
       x: test_dataset[x].to_tensor(default_value=0  ,      shape=[None, tokenizer.model_max_length]							)
       for x in ['''input_ids''', '''attention_mask''']
   }
   _A						     =					tf.data.Dataset.from_tensor_slices((test_features, test_dataset['''label'''])							)

   if SDP_ENABLED:
      _A						     =					tf_train_dataset.shard(sdp.size()  ,      sdp.rank()							)
      _A						     =					tf_test_dataset.shard(sdp.size()  ,      sdp.rank()							)
   _A						     =					tf_train_dataset.batch(_snake_case  ,      drop_remainder=_snake_case							)
   _A						     =					tf_test_dataset.batch(_snake_case  ,      drop_remainder=_snake_case							)

   return tf_train_dataset, tf_test_dataset


if __name__ == "__main__":
   UpperCAmelCase_		   =				argparse.ArgumentParser()

   # Hyperparameters sent by the client are passed as command-line arguments to the script.
   parser.add_argument("""--epochs""", type=int, default=3)
   parser.add_argument("""--per_device_train_batch_size""", type=int, default=1_6)
   parser.add_argument("""--per_device_eval_batch_size""", type=int, default=8)
   parser.add_argument("""--model_name_or_path""", type=str)
   parser.add_argument("""--learning_rate""", type=str, default=5E-5)
   parser.add_argument("""--do_train""", type=bool, default=True)
   parser.add_argument("""--do_eval""", type=bool, default=True)
   parser.add_argument("""--output_dir""", type=str)
   parser.add_argument("""--max_steps""", type=int, default=None)

   # Data, model, and output directories
   parser.add_argument("""--output_data_dir""", type=str, default=os.environ["""SM_OUTPUT_DATA_DIR"""])
   parser.add_argument("""--model_dir""", type=str, default=os.environ["""SM_MODEL_DIR"""])
   parser.add_argument("""--n_gpus""", type=str, default=os.environ["""SM_NUM_GPUS"""])

   UpperCAmelCase_   ,UpperCAmelCase_		   =				parser.parse_known_args()

   # Set up logging
   UpperCAmelCase_		   =				logging.getLogger(__name__)

   logging.basicConfig(
       level=logging.getLevelName("""INFO"""),
       handlers=[logging.StreamHandler(sys.stdout)],
       format="""%(asctime)s - %(name)s - %(levelname)s - %(message)s""",
   )

   if SDP_ENABLED:
      sdp.init()

      UpperCAmelCase_		   =				tf.config.experimental.list_physical_devices("""GPU""")
      for gpu in gpus:
         tf.config.experimental.set_memory_growth(gpu, True)
      if gpus:
         tf.config.experimental.set_visible_devices(gpus[sdp.local_rank()], """GPU""")

    # Load model and tokenizer
   UpperCAmelCase_		   =				TFAutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
   UpperCAmelCase_		   =				AutoTokenizer.from_pretrained(args.model_name_or_path)

   # get datasets
   UpperCAmelCase_   ,UpperCAmelCase_		   =				get_datasets(
       tokenizer=tokenizer,
       train_batch_size=args.per_device_train_batch_size,
       eval_batch_size=args.per_device_eval_batch_size,
   )

   # fine optimizer and loss
   UpperCAmelCase_		   =				tf.keras.optimizers.Adam(learning_rate=args.learning_rate)
   UpperCAmelCase_		   =				tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
   UpperCAmelCase_		   =				[tf.keras.metrics.SparseCategoricalAccuracy()]
   model.compile(optimizer=optimizer, loss=loss, metrics=metrics)

   # Training
   if args.do_train:
      # train_results = model.fit(tf_train_dataset, epochs=args.epochs, batch_size=args.train_batch_size)
      UpperCAmelCase_		   =				time.time()
      UpperCAmelCase_		   =				fit(
          model,
          loss,
          optimizer,
          tf_train_dataset,
          args.epochs,
          args.per_device_train_batch_size,
          max_steps=args.max_steps,
      )
      UpperCAmelCase_		   =				time.time() - start_train_time
      logger.info("""*** Train ***""")
      logger.info(f'train_runtime = {end_train_time}')

      UpperCAmelCase_		   =				os.path.join(args.output_dir, """train_results.txt""")

      if not SDP_ENABLED or sdp.rank() == 0:
         with open(output_eval_file, """w""") as writer:
            logger.info("""***** Train results *****""")
            logger.info(train_results)
            for key, value in train_results.items():
               logger.info(f'  {key} = {value}')
               writer.write(f'{key} = {value}\n')

    # Evaluation
   if args.do_eval and (not SDP_ENABLED or sdp.rank() == 0):
      UpperCAmelCase_		   =				model.evaluate(tf_test_dataset, batch_size=args.per_device_eval_batch_size, return_dict=True)
      logger.info("""*** Evaluate ***""")

      UpperCAmelCase_		   =				os.path.join(args.output_dir, """eval_results.txt""")

      with open(output_eval_file, """w""") as writer:
         logger.info("""***** Eval results *****""")
         logger.info(result)
         for key, value in result.items():
            logger.info(f'  {key} = {value}')
            writer.write(f'{key} = {value}\n')

    # Save result
   if SDP_ENABLED:
      if sdp.rank() == 0:
         model.save_pretrained(args.output_dir)
         tokenizer.save_pretrained(args.output_dir)
   else:
      model.save_pretrained(args.output_dir)
      tokenizer.save_pretrained(args.output_dir)



