






"""simple docstring"""




import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


_snake_case          =   os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(F'''{bindir}/../../examples/pytorch/translation'''):
 from run_translation import main  # noqa


set_seed(42)
_snake_case          =   'sshleifer/student_marian_en_ro_6_1'
_snake_case          =   'sshleifer/tiny-mbart'






@require_torch
class        UpperCamelCase			(	__snake_case      ):


 def _lowercase	(     self		:						int							,						UpperCAmelCase__		:						List[Any]=False							,						UpperCAmelCase__		:						List[Any]=None							,						UpperCAmelCase__		:						Tuple=True							,						UpperCAmelCase__		:						List[str]=True							,						UpperCAmelCase__		:						List[Any]=True							,						UpperCAmelCase__		:						Optional[int]=True							,						)					->       str:
   _a  :       str             =	self.run_trainer(
       eval_steps=1							,						max_len=12							,						model_name=UpperCAmelCase__							,						num_train_epochs=1							,						distributed=UpperCAmelCase__							,						extra_args_str=UpperCAmelCase__							,						predict_with_generate=UpperCAmelCase__							,						do_train=UpperCAmelCase__							,						do_eval=UpperCAmelCase__							,						do_predict=UpperCAmelCase__							,						)
   _a  :       List[Any]             =	TrainerState.load_from_json(os.path.join(UpperCAmelCase__							,						"""trainer_state.json"""							)							).log_history

   if not do_eval:
     return

   _a  :       Dict             =	[log for log in logs if '''eval_loss''' in log.keys()]

   _a  :       List[str]             =	eval_metrics[0]
   if predict_with_generate:
     assert "eval_bleu" in first_step_stats

     _a  :       Tuple             =	eval_metrics[-1]
     assert isinstance(last_step_stats["""eval_bleu"""]							,						UpperCAmelCase__							)
     assert not math.isnan(float(last_step_stats["""eval_loss"""]							)							), "eval_loss must not be `nan`"


 @require_torch_non_multi_gpu
 def _lowercase	(     self		:						Tuple							)					->       Union[str, Any]:
   self.run_seqaseq_quick()


 @require_torch_multi_gpu
 def _lowercase	(     self		:						List[Any]							)					->       List[str]:
   self.run_seqaseq_quick(distributed=UpperCAmelCase__							)


 @require_torch_multi_gpu
 def _lowercase	(     self		:						int							)					->       List[Any]:
   self.run_seqaseq_quick(distributed=UpperCAmelCase__							)


 @unittest.skip("""Requires an update of the env running those tests"""							)
 @require_torch_multi_gpu
 @require_fairscale
 def _lowercase	(     self		:						Any							)					->       Optional[Any]:
   self.run_seqaseq_quick(distributed=UpperCAmelCase__							,						extra_args_str="""--sharded_ddp simple"""							)


 @unittest.skip("""Requires an update of the env running those tests"""							)
 @require_torch_multi_gpu
 @require_fairscale
 def _lowercase	(     self		:						str							)					->       Union[str, Any]:
   self.run_seqaseq_quick(distributed=UpperCAmelCase__							,						extra_args_str="""--sharded_ddp simple --fp16"""							)


 @unittest.skip("""Requires an update of the env running those tests"""							)
 @require_torch_multi_gpu
 @require_fairscale
 def _lowercase	(     self		:						int							)					->       int:
   self.run_seqaseq_quick(distributed=UpperCAmelCase__							,						extra_args_str="""--sharded_ddp zero_dp_2"""							,						predict_with_generate=UpperCAmelCase__							)


 @unittest.skip("""Requires an update of the env running those tests"""							)
 @require_torch_multi_gpu
 @require_fairscale
 def _lowercase	(     self		:						List[str]							)					->       Union[str, Any]:
   self.run_seqaseq_quick(
       distributed=UpperCAmelCase__							,						extra_args_str="""--sharded_ddp zero_dp_2 --fp16"""							,						predict_with_generate=UpperCAmelCase__							)


 @require_apex
 @require_torch_gpu
 def _lowercase	(     self		:						Optional[Any]							)					->       int:
   self.run_seqaseq_quick(distributed=UpperCAmelCase__							,						extra_args_str="""--fp16 --fp16_backend=apex"""							)
   # test 2nd time - was getting eval_loss': nan'
   # to reproduce the problem set distributed=False
   self.run_seqaseq_quick(distributed=UpperCAmelCase__							,						extra_args_str="""--fp16 --fp16_backend=apex"""							)


 @parameterized.expand(["""base""", """low""", """high""", """mixed"""]							)
 @require_torch_multi_gpu
 def _lowercase	(     self		:						Optional[Any]							,						UpperCAmelCase__		:						Optional[int]							)					->       str:
   _a  :       Any             =	{
       # test with the default log_level - should be info and thus log info once
       '''base''': {'''extra_args_str''': '''''', '''n_matches''': 1},
       # test with low log_level and log_level_replica - should be noisy on all processes
       # now the info string should appear twice on 2 processes
       '''low''': {'''extra_args_str''': '''--log_level debug --log_level_replica debug''', '''n_matches''': 2},
       # test with high log_level and low log_level_replica
       # now the info string should appear once only on the replica
       '''high''': {'''extra_args_str''': '''--log_level error --log_level_replica debug''', '''n_matches''': 1},
       # test with high log_level and log_level_replica - should be quiet on all processes
       '''mixed''': {'''extra_args_str''': '''--log_level error --log_level_replica error''', '''n_matches''': 0},
   }

   _a  :       Any             =	experiments[experiment_id]
   _a  :       Any             =	{'''distributed''': True, '''predict_with_generate''': False, '''do_eval''': False, '''do_predict''': False}
   _a  :       Union[str, Any]             =	'''Running training'''
   with CaptureStderr() as cl:
     self.run_seqaseq_quick(**UpperCAmelCase__							,						extra_args_str=data["""extra_args_str"""]							)
   _a  :       int             =	len(re.findall(UpperCAmelCase__							,						cl.err							)							)
   self.assertEqual(UpperCAmelCase__							,						data["""n_matches"""]							)


 @slow
 def _lowercase	(     self		:						Optional[int]							)					->       Optional[Any]:
   _a  :       Tuple             =	self.run_trainer(
       eval_steps=2							,						max_len=128							,						model_name=UpperCAmelCase__							,						learning_rate=3E-4							,						num_train_epochs=10							,						distributed=UpperCAmelCase__							,						)

   # Check metrics
   _a  :       Union[str, Any]             =	TrainerState.load_from_json(os.path.join(UpperCAmelCase__							,						"""trainer_state.json"""							)							).log_history
   _a  :       Union[str, Any]             =	[log for log in logs if '''eval_loss''' in log.keys()]
   _a  :       List[str]             =	eval_metrics[0]
   _a  :       List[Any]             =	eval_metrics[-1]

   assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
   assert isinstance(last_step_stats["""eval_bleu"""]							,						UpperCAmelCase__							)

   # test if do_predict saves generations and metrics
   _a  :       int             =	os.listdir(UpperCAmelCase__							)
   _a  :       Optional[int]             =	{os.path.basename(UpperCAmelCase__							) for p in contents}
   assert "generated_predictions.txt" in contents
   assert "predict_results.json" in contents


 @slow
 @require_bitsandbytes
 def _lowercase	(     self		:						Any							)					->       Union[str, Any]:
   from transformers.training_args import OptimizerNames

   def train_and_return_metrics(UpperCAmelCase__		:						Union[str, Any]							) -> Tuple[int, float]:
     _a  :       Optional[int]             =	'''--skip_memory_metrics 0'''

     _a  :       List[Any]             =	self.run_trainer(
         max_len=128							,						model_name=UpperCAmelCase__							,						learning_rate=3E-4							,						num_train_epochs=1							,						optim=UpperCAmelCase__							,						distributed=UpperCAmelCase__							,						extra_args_str=UpperCAmelCase__							,						do_eval=UpperCAmelCase__							,						do_predict=UpperCAmelCase__							,						n_gpus_to_use=1							,						)

     # Check metrics
     _a  :       str             =	TrainerState.load_from_json(Path(UpperCAmelCase__							,						"""trainer_state.json"""							)							).log_history
     _a  :       Union[str, Any]             =	int(logs[0]["""train_mem_gpu_peaked_delta"""] / 2**20							)
     _a  :       int             =	int(logs[0]["""train_mem_gpu_alloc_delta"""] / 2**20							)

     _a  :       List[Any]             =	logs[0]['''train_loss''']
     return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

   _a  :       List[Any]             =	train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value							)
   _a  :       Optional[Any]             =	train_and_return_metrics(OptimizerNames.ADAMW_BNB.value							)

   _a  :       Dict             =	gpu_alloc_mem_orig - gpu_alloc_mem_bnb

   _a  :       Union[str, Any]             =	gpu_peak_mem_orig + gpu_alloc_mem_orig
   _a  :       Optional[Any]             =	gpu_peak_mem_bnb + gpu_alloc_mem_bnb
   _a  :       str             =	gpu_total_mem_orig - gpu_total_mem_bnb

   # sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
   # doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
   # in 2 bytes and the diff in optim memory usage is derived as so:
   #
   # - normal 25*8=~200MB (8 bytes per param)
   # - bnb    25*2= ~50MB (2 bytes per param)
   #
   # Thus we should expect ~150MB total memory saved.
   #
   # Peak memory should be the same - the total should be different by about that same margin
   #
   # After leaving a small margin to accommodate for differences between gpus let's check
   # that we have at least 120MB in savings
   _a  :       List[str]             =	120

   # uncomment the following if this test starts failing - requires py38 for a new print feature
   # gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
   # print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
   # print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
   # print(f"{gpu_alloc_mem_diff=}MB")
   # print(f"{gpu_peak_mem_diff=}MB")
   # print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
   # print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

   self.assertGreater(
       UpperCAmelCase__							,						UpperCAmelCase__							,						"""should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got"""
       f""" a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and"""
       f""" gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB"""							,						)

   self.assertGreater(
       UpperCAmelCase__							,						UpperCAmelCase__							,						"""should use ~150MB less total gpu memory with BNB, compared to without it for this model but got"""
       f""" a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and"""
       f""" gpu_total_mem_bnb={gpu_total_mem_bnb}MB"""							,						)

   self.assertEqual(
       UpperCAmelCase__							,						UpperCAmelCase__							,						f"""loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"""							)







 def _lowercase	(     self		:						List[Any]							,						UpperCAmelCase__		:						Optional[int]							,						UpperCAmelCase__		:						Any							,						UpperCAmelCase__		:						List[str]							,						UpperCAmelCase__		:						Union[str, Any] = 3E-3							,						UpperCAmelCase__		:						Optional[Any] = "adafactor"							,						UpperCAmelCase__		:						Optional[Any] = False							,						UpperCAmelCase__		:						str = None							,						UpperCAmelCase__		:						Optional[Any] = 0							,						UpperCAmelCase__		:						str = True							,						UpperCAmelCase__		:						Any = True							,						UpperCAmelCase__		:						Union[str, Any] = True							,						UpperCAmelCase__		:						Optional[int] = True							,						UpperCAmelCase__		:						Any = None							,						)					->       Tuple:
   _a  :       Tuple             =	self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
   _a  :       Dict             =	self.get_auto_remove_tmp_dir()
   _a  :       int             =	f"""\n            --model_name_or_path {model_name}\n            --train_file {data_dir}/train.json\n            --validation_file {data_dir}/val.json\n            --test_file {data_dir}/test.json\n            --output_dir {output_dir}\n            --overwrite_output_dir\n            --max_train_samples 8\n            --max_source_length {max_len}\n            --max_target_length {max_len}\n            --do_train\n            --num_train_epochs {str(UpperCAmelCase__							)}\n            --per_device_train_batch_size 4\n            --learning_rate {learning_rate}\n            --warmup_steps 8\n            --logging_steps 0\n            --logging_strategy no\n            --save_steps {str(UpperCAmelCase__							)}\n            --group_by_length\n            --label_smoothing_factor 0.1\n            --target_lang ro_RO\n            --source_lang en_XX\n        """.split()

   _a  :       Any             =	f"""\n            --do_eval\n            --per_device_eval_batch_size 4\n            --max_eval_samples 8\n            --val_max_target_length {max_len}\n            --evaluation_strategy steps\n            --eval_steps {str(UpperCAmelCase__							)}\n        """.split()

   _a  :       Optional[Any]             =	'''
            --do_predict
        '''.split()

   _a  :       Union[str, Any]             =	[]
   if do_train:
     args += args_train

   if do_eval:
     args += args_eval

   if do_predict:
     args += args_predict

   if predict_with_generate:
     args += "--predict_with_generate".split()

   if do_train:
     if optim == "adafactor":
       args += "--adafactor".split()
     else:
       args += f"""--optim {optim}""".split()

   if extra_args_str is not None:
     args += extra_args_str.split()

   if distributed:
     if n_gpus_to_use is None:
       _a  :       Dict             =	get_gpu_count()
     _a  :       Tuple             =	get_torch_dist_unique_port()
     _a  :       str             =	f"""\n                -m torch.distributed.run\n                --nproc_per_node={n_gpus_to_use}\n                --master_port={master_port}\n                {self.examples_dir_str}/pytorch/translation/run_translation.py\n            """.split()
     _a  :       str             =	[sys.executable] + distributed_args + args
     # keep for quick debug
     # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
     execute_subprocess_async(UpperCAmelCase__							,						env=self.get_env()							)
   else:
     _a  :       List[str]             =	['''run_translation.py'''] + args
     with patch.object(UpperCAmelCase__							,						"""argv"""							,						UpperCAmelCase__							):
       main()

   return output_dir



'''simple docstring'''




import os
import sys
import unittest


lowercase			:				Dict    		=			os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, 'utils'))

import get_test_info  # noqa: E402
from get_test_info import (  # noqa: E402
    get_model_to_test_mapping,
    get_model_to_tester_mapping,
    get_test_to_tester_mapping,
)


lowercase			:				Any    		=			os.path.join('tests', 'models', 'bert', 'test_modeling_bert.py')
lowercase			:				Optional[int]    		=			os.path.join('tests', 'models', 'blip', 'test_modeling_blip.py')





class  A				(							unittest.TestCase						):





     def 				__lowerCAmelCase							(							self	) ->				List[str]:






         """simple docstring"""






         A       :   Tuple					=						get_test_to_tester_mapping(SCREAMING_SNAKE_CASE	)
         A       :   Any					=						get_test_to_tester_mapping(SCREAMING_SNAKE_CASE	)

         A       :   List[Any]					=						{'''BertModelTest''': '''BertModelTester'''}

         A       :   int					=						{
             '''BlipModelTest''': '''BlipModelTester''',
             '''BlipTextImageModelTest''': '''BlipTextImageModelsModelTester''',
             '''BlipTextModelTest''': '''BlipTextModelTester''',
             '''BlipTextRetrievalModelTest''': '''BlipTextRetrievalModelTester''',
             '''BlipVQAModelTest''': '''BlipVQAModelTester''',
             '''BlipVisionModelTest''': '''BlipVisionModelTester''',
         }

         self.assertEqual(get_test_info.to_json(SCREAMING_SNAKE_CASE	)						,			SCREAMING_SNAKE_CASE	)
         self.assertEqual(get_test_info.to_json(SCREAMING_SNAKE_CASE	)						,			SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         A       :   Tuple					=						get_model_to_test_mapping(SCREAMING_SNAKE_CASE	)
         A       :   Tuple					=						get_model_to_test_mapping(SCREAMING_SNAKE_CASE	)

         A       :   List[str]					=						{
             '''BertForMaskedLM''': ['''BertModelTest'''],
             '''BertForMultipleChoice''': ['''BertModelTest'''],
             '''BertForNextSentencePrediction''': ['''BertModelTest'''],
             '''BertForPreTraining''': ['''BertModelTest'''],
             '''BertForQuestionAnswering''': ['''BertModelTest'''],
             '''BertForSequenceClassification''': ['''BertModelTest'''],
             '''BertForTokenClassification''': ['''BertModelTest'''],
             '''BertLMHeadModel''': ['''BertModelTest'''],
             '''BertModel''': ['''BertModelTest'''],
         }

         A       :   Union[str, Any]					=						{
             '''BlipForConditionalGeneration''': ['''BlipTextImageModelTest'''],
             '''BlipForImageTextRetrieval''': ['''BlipTextRetrievalModelTest'''],
             '''BlipForQuestionAnswering''': ['''BlipVQAModelTest'''],
             '''BlipModel''': ['''BlipModelTest'''],
             '''BlipTextModel''': ['''BlipTextModelTest'''],
             '''BlipVisionModel''': ['''BlipVisionModelTest'''],
         }

         self.assertEqual(get_test_info.to_json(SCREAMING_SNAKE_CASE	)						,			SCREAMING_SNAKE_CASE	)
         self.assertEqual(get_test_info.to_json(SCREAMING_SNAKE_CASE	)						,			SCREAMING_SNAKE_CASE	)





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






         """simple docstring"""






         A       :   int					=						get_model_to_tester_mapping(SCREAMING_SNAKE_CASE	)
         A       :   Union[str, Any]					=						get_model_to_tester_mapping(SCREAMING_SNAKE_CASE	)

         A       :   Dict					=						{
             '''BertForMaskedLM''': ['''BertModelTester'''],
             '''BertForMultipleChoice''': ['''BertModelTester'''],
             '''BertForNextSentencePrediction''': ['''BertModelTester'''],
             '''BertForPreTraining''': ['''BertModelTester'''],
             '''BertForQuestionAnswering''': ['''BertModelTester'''],
             '''BertForSequenceClassification''': ['''BertModelTester'''],
             '''BertForTokenClassification''': ['''BertModelTester'''],
             '''BertLMHeadModel''': ['''BertModelTester'''],
             '''BertModel''': ['''BertModelTester'''],
         }

         A       :   str					=						{
             '''BlipForConditionalGeneration''': ['''BlipTextImageModelsModelTester'''],
             '''BlipForImageTextRetrieval''': ['''BlipTextRetrievalModelTester'''],
             '''BlipForQuestionAnswering''': ['''BlipVQAModelTester'''],
             '''BlipModel''': ['''BlipModelTester'''],
             '''BlipTextModel''': ['''BlipTextModelTester'''],
             '''BlipVisionModel''': ['''BlipVisionModelTester'''],
         }

         self.assertEqual(get_test_info.to_json(SCREAMING_SNAKE_CASE	)						,			SCREAMING_SNAKE_CASE	)
         self.assertEqual(get_test_info.to_json(SCREAMING_SNAKE_CASE	)						,			SCREAMING_SNAKE_CASE	)




