



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


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


set_seed(42)
lowercase_						=       """sshleifer/student_marian_en_ro_6_1"""
lowercase_						=       """sshleifer/tiny-mbart"""
@require_torch
class 			_snake_case     (   lowercase__):




				def 							A__ (  self   :  Union[str, Any],   __lowercase   :  Optional[int]=False,   __lowercase   :  Union[str, Any]=None,   __lowercase   :  Tuple=True,   __lowercase   :  Union[str, Any]=True,   __lowercase   :  Union[str, Any]=True,   __lowercase   :  List[Any]=True,   ):
							lowercase__   					=			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,   )
							lowercase__   					=			TrainerState.load_from_json(os.path.join(__lowercase,   "trainer_state.json"  )  ).log_history

							if not do_eval:
										return

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

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

										lowercase__   					=			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 							A__ (  self   :  Any  ):
							self.run_seqaseq_quick()




				@require_torch_multi_gpu
				def 							A__ (  self   :  List[str]  ):
							self.run_seqaseq_quick(distributed=__lowercase  )




				@require_torch_multi_gpu
				def 							A__ (  self   :  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 							A__ (  self   :  Optional[int]  ):
							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 							A__ (  self   :  Union[str, Any]  ):
							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 							A__ (  self   :  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 							A__ (  self   :  Optional[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 							A__ (  self   :  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 							A__ (  self   :  Tuple,   __lowercase   :  Tuple  ):
							# as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
							lowercase__   					=			{
							    # 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},
							}

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




				@slow
				def 							A__ (  self   :  int  ):
							lowercase__   					=			self.run_trainer(
							    eval_steps=2,   max_len=128,   model_name=__lowercase,   learning_rate=3e-4,   num_train_epochs=10,   distributed=__lowercase,   )

							# Check metrics
							lowercase__   					=			TrainerState.load_from_json(os.path.join(__lowercase,   "trainer_state.json"  )  ).log_history
							lowercase__   					=			[log for log in logs if "eval_loss" in log.keys()]
							lowercase__   					=			eval_metrics[0]
							lowercase__   					=			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
							lowercase__   					=			os.listdir(__lowercase  )
							lowercase__   					=			{os.path.basename(__lowercase  ) for p in contents}
							assert "generated_predictions.txt" in contents
							assert "predict_results.json" in contents




				@slow
				@require_bitsandbytes
				def 							A__ (  self   :  List[str]  ):
							from transformers.training_args import OptimizerNames

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

										lowercase__   					=			self.run_trainer(
										    max_len=128,   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
										lowercase__   					=			TrainerState.load_from_json(Path(__lowercase,   "trainer_state.json"  )  ).log_history
										lowercase__   					=			int(logs[0]["train_mem_gpu_peaked_delta"] / 2**20  )
										lowercase__   					=			int(logs[0]["train_mem_gpu_alloc_delta"] / 2**20  )

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

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

							lowercase__   					=			gpu_alloc_mem_orig - gpu_alloc_mem_bnb

							lowercase__   					=			gpu_peak_mem_orig + gpu_alloc_mem_orig
							lowercase__   					=			gpu_peak_mem_bnb + gpu_alloc_mem_bnb
							lowercase__   					=			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
							lowercase__   					=			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(
							    __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 							A__ (  self   :  str,   __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,   ):
							lowercase__   					=			self.test_file_dir / "../fixtures/tests_samples/wmt_en_ro"
							lowercase__   					=			self.get_auto_remove_tmp_dir()
							lowercase__   					=			F'''
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(__lowercase  )}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(__lowercase  )}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        '''.split()

							lowercase__   					=			F'''
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(__lowercase  )}
        '''.split()

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

							lowercase__   					=			[]
							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:
													lowercase__   					=			get_gpu_count()
										lowercase__   					=			get_torch_dist_unique_port()
										lowercase__   					=			F'''
                -m torch.distributed.run
                --nproc_per_node={n_gpus_to_use}
                --master_port={master_port}
                {self.examples_dir_str}/pytorch/translation/run_translation.py
            '''.split()
										lowercase__   					=			[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:
										lowercase__   					=			["run_translation.py"] + args
										with patch.object(__lowercase,   "argv",   __lowercase  ):
													main()

							return output_dir









from collections import OrderedDict
from typing import Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...feature_extraction_utils import FeatureExtractionMixin
from ...onnx import OnnxConfig
from ...onnx.utils import compute_effective_axis_dimension
from ...tokenization_utils_base import PreTrainedTokenizerBase
from ...utils import TensorType, logging


lowercase_						=       logging.get_logger(__name__)

lowercase_						=       {
    """deepmind/language-perceiver""": """https://huggingface.co/deepmind/language-perceiver/resolve/main/config.json""",
    # See all Perceiver models at https://huggingface.co/models?filter=perceiver
}
class 			_snake_case     (   lowercase__):
				UpperCamelCase__			:   Any       						="""perceiver"""




				def __init__(  self   :  Any,   __lowercase   :  Optional[Any]=256,   __lowercase   :  List[str]=1280,   __lowercase   :  Dict=768,   __lowercase   :  int=1,   __lowercase   :  Dict=26,   __lowercase   :  Any=8,   __lowercase   :  List[Any]=8,   __lowercase   :  Dict=None,   __lowercase   :  List[Any]=None,   __lowercase   :  str="kv",   __lowercase   :  str=1,   __lowercase   :  Optional[Any]=1,   __lowercase   :  str="gelu",   __lowercase   :  List[str]=0.1,   __lowercase   :  int=0.02,   __lowercase   :  Union[str, Any]=1e-1_2,   __lowercase   :  Optional[Any]=True,   __lowercase   :  Optional[Any]=262,   __lowercase   :  str=2048,   __lowercase   :  Optional[Any]=56,   __lowercase   :  str=[368, 496],   __lowercase   :  str=16,   __lowercase   :  int=1920,   __lowercase   :  Dict=16,   __lowercase   :  List[Any]=[1, 16, 224, 224],   **__lowercase   :  str,   ):
							super().__init__(**__lowercase  )

							lowercase__   					=			num_latents
							lowercase__   					=			d_latents
							lowercase__   					=			d_model
							lowercase__   					=			num_blocks
							lowercase__   					=			num_self_attends_per_block
							lowercase__   					=			num_self_attention_heads
							lowercase__   					=			num_cross_attention_heads
							lowercase__   					=			qk_channels
							lowercase__   					=			v_channels
							lowercase__   					=			cross_attention_shape_for_attention
							lowercase__   					=			self_attention_widening_factor
							lowercase__   					=			cross_attention_widening_factor
							lowercase__   					=			hidden_act
							lowercase__   					=			attention_probs_dropout_prob
							lowercase__   					=			initializer_range
							lowercase__   					=			layer_norm_eps
							lowercase__   					=			use_query_residual
							# masked language modeling attributes
							lowercase__   					=			vocab_size
							lowercase__   					=			max_position_embeddings
							# image classification attributes
							lowercase__   					=			image_size
							# flow attributes
							lowercase__   					=			train_size
							# multimodal autoencoding attributes
							lowercase__   					=			num_frames
							lowercase__   					=			audio_samples_per_frame
							lowercase__   					=			samples_per_patch
							lowercase__   					=			output_shape
class 			_snake_case     (   lowercase__):




				@property
				def 							A__ (  self   :  Optional[int]  ):
							if self.task == "multiple-choice":
										lowercase__   					=			{0: "batch", 1: "choice", 2: "sequence"}
							else:
										lowercase__   					=			{0: "batch", 1: "sequence"}
							return OrderedDict(
							    [
							        ("inputs", dynamic_axis),
							        ("attention_mask", dynamic_axis),
							    ]  )




				@property
				def 							A__ (  self   :  Optional[Any]  ):
							return 1e-4




				def 							A__ (  self   :  Tuple,   __lowercase   :  Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"],   __lowercase   :  int = -1,   __lowercase   :  int = -1,   __lowercase   :  int = -1,   __lowercase   :  bool = False,   __lowercase   :  Optional[TensorType] = None,   __lowercase   :  int = 3,   __lowercase   :  int = 40,   __lowercase   :  int = 40,   ):
							# copied from `transformers.onnx.config.OnnxConfig` and slightly altered/simplified

							if isinstance(__lowercase,   __lowercase  ):
										# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
										lowercase__   					=			compute_effective_axis_dimension(
										    __lowercase,   fixed_dimension=OnnxConfig.default_fixed_batch,   num_token_to_add=0  )
										# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
										lowercase__   					=			preprocessor.num_special_tokens_to_add(__lowercase  )
										lowercase__   					=			compute_effective_axis_dimension(
										    __lowercase,   fixed_dimension=OnnxConfig.default_fixed_sequence,   num_token_to_add=__lowercase  )
										# Generate dummy inputs according to compute batch and sequence
										lowercase__   					=			[" ".join(["a"]  ) * seq_length] * batch_size
										lowercase__   					=			dict(preprocessor(__lowercase,   return_tensors=__lowercase  )  )
										lowercase__   					=			inputs.pop("input_ids"  )
										return inputs
							elif isinstance(__lowercase,   __lowercase  ) and preprocessor.model_input_names[0] == "pixel_values":
										# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
										lowercase__   					=			compute_effective_axis_dimension(__lowercase,   fixed_dimension=OnnxConfig.default_fixed_batch  )
										lowercase__   					=			self._generate_dummy_images(__lowercase,   __lowercase,   __lowercase,   __lowercase  )
										lowercase__   					=			dict(preprocessor(images=__lowercase,   return_tensors=__lowercase  )  )
										lowercase__   					=			inputs.pop("pixel_values"  )
										return inputs
							else:
										raise ValueError(
										    "Unable to generate dummy inputs for the model. Please provide a tokenizer or a preprocessor."  )




