

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)

_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_(			snake_case__:  Optional[Any]=False   )       -> Optional[Any]:
 # 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__							=				min(2  ,    get_gpu_count()   ) if distributed else 1
 UpperCAmelCase__							=				os.environ.get('DS_TEST_PORT'  ,    snake_case__   )
 return f"deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}".split()
def 						UpperCamelCase_(			)       -> Tuple:
 UpperCAmelCase__							=				f"{FIXTURE_DIRECTORY}/tests_samples"
 UpperCAmelCase__							=				f"{data_dir_samples}/wmt_en_ro"
 UpperCAmelCase__							=				f"{data_dir_samples}/xsum"
 UpperCAmelCase__							=				'\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.
 UpperCAmelCase__							=				{
     '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__							=				f"{ROOT_DIRECTORY}/examples/pytorch"

 UpperCAmelCase__							=				{
     '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__							=				get_launcher(distributed=snake_case__   )

 UpperCAmelCase__							=				{}
 for task, args in tasks.items():
  UpperCAmelCase__							=				args.split()
  for model in tasksamodels[task]:
   UpperCAmelCase__							=				globals()[f"{model.upper().replace('-'  ,    '_'   )}_TINY"]
   UpperCAmelCase__							=				f"--model_name_or_path {model_name}".split()
   UpperCAmelCase__							=				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_(			snake_case__:  Tuple  ,    snake_case__:  Union[str, Any]  ,    snake_case__:  List[Any]   )       -> Dict:
 # 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__							=				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     lowercase						(       _UpperCamelCase					):



 '''simple docstring'''


 def   UpperCamelCase__							(self       ,       __a       ,       __a ) ->					List[str]:






  """simple docstring"""





  if task not in task_cmds:
   raise ValueError(F"don't know of task {task}, have {task_cmds.keys()}" )

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

  UpperCAmelCase__							=				self.get_auto_remove_tmp_dir()
  UpperCAmelCase__							=				F"--output_dir {output_dir}".split()

  cmd += args_ds + args_out

  return cmd, output_dir




 @parameterized.expand(__a       ,       name_func=__a )
 def   UpperCamelCase__							(self       ,       __a       ,       __a ) ->					List[str]:






  """simple docstring"""





  UpperCAmelCase__				,    UpperCAmelCase__							=				self.get_task_cmd(__a       ,       __a )

  # 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(__a       ,       env=self.get_env() )

  # 2. test that the fp32 weights get reconsolidated
  UpperCAmelCase__							=				F"{output_dir}/checkpoint-1"
  UpperCAmelCase__							=				F"{chkpt_dir}/out.bin"
  UpperCAmelCase__							=				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(__a       ,       shell=__a )
  assert os.path.exists(__a ), 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 unittest

from transformers import load_tool
from transformers.tools.agent_types import AGENT_TYPE_MAPPING

from .test_tools_common import ToolTesterMixin, output_types



class     lowercase						(       unittest.TestCase ,				_UpperCamelCase					):



 '''simple docstring'''


 def   UpperCamelCase__							(self ) ->					List[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				load_tool('translation' )
  self.tool.setup()
  UpperCAmelCase__							=				load_tool('translation'       ,       remote=__a )




 def   UpperCamelCase__							(self ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__							=				self.tool('Hey, what\'s up?'       ,       src_lang='English'       ,       tgt_lang='French' )
  self.assertEqual(__a       ,       '- Hé, comment ça va?' )




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__							=				self.remote_tool('Hey, what\'s up?'       ,       src_lang='English'       ,       tgt_lang='French' )
  self.assertEqual(__a       ,       '- Hé, comment ça va?' )




 def   UpperCamelCase__							(self ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__							=				self.tool(text='Hey, what\'s up?'       ,       src_lang='English'       ,       tgt_lang='French' )
  self.assertEqual(__a       ,       '- Hé, comment ça va?' )




 def   UpperCamelCase__							(self ) ->					Optional[Any]:






  """simple docstring"""





  UpperCAmelCase__							=				self.remote_tool(text='Hey, what\'s up?'       ,       src_lang='English'       ,       tgt_lang='French' )
  self.assertEqual(__a       ,       '- Hé, comment ça va?' )




 def   UpperCamelCase__							(self ) ->					Any:






  """simple docstring"""





  UpperCAmelCase__							=				['Hey, what\'s up?', 'English', 'Spanish']
  UpperCAmelCase__							=				self.tool(*__a )

  # There is a single output
  if len(self.tool.outputs ) == 1:
   UpperCAmelCase__							=				[outputs]

  self.assertListEqual(output_types(__a )       ,       self.tool.outputs )




 def   UpperCamelCase__							(self ) ->					Optional[int]:






  """simple docstring"""





  UpperCAmelCase__							=				['Hey, what\'s up?', 'English', 'Spanish']
  UpperCAmelCase__							=				self.tool(*__a )

  if not isinstance(__a       ,       __a ):
   UpperCAmelCase__							=				[outputs]

  self.assertEqual(len(__a )       ,       len(self.tool.outputs ) )

  for output, output_type in zip(__a       ,       self.tool.outputs ):
   UpperCAmelCase__							=				AGENT_TYPE_MAPPING[output_type]
   self.assertTrue(isinstance(__a       ,       __a ) )




 def   UpperCamelCase__							(self ) ->					Dict:






  """simple docstring"""





  UpperCAmelCase__							=				['Hey, what\'s up?', 'English', 'Spanish']

  UpperCAmelCase__							=				[]

  for _input, input_type in zip(__a       ,       self.tool.inputs ):
   if isinstance(__a       ,       __a ):
    _inputs.append([AGENT_TYPE_MAPPING[_input_type](_input ) for _input_type in input_type] )
   else:
    _inputs.append(AGENT_TYPE_MAPPING[input_type](_input ) )

        # Should not raise an error
  UpperCAmelCase__							=				self.tool(*__a )

  if not isinstance(__a       ,       __a ):
   UpperCAmelCase__							=				[outputs]

  self.assertEqual(len(__a )       ,       len(self.tool.outputs ) )



