
import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import BertGenerationDecoder, BertGenerationEncoder






class  __magic_name__						:
				'''simple docstring'''







				def __init__(							self:   List[Any]		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   Optional[Any]=13		,					_lowerCamelCase:   Tuple=7		,					_lowerCamelCase:   Union[str, Any]=True		,					_lowerCamelCase:   int=True		,					_lowerCamelCase:   str=99		,					_lowerCamelCase:   List[str]=32		,					_lowerCamelCase:   Tuple=5		,					_lowerCamelCase:   Optional[Any]=4		,					_lowerCamelCase:   Tuple=37		,					_lowerCamelCase:   Dict="gelu"		,					_lowerCamelCase:   Optional[int]=0.1		,					_lowerCamelCase:   Optional[Any]=0.1		,					_lowerCamelCase:   int=50		,					_lowerCamelCase:   Optional[int]=0.02		,					_lowerCamelCase:   Any=True		,					_lowerCamelCase:   List[Any]=None		,					):
											SCREAMING_SNAKE_CASE_   	=    parent
											SCREAMING_SNAKE_CASE_   	=    batch_size
											SCREAMING_SNAKE_CASE_   	=    seq_length
											SCREAMING_SNAKE_CASE_   	=    is_training
											SCREAMING_SNAKE_CASE_   	=    use_input_mask
											SCREAMING_SNAKE_CASE_   	=    vocab_size
											SCREAMING_SNAKE_CASE_   	=    hidden_size
											SCREAMING_SNAKE_CASE_   	=    num_hidden_layers
											SCREAMING_SNAKE_CASE_   	=    num_attention_heads
											SCREAMING_SNAKE_CASE_   	=    intermediate_size
											SCREAMING_SNAKE_CASE_   	=    hidden_act
											SCREAMING_SNAKE_CASE_   	=    hidden_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    max_position_embeddings
											SCREAMING_SNAKE_CASE_   	=    initializer_range
											SCREAMING_SNAKE_CASE_   	=    use_labels
											SCREAMING_SNAKE_CASE_   	=    scope


				def     _A							(							self:   int		):
											SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size		)

											SCREAMING_SNAKE_CASE_   	=    None
											if self.use_input_mask:
																		SCREAMING_SNAKE_CASE_   	=    random_attention_mask([self.batch_size, self.seq_length]		)

											if self.use_labels:
																		SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size		)

											SCREAMING_SNAKE_CASE_   	=    self.get_config()

											return config, input_ids, input_mask, token_labels


				def     _A							(							self:   List[str]		):
											return BertGenerationConfig(
											    vocab_size=self.vocab_size		,					hidden_size=self.hidden_size		,					num_hidden_layers=self.num_hidden_layers		,					num_attention_heads=self.num_attention_heads		,					intermediate_size=self.intermediate_size		,					hidden_act=self.hidden_act		,					hidden_dropout_prob=self.hidden_dropout_prob		,					attention_probs_dropout_prob=self.attention_probs_dropout_prob		,					max_position_embeddings=self.max_position_embeddings		,					is_decoder=_lowerCamelCase		,					initializer_range=self.initializer_range		,					)


				def     _A							(							self:   Tuple		):
											(
											    (
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,		
											)   	=    self.prepare_config_and_inputs()

											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
											SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length]		,					vocab_size=2		)

											return (
											    config,
											    input_ids,
											    input_mask,
											    token_labels,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)


				def     _A							(							self:   str		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Tuple		,					**_lowerCamelCase:   Optional[int]		,					):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder(config=_lowerCamelCase		)
											model.to(_lowerCamelCase		)
											model.eval()
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		,					attention_mask=_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)
											self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)		)


				def     _A							(							self:   Any		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   Optional[int]		,					_lowerCamelCase:   Tuple		,					**_lowerCamelCase:   Any		,					):
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder(config=_lowerCamelCase		)
											model.to(_lowerCamelCase		)
											model.eval()
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					)
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					)
											self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)		)


				def     _A							(							self:   Optional[Any]		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   Optional[int]		,					_lowerCamelCase:   List[str]		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   str		,					**_lowerCamelCase:   Dict		,					):
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    BertGenerationDecoder(config=_lowerCamelCase		).to(_lowerCamelCase		).eval()

											# first forward pass
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					use_cache=_lowerCamelCase		,					)
											SCREAMING_SNAKE_CASE_   	=    outputs.past_key_values

											# create hypothetical multiple next token and extent to next_input_ids
											SCREAMING_SNAKE_CASE_   	=    ids_tensor((self.batch_size, 3)		,					config.vocab_size		)
											SCREAMING_SNAKE_CASE_   	=    ids_tensor((self.batch_size, 3)		,					vocab_size=2		)

											# append to next input_ids and
											SCREAMING_SNAKE_CASE_   	=    torch.cat([input_ids, next_tokens]		,					dim=-1		)
											SCREAMING_SNAKE_CASE_   	=    torch.cat([input_mask, next_mask]		,					dim=-1		)

											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					output_hidden_states=_lowerCamelCase		,					)['''hidden_states'''][0]
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					past_key_values=_lowerCamelCase		,					output_hidden_states=_lowerCamelCase		,					)['''hidden_states'''][0]

											# select random slice
											SCREAMING_SNAKE_CASE_   	=    ids_tensor((1,)		,					output_from_past.shape[-1]		).item()
											SCREAMING_SNAKE_CASE_   	=    output_from_no_past[:, -3:, random_slice_idx].detach()
											SCREAMING_SNAKE_CASE_   	=    output_from_past[:, :, random_slice_idx].detach()

											self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

											# test that outputs are equal for slice
											self.parent.assertTrue(torch.allclose(_lowerCamelCase		,					_lowerCamelCase		,					atol=1E-3		)		)


				def     _A							(							self:   List[Any]		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   int		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Optional[Any]		,					*_lowerCamelCase:   List[Any]		,					):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationDecoder(_lowerCamelCase		)
											model.to(_lowerCamelCase		)
											model.eval()
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		,					attention_mask=_lowerCamelCase		,					labels=_lowerCamelCase		)
											self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)		)





				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    self.prepare_config_and_inputs()
											SCREAMING_SNAKE_CASE_   	=    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
											return config, inputs_dict






@require_torch
class  __magic_name__						(  __UpperCAmelCase    ,  __UpperCAmelCase    ,  __UpperCAmelCase    ,  unittest.TestCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							Optional[Any]	      =					(BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
				SCREAMING_SNAKE_CASE__					:							Any	      =					(BertGenerationDecoder,) if is_torch_available() else ()
				SCREAMING_SNAKE_CASE__					:							Dict	      =					(
				    {"feature-extraction": BertGenerationEncoder, "text-generation": BertGenerationDecoder}
				    if is_torch_available()
				    else {}
				)


				def     _A							(							self:   Union[str, Any]		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoderTester(self		)
											SCREAMING_SNAKE_CASE_   	=    ConfigTester(self		,					config_class=_lowerCamelCase		,					hidden_size=37		)


				def     _A							(							self:   Any		):
											self.config_tester.run_common_tests()


				def     _A							(							self:   int		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_lowerCamelCase		)


				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs()
											SCREAMING_SNAKE_CASE_   	=    '''bert'''
											self.model_tester.create_and_check_model(_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		)


				def     _A							(							self:   Optional[int]		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase		)


				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCamelCase		)


				def     _A							(							self:   Optional[Any]		):
											# This regression test was failing with PyTorch < 1.3
											(
											    (
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,	(
											    SCREAMING_SNAKE_CASE_	
											)	,		
											)   	=    self.model_tester.prepare_config_and_inputs_for_decoder()

											SCREAMING_SNAKE_CASE_   	=    None

											self.model_tester.create_and_check_model_as_decoder(
											    _lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					)


				def     _A							(							self:   Optional[Any]		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_for_causal_lm(*_lowerCamelCase		)





				@slow
				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''		)
											self.assertIsNotNone(_lowerCamelCase		)






@require_torch
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







				@slow
				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]]		)
											with torch.no_grad():
																		SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)[0]
											SCREAMING_SNAKE_CASE_   	=    torch.Size([1, 8, 10_24]		)
											self.assertEqual(output.shape		,					_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor(
											    [[[0.17_75, 0.00_83, -0.03_21], [1.60_02, 0.12_87, 0.39_12], [2.14_73, 0.57_91, 0.60_66]]]		)
											self.assertTrue(torch.allclose(output[:, :3, :3]		,					_lowerCamelCase		,					atol=1E-4		)		)






@require_torch
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







				@slow
				def     _A							(							self:   Optional[int]		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]]		)
											with torch.no_grad():
																		SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)[0]
											SCREAMING_SNAKE_CASE_   	=    torch.Size([1, 8, 5_03_58]		)
											self.assertEqual(output.shape		,					_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor(
											    [[[-0.57_88, -2.59_94, -3.70_54], [0.04_38, 4.79_97, 1.87_95], [1.58_62, 6.64_09, 4.46_38]]]		)
											self.assertTrue(torch.allclose(output[:, :3, :3]		,					_lowerCamelCase		,					atol=1E-4		)		)








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


__SCREAMING_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)
__SCREAMING_SNAKE_CASE 						="""sshleifer/student_marian_en_ro_6_1"""
__SCREAMING_SNAKE_CASE 						="""sshleifer/tiny-mbart"""






@require_torch
class  __magic_name__						(  __UpperCAmelCase):
				'''simple docstring'''







				def     _A							(							self:   List[Any]		,					_lowerCamelCase:   List[Any]=False		,					_lowerCamelCase:   List[Any]=None		,					_lowerCamelCase:   int=True		,					_lowerCamelCase:   Optional[int]=True		,					_lowerCamelCase:   Any=True		,					_lowerCamelCase:   Any=True		,					):
											SCREAMING_SNAKE_CASE_   	=    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_   	=    TrainerState.load_from_json(os.path.join(_lowerCamelCase		,					'''trainer_state.json'''		)		).log_history

											if not do_eval:
																		return

											SCREAMING_SNAKE_CASE_   	=    [log for log in logs if '''eval_loss''' in log.keys()]

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

																		SCREAMING_SNAKE_CASE_   	=    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     _A							(							self:   Optional[int]		):
											self.run_seqaseq_quick()


				@require_torch_multi_gpu
				def     _A							(							self:   Optional[Any]		):
											self.run_seqaseq_quick(distributed=_lowerCamelCase		)


				@require_torch_multi_gpu
				def     _A							(							self:   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     _A							(							self:   Optional[int]		):
											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     _A							(							self:   Optional[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     _A							(							self:   Union[str, 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     _A							(							self:   Optional[Any]		):
											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     _A							(							self:   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     _A							(							self:   Union[str, Any]		,					_lowerCamelCase:   Dict		):
											# as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
											SCREAMING_SNAKE_CASE_   	=    {
											    # 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_   	=    experiments[experiment_id]
											SCREAMING_SNAKE_CASE_   	=    {'''distributed''': True, '''predict_with_generate''': False, '''do_eval''': False, '''do_predict''': False}
											SCREAMING_SNAKE_CASE_   	=    '''Running training'''
											with CaptureStderr() as cl:
																		self.run_seqaseq_quick(**_lowerCamelCase		,					extra_args_str=data['''extra_args_str''']		)
											SCREAMING_SNAKE_CASE_   	=    len(re.findall(_lowerCamelCase		,					cl.err		)		)
											self.assertEqual(_lowerCamelCase		,					data['''n_matches''']		)


				@slow
				def     _A							(							self:   Any		):
											SCREAMING_SNAKE_CASE_   	=    self.run_trainer(
											    eval_steps=2		,					max_len=1_28		,					model_name=_lowerCamelCase		,					learning_rate=3E-4		,					num_train_epochs=10		,					distributed=_lowerCamelCase		,					)

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


				@slow
				@require_bitsandbytes
				def     _A							(							self:   Optional[int]		):
											from transformers.training_args import OptimizerNames

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

																		SCREAMING_SNAKE_CASE_   	=    self.run_trainer(
																		    max_len=1_28		,					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_   	=    TrainerState.load_from_json(Path(_lowerCamelCase		,					'''trainer_state.json'''		)		).log_history
																		SCREAMING_SNAKE_CASE_   	=    int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**20		)
																		SCREAMING_SNAKE_CASE_   	=    int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**20		)

																		SCREAMING_SNAKE_CASE_   	=    logs[0]['''train_loss''']
																		return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

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

											SCREAMING_SNAKE_CASE_   	=    gpu_alloc_mem_orig - gpu_alloc_mem_bnb

											SCREAMING_SNAKE_CASE_   	=    gpu_peak_mem_orig + gpu_alloc_mem_orig
											SCREAMING_SNAKE_CASE_   	=    gpu_peak_mem_bnb + gpu_alloc_mem_bnb
											SCREAMING_SNAKE_CASE_   	=    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_   	=    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(
											    _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     _A							(							self:   Union[str, 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		,					):
											SCREAMING_SNAKE_CASE_   	=    self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
											SCREAMING_SNAKE_CASE_   	=    self.get_auto_remove_tmp_dir()
											SCREAMING_SNAKE_CASE_   	=    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(_lowerCamelCase		)}\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(_lowerCamelCase		)}\n            --group_by_length\n            --label_smoothing_factor 0.1\n            --target_lang ro_RO\n            --source_lang en_XX\n        ".split()

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

											SCREAMING_SNAKE_CASE_   	=    '''
            --do_predict
        '''.split()

											SCREAMING_SNAKE_CASE_   	=    []
											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_   	=    get_gpu_count()
																		SCREAMING_SNAKE_CASE_   	=    get_torch_dist_unique_port()
																		SCREAMING_SNAKE_CASE_   	=    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()
																		SCREAMING_SNAKE_CASE_   	=    [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_   	=    ['''run_translation.py'''] + args
																		with patch.object(_lowerCamelCase		,					'''argv'''		,					_lowerCamelCase		):
																									main()

											return output_dir


