

import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device


if is_torch_available():
  import torch

  from transformers import XLMProphetNetForConditionalGeneration, XLMProphetNetTokenizer






@require_torch
class 							A_				(	unittest.TestCase      ):


     @slow
     def 						UpperCAmelCase						(							self    :	List[str]      )							->			List[str]:
      __lowerCAmelCase:     int     			=   XLMProphetNetForConditionalGeneration.from_pretrained('microsoft/xprophetnet-large-wiki100-cased'      )
      model.to(UpperCAmelCase      )

      # encoder-decoder outputs
      __lowerCAmelCase:     Optional[int]     			=   torch.tensor([[1_7, 9_6_2_0_8, 1_0_3_4_7_1, 2]]      ).to(UpperCAmelCase      )
      __lowerCAmelCase:     List[Any]     			=   torch.tensor(
          [[2, 2_5_0, 9_9_5_3, 3_4, 6_9_4_8_9, 1_6_2_0, 3_2, 1_1_8_4_2_4, 6_2_4, 2_1_0, 1_0_5, 2_9_1_3, 1_0_3_2, 3_5_1]]      ).to(UpperCAmelCase      )
      __lowerCAmelCase:     Any     			=   model(
          input_ids=UpperCAmelCase    ,						attention_mask=UpperCAmelCase    ,						encoder_outputs=UpperCAmelCase    ,						decoder_input_ids=UpperCAmelCase      )
      __lowerCAmelCase:     List[Any]     			=   output[0]
      __lowerCAmelCase:     Any     			=   torch.Size((1, 1_4, 2_5_0_0_1_2)      )
      self.assertEqual(output_predited_logis.shape    ,						UpperCAmelCase      )
      __lowerCAmelCase:     str     			=   torch.tensor(
          [[[-6.3986, -8.2391, 12.5189], [-6.3289, -8.0864, 12.6211], [-6.2418, -8.0445, 12.7968]]]      ).to(UpperCAmelCase      )
      self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3]    ,						UpperCAmelCase    ,						atol=1E-4      )      )

      # encoder outputs
      __lowerCAmelCase:     Optional[Any]     			=   model.prophetnet.encoder(UpperCAmelCase      )[0]
      __lowerCAmelCase:     int     			=   torch.tensor(
          [[[-1.4260, -0.7628, 0.8453], [-1.4719, -0.1391, 0.7807], [-1.7678, 0.0114, 0.4646]]]      ).to(UpperCAmelCase      )
      __lowerCAmelCase:     str     			=   torch.Size((1, 4, 1_0_2_4)      )
      self.assertEqual(encoder_outputs.shape    ,						UpperCAmelCase      )
      self.assertTrue(torch.allclose(encoder_outputs[:, :3, :3]    ,						UpperCAmelCase    ,						atol=1E-4      )      )

      # decoder outputs
      __lowerCAmelCase:     Optional[int]     			=   model.prophetnet.decoder(
          UpperCAmelCase    ,						encoder_hidden_states=UpperCAmelCase    ,						)
      __lowerCAmelCase:     Tuple     			=   decoder_outputs[1].view(1    ,						model.config.ngram    ,						1_4    ,						-1      )
      __lowerCAmelCase:     Union[str, Any]     			=   model.lm_head(UpperCAmelCase      )
      __lowerCAmelCase:     Union[str, Any]     			=   predicting_streams_logits[:, 0]
      self.assertTrue(torch.allclose(next_first_stream_logits[:, :3, :3]    ,						UpperCAmelCase    ,						atol=1E-4      )      )


     @slow
     def 						UpperCAmelCase						(							self    :	Optional[Any]      )							->			Optional[Any]:
      __lowerCAmelCase:     Tuple     			=   XLMProphetNetForConditionalGeneration.from_pretrained(
          'microsoft/xprophetnet-large-wiki100-cased-xglue-ntg'      )
      model.to(UpperCAmelCase      )

      __lowerCAmelCase:     Dict     			=   torch.tensor([[1_7, 9_6_2_0_8, 1_0_3_4_7_1, 2]]      ).to(UpperCAmelCase      )
      __lowerCAmelCase:     Optional[Any]     			=   torch.tensor(
          [[2, 2_5_0, 9_9_5_3, 3_4, 6_9_4_8_9, 1_6_2_0, 3_2, 1_1_8_4_2_4, 6_2_4, 2_1_0, 1_0_5, 2_9_1_3, 1_0_3_2, 3_5_1]]      ).to(UpperCAmelCase      )
      __lowerCAmelCase:     List[Any]     			=   model(
          input_ids=UpperCAmelCase    ,						attention_mask=UpperCAmelCase    ,						encoder_outputs=UpperCAmelCase    ,						decoder_input_ids=UpperCAmelCase      )
      __lowerCAmelCase:     str     			=   output[0]
      __lowerCAmelCase:     Optional[int]     			=   torch.Size((1, 1_4, 2_5_0_0_1_2)      )
      self.assertEqual(output_predited_logis.shape    ,						UpperCAmelCase      )
      # compare the actual values for a slice.
      __lowerCAmelCase:     Optional[Any]     			=   torch.tensor(
          [[[-9.2253, -9.7173, -6.3529], [-7.6701, -9.0145, -1.9382], [-8.0195, -7.0004, -0.1523]]]      ).to(UpperCAmelCase      )

      self.assertTrue(torch.allclose(output_predited_logis[:, :3, :3]    ,						UpperCAmelCase    ,						atol=1E-4      )      )


     @slow
     def 						UpperCAmelCase						(							self    :	Optional[Any]      )							->			Dict:
      __lowerCAmelCase:     Optional[int]     			=   XLMProphetNetForConditionalGeneration.from_pretrained(
          'microsoft/xprophetnet-large-wiki100-cased-xglue-ntg'      )
      model.to(UpperCAmelCase      )
      __lowerCAmelCase:     str     			=   5_1_2

      __lowerCAmelCase:     List[Any]     			=   XLMProphetNetTokenizer.from_pretrained('microsoft/xprophetnet-large-wiki100-cased-xglue-ntg'      )

      __lowerCAmelCase:     Optional[Any]     			=   (
          'Microsoft Corporation intends to officially end free support for the Windows 7 operating system after'
          ' January 14, 2020, according to the official portal of the organization. From that day, users of this'
          ' system will not be able to receive security updates, which could make their computers vulnerable to'
          ' cyber attacks.'
      )
      __lowerCAmelCase:     str     			=   (
          'орпорация Microsoft намерена официально прекратить бесплатную поддержку операционной системы Windows 7'
          ' после 14 января 2020 года, сообщается на официальном портале организации . С указанного дня пользователи'
          ' этой системы не смогут получать обновления безопасности, из-за чего их компьютеры могут стать уязвимыми'
          ' к кибератакам.'
      )
      __lowerCAmelCase:     List[str]     			=   (
          '根据该组织的官方门户网站，微软公司打算在2020年1月14日之后正式终止对Windows 7操作系统的免费支持。从那时起，该系统的用户将无法接收安全更新，这可能会使他们的计算机容易受到网络攻击。'
      )

      __lowerCAmelCase:     List[Any]     			=   tokenizer(
          [EN_SENTENCE, RU_SENTENCE, ZH_SENTENCE]    ,						padding=UpperCAmelCase    ,						max_length=2_5_5    ,						return_tensors='pt'      ).input_ids
      __lowerCAmelCase:     Optional[int]     			=   input_ids.to(UpperCAmelCase      )

      __lowerCAmelCase:     Union[str, Any]     			=   model.generate(
          UpperCAmelCase    ,						num_beams=1_0    ,						length_penalty=1.0    ,						no_repeat_ngram_size=3    ,						early_stopping=UpperCAmelCase      )
      __lowerCAmelCase:     Tuple     			=   [tokenizer.decode(UpperCAmelCase    ,						skip_special_tokens=UpperCAmelCase      ) for g in summary_ids]
      __lowerCAmelCase:     int     			=   'Microsoft to end Windows 7 free support after January 14, 2020'
      __lowerCAmelCase:     Any     			=   'Microsoft намерена прекратить бесплатную поддержку Windows 7 после 14 января 2020 года'
      __lowerCAmelCase:     List[str]     			=   '微软打算终止对Windows 7操作系统的免费支持'
      self.assertListEqual(
          [EXPECTED_TITLE_EN, EXPECTED_TITLE_RU, EXPECTED_TITLE_ZH]    ,						UpperCAmelCase    ,						)

      __lowerCAmelCase:     List[Any]     			=   model.generate(
          UpperCAmelCase    ,						num_beams=1    ,						length_penalty=1.0    ,						no_repeat_ngram_size=3    ,						early_stopping=UpperCAmelCase      )
      __lowerCAmelCase:     int     			=   [
          tokenizer.convert_ids_to_tokens(UpperCAmelCase    ,						skip_special_tokens=UpperCAmelCase      ) for g in summary_ids_beama
      ]
      __lowerCAmelCase:     str     			=   '▁Microsoft ▁to ▁end ▁free ▁support ▁for ▁Windows ▁7'.split(' '      )
      __lowerCAmelCase:     Tuple     			=   '▁Microsoft ▁намерен а ▁прекрати ть ▁бес плат ную ▁поддержку ▁Windows ▁7 ▁после ▁14 ▁января ▁2020 ▁года'.split(
          ' '      )
      __lowerCAmelCase:     Tuple     			=   '微软 公司 打算 终止 对 Windows ▁7 操作 系统的 免费 支持'.split(' '      )
      self.assertListEqual(
          [EXPECTED_TITLE_EN_BEAM1_TOK, EXPECTED_TITLE_RU_BEAM1_TOK, EXPECTED_TITLE_ZH_BEAM1_TOK]    ,						UpperCAmelCase    ,						)








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)

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

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

_a	 =     '''t5-small'''

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


# *** To Fix ***


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

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

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

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

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







def _a				( SCREAMING_SNAKE_CASE	:    List[str]=False	)					->					Union[str, Any]:

 """simple docstring"""






 __lowerCAmelCase:     Union[str, Any]     			=   min(2					,      get_gpu_count()	) if distributed else 1
 __lowerCAmelCase:     Any     			=   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 _a				( )					->					Optional[Any]:

 """simple docstring"""






 __lowerCAmelCase:     int     			=   f'''{FIXTURE_DIRECTORY}/tests_samples'''
 __lowerCAmelCase:     List[str]     			=   f'''{data_dir_samples}/wmt_en_ro'''
 __lowerCAmelCase:     str     			=   f'''{data_dir_samples}/xsum'''
 __lowerCAmelCase:     int     			=   '\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.
 __lowerCAmelCase:     Dict     			=   {
     '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:     Optional[Any]     			=   f'''{ROOT_DIRECTORY}/examples/pytorch'''

 __lowerCAmelCase:     Optional[int]     			=   {
     '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:     List[Any]     			=   get_launcher(distributed=SCREAMING_SNAKE_CASE	)

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


_a	 =     make_task_cmds()

_a	 =     '''zero2'''
_a	 =     '''zero3'''

_a	 =     [ZEROa, ZEROa]







def _a				( SCREAMING_SNAKE_CASE	:    Any					,      SCREAMING_SNAKE_CASE	:    List[Any]					,      SCREAMING_SNAKE_CASE	:    Optional[int]	)					->					Dict:

 """simple docstring"""






 __lowerCAmelCase:     Any     			=   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
_a	 =     list(itertools.product(stages, task_cmds.keys()))






@slow
@require_deepspeed
@require_torch_gpu
class 							A_				(	snake_case__      ):


     def 						UpperCAmelCase						(							self    :	int    ,						UpperCAmelCase    :	Any    ,						UpperCAmelCase    :	Optional[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()}'''      )

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

      __lowerCAmelCase:     Dict     			=   self.get_auto_remove_tmp_dir()
      __lowerCAmelCase:     Optional[Any]     			=   F'''--output_dir {output_dir}'''.split()

      cmd += args_ds + args_out

      return cmd, output_dir


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

      __lowerCAmelCase			,							__lowerCAmelCase:     str     			=   self.get_task_cmd(UpperCAmelCase    ,						UpperCAmelCase      )

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

      # 2. test that the fp32 weights get reconsolidated
      __lowerCAmelCase:     Optional[Any]     			=   F'''{output_dir}/checkpoint-1'''
      __lowerCAmelCase:     Dict     			=   F'''{chkpt_dir}/out.bin'''
      __lowerCAmelCase:     Optional[Any]     			=   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(UpperCAmelCase    ,						shell=UpperCAmelCase      )
      assert os.path.exists(UpperCAmelCase      ), 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





