


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


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


set_seed(42)
a_       							=    "sshleifer/student_marian_en_ro_6_1"
a_       							=    "sshleifer/tiny-mbart"
@require_torch
class snake_case      (  _UpperCamelCase):







		def  a_ (					self      :       str    ,   a__      :       str=False    ,   a__      :       str=None    ,   a__      :       Tuple=True    ,   a__      :       Optional[Any]=True    ,   a__      :       Union[str, Any]=True    ,   a__      :       Union[str, Any]=True    ,   )       ->						Any:





						'''simple docstring'''







						_A					  =			self.run_trainer(
						    eval_steps=1    ,   max_len=12    ,   model_name=a__    ,   num_train_epochs=1    ,   distributed=a__    ,   extra_args_str=a__    ,   predict_with_generate=a__    ,   do_train=a__    ,   do_eval=a__    ,   do_predict=a__    ,   )
						_A					  =			TrainerState.load_from_json(os.path.join(a__    ,   "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"]    ,   a__		)
										assert not math.isnan(float(last_step_stats["eval_loss"]		)		), "eval_loss must not be `nan`"







		@require_torch_non_multi_gpu
		def  a_ (					self      :       Dict		)       ->						Union[str, Any]:





						'''simple docstring'''







						self.run_seqaseq_quick()







		@require_torch_multi_gpu
		def  a_ (					self      :       Dict		)       ->						Optional[Any]:





						'''simple docstring'''







						self.run_seqaseq_quick(distributed=a__		)







		@require_torch_multi_gpu
		def  a_ (					self      :       str		)       ->						List[str]:





						'''simple docstring'''







						self.run_seqaseq_quick(distributed=a__		)







		@unittest.skip("Requires an update of the env running those tests"		)
		@require_torch_multi_gpu
		@require_fairscale
		def  a_ (					self      :       Any		)       ->						List[Any]:





						'''simple docstring'''







						self.run_seqaseq_quick(distributed=a__    ,   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      :       Dict		)       ->						Union[str, Any]:





						'''simple docstring'''







						self.run_seqaseq_quick(distributed=a__    ,   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      :       str		)       ->						Union[str, Any]:





						'''simple docstring'''







						self.run_seqaseq_quick(distributed=a__    ,   extra_args_str="--sharded_ddp zero_dp_2"    ,   predict_with_generate=a__		)







		@unittest.skip("Requires an update of the env running those tests"		)
		@require_torch_multi_gpu
		@require_fairscale
		def  a_ (					self      :       int		)       ->						Optional[int]:





						'''simple docstring'''







						self.run_seqaseq_quick(
						    distributed=a__    ,   extra_args_str="--sharded_ddp zero_dp_2 --fp16"    ,   predict_with_generate=a__		)







		@require_apex
		@require_torch_gpu
		def  a_ (					self      :       str		)       ->						int:





						'''simple docstring'''







						self.run_seqaseq_quick(distributed=a__    ,   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=a__    ,   extra_args_str="--fp16 --fp16_backend=apex"		)







		@parameterized.expand(["base", "low", "high", "mixed"]		)
		@require_torch_multi_gpu
		def  a_ (					self      :       Optional[Any]    ,   a__      :       Tuple		)       ->						List[str]:





						'''simple docstring'''







						_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(**a__    ,   extra_args_str=data["extra_args_str"]		)
						_A					  =			len(re.findall(a__    ,   cl.err		)		)
						self.assertEqual(a__    ,   data["n_matches"]		)







		@slow
		def  a_ (					self      :       Tuple		)       ->						Union[str, Any]:





						'''simple docstring'''







						_A					  =			self.run_trainer(
						    eval_steps=2    ,   max_len=1_28    ,   model_name=a__    ,   learning_rate=3E-4    ,   num_train_epochs=10    ,   distributed=a__    ,   )

						# Check metrics
						_A					  =			TrainerState.load_from_json(os.path.join(a__    ,   "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"]    ,   a__		)

						# test if do_predict saves generations and metrics
						_A					  =			os.listdir(a__		)
						_A					  =			{os.path.basename(a__		) for p in contents}
						assert "generated_predictions.txt" in contents
						assert "predict_results.json" in contents







		@slow
		@require_bitsandbytes
		def  a_ (					self      :       List[Any]		)       ->						str:





						'''simple docstring'''







						from transformers.training_args import OptimizerNames

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

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

										# Check metrics
										_A					  =			TrainerState.load_from_json(Path(a__    ,   "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(
						    a__    ,   a__    ,   "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(
						    a__    ,   a__    ,   "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(
						    a__    ,   a__    ,   F"""loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"""		)







		def  a_ (					self      :       Any    ,   a__      :       int    ,   a__      :       str    ,   a__      :       int    ,   a__      :       float = 3E-3    ,   a__      :       str = "adafactor"    ,   a__      :       bool = False    ,   a__      :       str = None    ,   a__      :       int = 0    ,   a__      :       bool = True    ,   a__      :       bool = True    ,   a__      :       bool = True    ,   a__      :       bool = True    ,   a__      :       int = None    ,   )       ->						str:





						'''simple docstring'''







						_A					  =			self.test_file_dir / "../fixtures/tests_samples/wmt_en_ro"
						_A					  =			self.get_auto_remove_tmp_dir()
						_A					  =			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(a__		)}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(a__		)}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        """.split()

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

						return output_dir



"""simple docstring"""


import os
import time
import warnings
from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional, Union

import torch
from filelock import FileLock
from torch.utils.data import Dataset

from ...tokenization_utils_base import PreTrainedTokenizerBase
from ...utils import logging
from ..processors.glue import glue_convert_examples_to_features, glue_output_modes, glue_processors
from ..processors.utils import InputFeatures


a_       							=    logging.get_logger(__name__)
@dataclass
class snake_case      :
		__UpperCamelCase           =			field(metadata={'help': 'The name of the task to train on: ' + ', '.join(glue_processors.keys())})
		__UpperCamelCase           =			field(
		    metadata={'help': 'The input data dir. Should contain the .tsv files (or other data files) for the task.'})
		__UpperCamelCase           =			field(
		    default=128     ,				metadata={
		        'help': (
		            'The maximum total input sequence length after tokenization. Sequences longer '
		            'than this will be truncated, sequences shorter will be padded.'
		        )
		    }     ,				)
		__UpperCamelCase           =			field(
		    default=_UpperCamelCase     ,				metadata={'help': 'Overwrite the cached training and evaluation sets'})







		def  a_ (					self      :       Optional[Any]		)       ->						Tuple:





						'''simple docstring'''







						_A					  =			self.task_name.lower()





class snake_case      (  _UpperCamelCase):
		__UpperCamelCase           =			'train'
		__UpperCamelCase           =			'dev'
		__UpperCamelCase           =			'test'





class snake_case      (  _UpperCamelCase):
		__UpperCamelCase           =			42
		__UpperCamelCase           =			42
		__UpperCamelCase           =			42







		def __init__(					self      :       Optional[int]    ,   a__      :       GlueDataTrainingArguments    ,   a__      :       PreTrainedTokenizerBase    ,   a__      :       Optional[int] = None    ,   a__      :       Union[str, Split] = Split.train    ,   a__      :       Optional[str] = None    ,   )       ->						Tuple:





						'''simple docstring'''







						warnings.warn(
						    "This dataset will be removed from the library soon, preprocessing should be handled with the 🤗 Datasets "
						    "library. You can have a look at this example script for pointers: "
						    "https://github.com/huggingface/transformers/blob/main/examples/pytorch/text-classification/run_glue.py"    ,   a__    ,   )
						_A					  =			args
						_A					  =			glue_processors[args.task_name]()
						_A					  =			glue_output_modes[args.task_name]
						if isinstance(a__    ,   a__		):
										try:
														_A					  =			Split[mode]
										except KeyError:
														raise KeyError("mode is not a valid split name"		)
        # Load data features from cache or dataset file
						_A					  =			os.path.join(
						    cache_dir if cache_dir is not None else args.data_dir    ,   F"""cached_{mode.value}_{tokenizer.__class__.__name__}_{args.max_seq_length}_{args.task_name}"""    ,   )
						_A					  =			self.processor.get_labels()
						if args.task_name in ["mnli", "mnli-mm"] and tokenizer.__class__.__name__ in (
						    "RobertaTokenizer",
						    "RobertaTokenizerFast",
						    "XLMRobertaTokenizer",
						    "BartTokenizer",
						    "BartTokenizerFast",
						):
										# HACK(label indices are swapped in RoBERTa pretrained model)
										_A  ,		_A					  =			label_list[2], label_list[1]
						_A					  =			label_list

						# Make sure only the first process in distributed training processes the dataset,
						# and the others will use the cache.
						_A					  =			cached_features_file + ".lock"
						with FileLock(a__		):
										if os.path.exists(a__		) and not args.overwrite_cache:
														_A					  =			time.time()
														_A					  =			torch.load(a__		)
														logger.info(
														    F"""Loading features from cached file {cached_features_file} [took %.3f s]"""    ,   time.time() - start		)
										else:
														logger.info(F"""Creating features from dataset file at {args.data_dir}"""		)

														if mode == Split.dev:
																		_A					  =			self.processor.get_dev_examples(args.data_dir		)
														elif mode == Split.test:
																		_A					  =			self.processor.get_test_examples(args.data_dir		)
														else:
																		_A					  =			self.processor.get_train_examples(args.data_dir		)
														if limit_length is not None:
																		_A					  =			examples[:limit_length]
														_A					  =			glue_convert_examples_to_features(
														    a__    ,   a__    ,   max_length=args.max_seq_length    ,   label_list=a__    ,   output_mode=self.output_mode    ,   )
														_A					  =			time.time()
														torch.save(self.features    ,   a__		)
														# ^ This seems to take a lot of time so I want to investigate why and how we can improve.
														logger.info(
														    F"""Saving features into cached file {cached_features_file} [took {time.time() - start:.3f} s]"""		)







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





						'''simple docstring'''







						return len(self.features		)







		def __getitem__(					self      :       Tuple    ,   a__      :       Union[str, Any]		)       ->						InputFeatures:





						'''simple docstring'''







						return self.features[i]







		def  a_ (					self      :       Optional[int]		)       ->						List[Any]:





						'''simple docstring'''







						return self.label_list