


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	: List[Any]  				=os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(f"""{bindir}/../../examples/pytorch/translation"""):
						from run_translation import main  # noqa


set_seed(4_2)
__lowerCAmelCase	: Any  				='sshleifer/student_marian_en_ro_6_1'
__lowerCAmelCase	: Optional[int]  				='sshleifer/tiny-mbart'
@require_torch
class        _lowercase      (			A__       ):






				'''simple docstring'''



				def 		__magic_name__(       self			:str			,			lowerCAmelCase__			:List[Any]=False			,			lowerCAmelCase__			:Tuple=None			,			lowerCAmelCase__			:Any=True			,			lowerCAmelCase__			:Tuple=True			,			lowerCAmelCase__			:Dict=True			,			lowerCAmelCase__			:Tuple=True			,			)		->			Union[str, Any]:
										__SCREAMING_SNAKE_CASE						:     str					=		self.run_trainer(
										    eval_steps=1			,			max_len=12			,			model_name=lowerCAmelCase__			,			num_train_epochs=1			,			distributed=lowerCAmelCase__			,			extra_args_str=lowerCAmelCase__			,			predict_with_generate=lowerCAmelCase__			,			do_train=lowerCAmelCase__			,			do_eval=lowerCAmelCase__			,			do_predict=lowerCAmelCase__			,			)
										__SCREAMING_SNAKE_CASE						:     Dict					=		TrainerState.load_from_json(os.path.join(lowerCAmelCase__			,			'''trainer_state.json'''	)	).log_history

										if not do_eval:
																return

										__SCREAMING_SNAKE_CASE						:     str					=		[log for log in logs if '''eval_loss''' in log.keys()]

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

																__SCREAMING_SNAKE_CASE						:     List[Any]					=		eval_metrics[-1]
																assert isinstance(last_step_stats['''eval_bleu''']			,			lowerCAmelCase__	)
																assert not math.isnan(float(last_step_stats['''eval_loss''']	)	), "eval_loss must not be `nan`"


				@require_torch_non_multi_gpu
				def 		__magic_name__(       self			:Optional[Any]	)		->			Dict:
										self.run_seqaseq_quick()


				@require_torch_multi_gpu
				def 		__magic_name__(       self			:str	)		->			Tuple:
										self.run_seqaseq_quick(distributed=lowerCAmelCase__	)


				@require_torch_multi_gpu
				def 		__magic_name__(       self			:Tuple	)		->			List[Any]:
										self.run_seqaseq_quick(distributed=lowerCAmelCase__	)


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


				@unittest.skip('''Requires an update of the env running those tests'''	)
				@require_torch_multi_gpu
				@require_fairscale
				def 		__magic_name__(       self			:Any	)		->			List[Any]:
										self.run_seqaseq_quick(distributed=lowerCAmelCase__			,			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 		__magic_name__(       self			:int	)		->			Any:
										self.run_seqaseq_quick(distributed=lowerCAmelCase__			,			extra_args_str='''--sharded_ddp zero_dp_2'''			,			predict_with_generate=lowerCAmelCase__	)


				@unittest.skip('''Requires an update of the env running those tests'''	)
				@require_torch_multi_gpu
				@require_fairscale
				def 		__magic_name__(       self			:str	)		->			Dict:
										self.run_seqaseq_quick(
										    distributed=lowerCAmelCase__			,			extra_args_str='''--sharded_ddp zero_dp_2 --fp16'''			,			predict_with_generate=lowerCAmelCase__	)


				@require_apex
				@require_torch_gpu
				def 		__magic_name__(       self			:Dict	)		->			List[Any]:
										# 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=lowerCAmelCase__			,			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=lowerCAmelCase__			,			extra_args_str='''--fp16 --fp16_backend=apex'''	)


				@parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']	)
				@require_torch_multi_gpu
				def 		__magic_name__(       self			:int			,			lowerCAmelCase__			:int	)		->			Tuple:
										# as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
										__SCREAMING_SNAKE_CASE						:     Union[str, 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},
										}

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


				@slow
				def 		__magic_name__(       self			:Any	)		->			List[Any]:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		self.run_trainer(
										    eval_steps=2			,			max_len=128			,			model_name=lowerCAmelCase__			,			learning_rate=3E-4			,			num_train_epochs=10			,			distributed=lowerCAmelCase__			,			)

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

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

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


				@slow
				@require_bitsandbytes
				def 		__magic_name__(       self			:Any	)		->			Optional[Any]:
										from transformers.training_args import OptimizerNames

										def train_and_return_metrics(lowerCAmelCase__			:str	) -> Tuple[int, float]:
																__SCREAMING_SNAKE_CASE						:     str					=		'''--skip_memory_metrics 0'''

																__SCREAMING_SNAKE_CASE						:     int					=		self.run_trainer(
																    max_len=128			,			model_name=lowerCAmelCase__			,			learning_rate=3E-4			,			num_train_epochs=1			,			optim=lowerCAmelCase__			,			distributed=lowerCAmelCase__			,			extra_args_str=lowerCAmelCase__			,			do_eval=lowerCAmelCase__			,			do_predict=lowerCAmelCase__			,			n_gpus_to_use=1			,			)

																# Check metrics
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		TrainerState.load_from_json(Path(lowerCAmelCase__			,			'''trainer_state.json'''	)	).log_history
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**20	)
																__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**20	)

																__SCREAMING_SNAKE_CASE						:     List[str]					=		logs[0]['''train_loss''']
																return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Dict					=		train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value	)
										__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		train_and_return_metrics(OptimizerNames.ADAMW_BNB.value	)

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		gpu_alloc_mem_orig - gpu_alloc_mem_bnb

										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		gpu_peak_mem_orig + gpu_alloc_mem_orig
										__SCREAMING_SNAKE_CASE						:     str					=		gpu_peak_mem_bnb + gpu_alloc_mem_bnb
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		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
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		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(
										    lowerCAmelCase__			,			lowerCAmelCase__			,			'''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(
										    lowerCAmelCase__			,			lowerCAmelCase__			,			'''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(
										    lowerCAmelCase__			,			lowerCAmelCase__			,			f'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''	)



				def 		__magic_name__(       self			:Optional[Any]			,			lowerCAmelCase__			:int			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:int			,			lowerCAmelCase__			:float = 3E-3			,			lowerCAmelCase__			:str = "adafactor"			,			lowerCAmelCase__			:bool = False			,			lowerCAmelCase__			:str = None			,			lowerCAmelCase__			:int = 0			,			lowerCAmelCase__			:bool = True			,			lowerCAmelCase__			:bool = True			,			lowerCAmelCase__			:bool = True			,			lowerCAmelCase__			:bool = True			,			lowerCAmelCase__			:int = None			,			)		->			List[Any]:
										__SCREAMING_SNAKE_CASE						:     Any					=		self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		self.get_auto_remove_tmp_dir()
										__SCREAMING_SNAKE_CASE						:     List[str]					=		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(lowerCAmelCase__	)}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(lowerCAmelCase__	)}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        '''.split()

										__SCREAMING_SNAKE_CASE						:     Any					=		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(lowerCAmelCase__	)}
        '''.split()

										__SCREAMING_SNAKE_CASE						:     str					=		'''
            --do_predict
        '''.split()

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

										return output_dir




import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging


if TYPE_CHECKING:
						from transformers.pipelines.conversational import Conversation


__lowerCAmelCase	: Optional[int]  				=logging.get_logger(__name__)

__lowerCAmelCase	: Optional[Any]  				={'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

__lowerCAmelCase	: List[str]  				={
    'tokenizer_file': {
        'EleutherAI/gpt-neox-20b': 'https://huggingface.co/EleutherAI/gpt-neox-20b/resolve/main/tokenizer.json',
    },
}

__lowerCAmelCase	: Optional[int]  				={
    'gpt-neox-20b': 2_0_4_8,
}
class        _lowercase      (			A__       ):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			str						=		VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE__				:			Dict						=		PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE__				:			Any						=		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE__				:			Optional[Any]						=		['''input_ids''', '''attention_mask''']


				def __init__(       self			:int			,			lowerCAmelCase__			:Any=None			,			lowerCAmelCase__			:Optional[Any]=None			,			lowerCAmelCase__			:List[Any]=None			,			lowerCAmelCase__			:str="<|endoftext|>"			,			lowerCAmelCase__			:str="<|endoftext|>"			,			lowerCAmelCase__			:Dict="<|endoftext|>"			,			lowerCAmelCase__			:Union[str, Any]=False			,			**lowerCAmelCase__			:List[str]			,			)		->			Any:
										super().__init__(
										    lowerCAmelCase__			,			lowerCAmelCase__			,			tokenizer_file=lowerCAmelCase__			,			unk_token=lowerCAmelCase__			,			bos_token=lowerCAmelCase__			,			eos_token=lowerCAmelCase__			,			add_prefix_space=lowerCAmelCase__			,			**lowerCAmelCase__			,			)

										__SCREAMING_SNAKE_CASE						:     List[Any]					=		json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()	)
										if pre_tok_state.get('''add_prefix_space'''			,			lowerCAmelCase__	) != add_prefix_space:
																__SCREAMING_SNAKE_CASE						:     List[str]					=		getattr(lowerCAmelCase__			,			pre_tok_state.pop('''type'''	)	)
																__SCREAMING_SNAKE_CASE						:     str					=		add_prefix_space
																__SCREAMING_SNAKE_CASE						:     Any					=		pre_tok_class(**lowerCAmelCase__	)

										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		add_prefix_space


				def 		__magic_name__(       self			:Union[str, Any]			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:Optional[str] = None	)		->			Tuple[str]:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		self._tokenizer.model.save(lowerCAmelCase__			,			name=lowerCAmelCase__	)
										return tuple(lowerCAmelCase__	)



				def 		__magic_name__(       self			:Optional[Any]			,			lowerCAmelCase__			:"Conversation"	)		->			List[int]:
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		[]
										for is_user, text in conversation.iter_texts():
																input_ids.extend(self.encode(lowerCAmelCase__			,			add_special_tokens=lowerCAmelCase__	) + [self.eos_token_id]	)

										if len(lowerCAmelCase__	) > self.model_max_length:
																__SCREAMING_SNAKE_CASE						:     List[str]					=		input_ids[-self.model_max_length :]
										return input_ids
