





'''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


__lowerCAmelCase				:      Dict						      =os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(f"""{bindir}/../../examples/pytorch/translation"""):
  from run_translation import main  # noqa


set_seed(42)
__lowerCAmelCase				:      Optional[Any]						      ="sshleifer/student_marian_en_ro_6_1"
__lowerCAmelCase				:      List[str]						      ="sshleifer/tiny-mbart"






@require_torch
class   UpperCAmelCase      (			UpperCamelCase__							):
 def 	UpperCAmelCase_				(							self :Any		,			lowercase_ :int=False		,			lowercase_ :Any=None		,			lowercase_ :int=True		,			lowercase_ :Optional[int]=True		,			lowercase_ :Optional[Any]=True		,			lowercase_ :int=True		,			)->				Optional[Any]:
     A__							=     self.run_trainer(
         eval_steps=1		,			max_len=12		,			model_name=lowercase_		,			num_train_epochs=1		,			distributed=lowercase_		,			extra_args_str=lowercase_		,			predict_with_generate=lowercase_		,			do_train=lowercase_		,			do_eval=lowercase_		,			do_predict=lowercase_		,			)
     A__							=     TrainerState.load_from_json(os.path.join(lowercase_		,			"trainer_state.json"				)				).log_history

     if not do_eval:
         return

     A__							=     [log for log in logs if "eval_loss" in log.keys()]

     A__							=     eval_metrics[0]
     if predict_with_generate:
         assert "eval_bleu" in first_step_stats

         A__							=     eval_metrics[-1]
         assert isinstance(last_step_stats["eval_bleu"]		,			lowercase_				)
         assert not math.isnan(float(last_step_stats["eval_loss"]				)				), "eval_loss must not be `nan`"
 @require_torch_non_multi_gpu
 def 	UpperCAmelCase_				(							self :Any				)->				List[Any]:
     self.run_seqaseq_quick()
 @require_torch_multi_gpu
 def 	UpperCAmelCase_				(							self :str				)->				Optional[int]:
     self.run_seqaseq_quick(distributed=lowercase_				)
 @require_torch_multi_gpu
 def 	UpperCAmelCase_				(							self :Tuple				)->				Union[str, Any]:
     self.run_seqaseq_quick(distributed=lowercase_				)
 @unittest.skip("Requires an update of the env running those tests"				)
 @require_torch_multi_gpu
 @require_fairscale
 def 	UpperCAmelCase_				(							self :int				)->				Dict:
     self.run_seqaseq_quick(distributed=lowercase_		,			extra_args_str="--sharded_ddp simple"				)
 @unittest.skip("Requires an update of the env running those tests"				)
 @require_torch_multi_gpu
 @require_fairscale
 def 	UpperCAmelCase_				(							self :Tuple				)->				Tuple:
     self.run_seqaseq_quick(distributed=lowercase_		,			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 	UpperCAmelCase_				(							self :str				)->				List[str]:
     self.run_seqaseq_quick(distributed=lowercase_		,			extra_args_str="--sharded_ddp zero_dp_2"		,			predict_with_generate=lowercase_				)
 @unittest.skip("Requires an update of the env running those tests"				)
 @require_torch_multi_gpu
 @require_fairscale
 def 	UpperCAmelCase_				(							self :Optional[int]				)->				List[Any]:
     self.run_seqaseq_quick(
         distributed=lowercase_		,			extra_args_str="--sharded_ddp zero_dp_2 --fp16"		,			predict_with_generate=lowercase_				)
 @require_apex
 @require_torch_gpu
 def 	UpperCAmelCase_				(							self :Optional[int]				)->				List[str]:
     # XXX: apex breaks the trainer if it's run twice e.g. run_seq2seq.main() from the same
     # program and it breaks other tests that run from the same pytest worker, therefore until this is
     # sorted out it must be run only in an external program, that is distributed=True in this
     # test and only under one or more gpus - if we want cpu will need to make a special test
     #
     # specifically to the problem traced it to self.optimizer.step() - if it's run 2nd time via
     # 2nd main() call it botches the future eval.
     #
     self.run_seqaseq_quick(distributed=lowercase_		,			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=lowercase_		,			extra_args_str="--fp16 --fp16_backend=apex"				)
 @parameterized.expand(["base", "low", "high", "mixed"]				)
 @require_torch_multi_gpu
 def 	UpperCAmelCase_				(							self :Tuple		,			lowercase_ :Union[str, Any]				)->				str:
     # as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
     A__							=     {
         # 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__							=     experiments[experiment_id]
     A__							=     {"distributed": True, "predict_with_generate": False, "do_eval": False, "do_predict": False}
     A__							=     "Running training"
     with CaptureStderr() as cl:
         self.run_seqaseq_quick(**lowercase_		,			extra_args_str=data["extra_args_str"]				)
     A__							=     len(re.findall(lowercase_		,			cl.err				)				)
     self.assertEqual(lowercase_		,			data["n_matches"]				)
 @slow
 def 	UpperCAmelCase_				(							self :Tuple				)->				Tuple:
     A__							=     self.run_trainer(
         eval_steps=2		,			max_len=1_28		,			model_name=lowercase_		,			learning_rate=3E-4		,			num_train_epochs=10		,			distributed=lowercase_		,			)

     # Check metrics
     A__							=     TrainerState.load_from_json(os.path.join(lowercase_		,			"trainer_state.json"				)				).log_history
     A__							=     [log for log in logs if "eval_loss" in log.keys()]
     A__							=     eval_metrics[0]
     A__							=     eval_metrics[-1]

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

     # test if do_predict saves generations and metrics
     A__							=     os.listdir(lowercase_				)
     A__							=     {os.path.basename(lowercase_				) for p in contents}
     assert "generated_predictions.txt" in contents
     assert "predict_results.json" in contents
 @slow
 @require_bitsandbytes
 def 	UpperCAmelCase_				(							self :Any				)->				str:
     from transformers.training_args import OptimizerNames

     def train_and_return_metrics(lowercase_ :str				) -> Tuple[int, float]:
         A__							=     "--skip_memory_metrics 0"

         A__							=     self.run_trainer(
             max_len=1_28		,			model_name=lowercase_		,			learning_rate=3E-4		,			num_train_epochs=1		,			optim=lowercase_		,			distributed=lowercase_		,			extra_args_str=lowercase_		,			do_eval=lowercase_		,			do_predict=lowercase_		,			n_gpus_to_use=1		,			)

         # Check metrics
         A__							=     TrainerState.load_from_json(Path(lowercase_		,			"trainer_state.json"				)				).log_history
         A__							=     int(logs[0]["train_mem_gpu_peaked_delta"] / 2**20				)
         A__							=     int(logs[0]["train_mem_gpu_alloc_delta"] / 2**20				)

         A__							=     logs[0]["train_loss"]
         return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

     A__,      A__,      A__							=     train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value				)
     A__,      A__,      A__							=     train_and_return_metrics(OptimizerNames.ADAMW_BNB.value				)

     A__							=     gpu_alloc_mem_orig - gpu_alloc_mem_bnb

     A__							=     gpu_peak_mem_orig + gpu_alloc_mem_orig
     A__							=     gpu_peak_mem_bnb + gpu_alloc_mem_bnb
     A__							=     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__							=     1_20

     # 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(
         lowercase_		,			lowercase_		,			"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(
         lowercase_		,			lowercase_		,			"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(
         lowercase_		,			lowercase_		,			F"loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"				)



 def 	UpperCAmelCase_				(							self :Any		,			lowercase_ :int		,			lowercase_ :str		,			lowercase_ :int		,			lowercase_ :float = 3E-3		,			lowercase_ :str = "adafactor"		,			lowercase_ :bool = False		,			lowercase_ :str = None		,			lowercase_ :int = 0		,			lowercase_ :bool = True		,			lowercase_ :bool = True		,			lowercase_ :bool = True		,			lowercase_ :bool = True		,			lowercase_ :int = None		,			)->				Any:
     A__							=     self.test_file_dir / "../fixtures/tests_samples/wmt_en_ro"
     A__							=     self.get_auto_remove_tmp_dir()
     A__							=     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(lowercase_				)}\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(lowercase_				)}\n            --group_by_length\n            --label_smoothing_factor 0.1\n            --target_lang ro_RO\n            --source_lang en_XX\n        ".split()

     A__							=     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(lowercase_				)}\n        ".split()

     A__							=     "\n            --do_predict\n        ".split()

     A__							=     []
     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__							=     get_gpu_count()
         A__							=     get_torch_dist_unique_port()
         A__							=     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__							=     [sys.executable] + distributed_args + args
         # keep for quick debug
         # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
         execute_subprocess_async(lowercase_		,			env=self.get_env()				)
     else:
         A__							=     ["run_translation.py"] + args
         with patch.object(lowercase_		,			"argv"		,			lowercase_				):
             main()

     return output_dir










'''simple docstring'''

import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


__lowerCAmelCase				:      int						      =5_0000
__lowerCAmelCase				:      Union[str, Any]						      =5000

__lowerCAmelCase      ,	__lowerCAmelCase				:      Tuple						      =os.path.split(__file__)
__lowerCAmelCase				:      int						      =os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   int			):
    for i in range(_lowerCamelCase			):
        A__							=     dataset[i]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[int]		,			_lowerCamelCase   :   Any			):
    for i in range(0		,			len(_lowerCamelCase			)		,			_lowerCamelCase			):
        A__							=     dataset[i : i + batch_size]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[Any]		,			_lowerCamelCase   :   List[str]			):
    with dataset.formatted_as(type=_lowerCamelCase			):
        for i in range(_lowerCamelCase			):
            A__							=     dataset[i]





@get_duration
def 		UpperCamelCase							(				_lowerCamelCase   :   datasets.Dataset		,			_lowerCamelCase   :   Optional[int]		,			_lowerCamelCase   :   int		,			_lowerCamelCase   :   Dict			):
    with dataset.formatted_as(type=_lowerCamelCase			):
        for i in range(0		,			_lowerCamelCase		,			_lowerCamelCase			):
            A__							=     dataset[i : i + batch_size]





def 		UpperCamelCase							(				):
    A__							=     {"num examples": SPEED_TEST_N_EXAMPLES}
    A__							=     [
        (read, {"length": SMALL_TEST}),
        (read, {"length": SPEED_TEST_N_EXAMPLES}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1_00}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10_00}),
        (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
        (read_formatted, {"type": "pandas", "length": SMALL_TEST}),
        (read_formatted, {"type": "torch", "length": SMALL_TEST}),
        (read_formatted, {"type": "tensorflow", "length": SMALL_TEST}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10_00}),
    ]

    A__							=     [
        (read, {"length": SMALL_TEST}),
        (read, {"length": SPEED_TEST_N_EXAMPLES}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 1_00}),
        (read_batch, {"length": SPEED_TEST_N_EXAMPLES, "batch_size": 10_00}),
        (read_formatted, {"type": "numpy", "length": SMALL_TEST}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10}),
        (read_formatted_batch, {"type": "numpy", "length": SMALL_TEST, "batch_size": 10_00}),
    ]
    with tempfile.TemporaryDirectory() as tmp_dir:
        print("generating dataset"			)
        A__							=     datasets.Features(
            {"list": datasets.Sequence(datasets.Value("float32"			)			), "numbers": datasets.Value("float32"			)}			)
        A__							=     generate_example_dataset(
            os.path.join(_lowerCamelCase		,			"dataset.arrow"			)		,			_lowerCamelCase		,			num_examples=_lowerCamelCase		,			seq_shapes={"list": (1_00,)}		,			)
        print("first set of iterations"			)
        for func, kwargs in functions:
            print(func.__name__		,			str(_lowerCamelCase			)			)
            A__							=     func(_lowerCamelCase		,			**_lowerCamelCase			)

        print("shuffling dataset"			)
        A__							=     dataset.shuffle()
        print("Second set of iterations (after shuffling"			)
        for func, kwargs in functions_shuffled:
            print("shuffled "		,			func.__name__		,			str(_lowerCamelCase			)			)
            A__							=     func(
                _lowerCamelCase		,			**_lowerCamelCase			)

    with open(_lowerCamelCase		,			"wb"			) as f:
        f.write(json.dumps(_lowerCamelCase			).encode("utf-8"			)			)


if __name__ == "__main__":  # useful to run the profiler
  benchmark_iterating()



