






"""simple docstring"""



from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import (
						    TF_MODEL_FOR_PRETRAINING_MAPPING,
						    TFMobileBertForMaskedLM,
						    TFMobileBertForMultipleChoice,
						    TFMobileBertForNextSentencePrediction,
						    TFMobileBertForPreTraining,
						    TFMobileBertForQuestionAnswering,
						    TFMobileBertForSequenceClassification,
						    TFMobileBertForTokenClassification,
						    TFMobileBertModel,
						)





@require_tf
class 					_UpperCAmelCase       ( _A   ,						_A   ,						unittest.TestCase							):
					SCREAMING_SNAKE_CASE_		:				int							    =   (
					    (
					        TFMobileBertModel,
					        TFMobileBertForMaskedLM,
					        TFMobileBertForNextSentencePrediction,
					        TFMobileBertForPreTraining,
					        TFMobileBertForQuestionAnswering,
					        TFMobileBertForSequenceClassification,
					        TFMobileBertForTokenClassification,
					        TFMobileBertForMultipleChoice,
					    )
					    if is_tf_available()
					    else ()
					)
					SCREAMING_SNAKE_CASE_		:				Any							    =   (
					    {
					        "feature-extraction": TFMobileBertModel,
					        "fill-mask": TFMobileBertForMaskedLM,
					        "question-answering": TFMobileBertForQuestionAnswering,
					        "text-classification": TFMobileBertForSequenceClassification,
					        "token-classification": TFMobileBertForTokenClassification,
					        "zero-shot": TFMobileBertForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					SCREAMING_SNAKE_CASE_		:				str							    =   False
					SCREAMING_SNAKE_CASE_		:				int							    =   False




					def   A  (				self     :						Dict    ,						A     :						Dict    ,						A     :						List[Any]    ,						A     :						Tuple=False   )  ->  int:
							lowercase_ : str							       =   super()._prepare_for_class(A    ,						A    ,						return_labels=A   )

							if return_labels:
									if model_class in get_values(A   ):
											lowercase_ : List[str]							       =   tf.zeros(self.model_tester.batch_size    ,						dtype=tf.intaa   )

							return inputs_dict





					class 					_UpperCAmelCase       ( _A							):




										def __init__(				self     :						Any    ,						A     :						int    ,						A     :						int=13    ,						A     :						Any=7    ,						A     :						Optional[Any]=True    ,						A     :						List[str]=True    ,						A     :						List[Any]=True    ,						A     :						Dict=True    ,						A     :						List[str]=99    ,						A     :						str=32    ,						A     :						str=32    ,						A     :						List[Any]=2    ,						A     :						Tuple=4    ,						A     :						Optional[Any]=37    ,						A     :						Tuple="gelu"    ,						A     :						Optional[int]=0.1    ,						A     :						Tuple=0.1    ,						A     :						Optional[int]=5_12    ,						A     :						List[Any]=16    ,						A     :						Optional[Any]=2    ,						A     :						Any=0.02    ,						A     :						List[str]=3    ,						A     :						Dict=4    ,						A     :						Tuple=None    ,						)  ->  List[Any]:
												lowercase_ : Optional[int]							       =   parent
												lowercase_ : Optional[int]							       =   batch_size
												lowercase_ : Union[str, Any]							       =   seq_length
												lowercase_ : List[str]							       =   is_training
												lowercase_ : str							       =   use_input_mask
												lowercase_ : Tuple							       =   use_token_type_ids
												lowercase_ : Any							       =   use_labels
												lowercase_ : Optional[int]							       =   vocab_size
												lowercase_ : Tuple							       =   hidden_size
												lowercase_ : List[str]							       =   num_hidden_layers
												lowercase_ : Dict							       =   num_attention_heads
												lowercase_ : Any							       =   intermediate_size
												lowercase_ : Tuple							       =   hidden_act
												lowercase_ : Optional[Any]							       =   hidden_dropout_prob
												lowercase_ : Optional[Any]							       =   attention_probs_dropout_prob
												lowercase_ : str							       =   max_position_embeddings
												lowercase_ : str							       =   type_vocab_size
												lowercase_ : Tuple							       =   type_sequence_label_size
												lowercase_ : str							       =   initializer_range
												lowercase_ : str							       =   num_labels
												lowercase_ : Optional[Any]							       =   num_choices
												lowercase_ : Any							       =   scope
												lowercase_ : int							       =   embedding_size




										def   A  (				self     :						List[Any]   )  ->  Any:
												lowercase_ : List[Any]							       =   ids_tensor([self.batch_size, self.seq_length]    ,						self.vocab_size   )

												lowercase_ : Any							       =   None
												if self.use_input_mask:
														lowercase_ : List[Any]							       =   random_attention_mask([self.batch_size, self.seq_length]   )

												lowercase_ : int							       =   None
												if self.use_token_type_ids:
														lowercase_ : Tuple							       =   ids_tensor([self.batch_size, self.seq_length]    ,						self.type_vocab_size   )

												lowercase_ : str							       =   None
												lowercase_ : Dict							       =   None
												lowercase_ : Any							       =   None
												if self.use_labels:
														lowercase_ : List[str]							       =   ids_tensor([self.batch_size]    ,						self.type_sequence_label_size   )
														lowercase_ : List[str]							       =   ids_tensor([self.batch_size, self.seq_length]    ,						self.num_labels   )
														lowercase_ : int							       =   ids_tensor([self.batch_size]    ,						self.num_choices   )

												lowercase_ : Optional[Any]							       =   MobileBertConfig(
												    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    ,						type_vocab_size=self.type_vocab_size    ,						initializer_range=self.initializer_range    ,						embedding_size=self.embedding_size    ,						)

												return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels




										def   A  (				self     :						List[str]    ,						A     :						List[str]    ,						A     :						Optional[int]    ,						A     :						str    ,						A     :						Union[str, Any]    ,						A     :						List[str]    ,						A     :						Optional[Any]    ,						A     :						Union[str, Any]   )  ->  Optional[Any]:
												lowercase_ : Optional[int]							       =   TFMobileBertModel(config=A   )
												lowercase_ : List[Any]							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
												lowercase_ : Tuple							       =   model(A   )

												lowercase_ : Optional[int]							       =   [input_ids, input_mask]
												lowercase_ : int							       =   model(A   )

												lowercase_ : Optional[int]							       =   model(A   )

												self.parent.assertEqual(
												    result.last_hidden_state.shape    ,						(self.batch_size, self.seq_length, self.hidden_size)   )
												self.parent.assertEqual(result.pooler_output.shape    ,						(self.batch_size, self.hidden_size)   )




										def   A  (				self     :						List[Any]    ,						A     :						int    ,						A     :						int    ,						A     :						int    ,						A     :						int    ,						A     :						Optional[int]    ,						A     :						List[str]    ,						A     :						Dict   )  ->  List[Any]:
												lowercase_ : List[Any]							       =   TFMobileBertForMaskedLM(config=A   )
												lowercase_ : Tuple							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
												lowercase_ : Optional[int]							       =   model(A   )
												self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)   )




										def   A  (				self     :						List[Any]    ,						A     :						Optional[Any]    ,						A     :						int    ,						A     :						Dict    ,						A     :						Union[str, Any]    ,						A     :						List[Any]    ,						A     :						List[str]    ,						A     :						List[Any]   )  ->  Union[str, Any]:
												lowercase_ : Optional[int]							       =   TFMobileBertForNextSentencePrediction(config=A   )
												lowercase_ : int							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
												lowercase_ : Optional[int]							       =   model(A   )
												self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, 2)   )




										def   A  (				self     :						Any    ,						A     :						List[str]    ,						A     :						Union[str, Any]    ,						A     :						Any    ,						A     :						Any    ,						A     :						Tuple    ,						A     :						Optional[Any]    ,						A     :						str   )  ->  List[str]:
												lowercase_ : Optional[Any]							       =   TFMobileBertForPreTraining(config=A   )
												lowercase_ : List[str]							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
												lowercase_ : List[Any]							       =   model(A   )
												self.parent.assertEqual(
												    result.prediction_logits.shape    ,						(self.batch_size, self.seq_length, self.vocab_size)   )
												self.parent.assertEqual(result.seq_relationship_logits.shape    ,						(self.batch_size, 2)   )




										def   A  (				self     :						Dict    ,						A     :						Optional[int]    ,						A     :						List[Any]    ,						A     :						List[Any]    ,						A     :						Union[str, Any]    ,						A     :						Optional[Any]    ,						A     :						Any    ,						A     :						str   )  ->  str:
												lowercase_ : Tuple							       =   self.num_labels
												lowercase_ : Optional[Any]							       =   TFMobileBertForSequenceClassification(config=A   )
												lowercase_ : Optional[int]							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
												lowercase_ : Optional[int]							       =   model(A   )
												self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_labels)   )




										def   A  (				self     :						List[str]    ,						A     :						Optional[int]    ,						A     :						Union[str, Any]    ,						A     :						int    ,						A     :						Optional[int]    ,						A     :						List[Any]    ,						A     :						Any    ,						A     :						Tuple   )  ->  str:
												lowercase_ : Tuple							       =   self.num_choices
												lowercase_ : List[str]							       =   TFMobileBertForMultipleChoice(config=A   )
												lowercase_ : Tuple							       =   tf.tile(tf.expand_dims(A    ,						1   )    ,						(1, self.num_choices, 1)   )
												lowercase_ : Dict							       =   tf.tile(tf.expand_dims(A    ,						1   )    ,						(1, self.num_choices, 1)   )
												lowercase_ : List[str]							       =   tf.tile(tf.expand_dims(A    ,						1   )    ,						(1, self.num_choices, 1)   )
												lowercase_ : Optional[int]							       =   {
												    '''input_ids''': multiple_choice_inputs_ids,
												    '''attention_mask''': multiple_choice_input_mask,
												    '''token_type_ids''': multiple_choice_token_type_ids,
												}
												lowercase_ : Union[str, Any]							       =   model(A   )
												self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_choices)   )




										def   A  (				self     :						Tuple    ,						A     :						Optional[int]    ,						A     :						Optional[Any]    ,						A     :						Any    ,						A     :						Dict    ,						A     :						Optional[int]    ,						A     :						Dict    ,						A     :						str   )  ->  str:
												lowercase_ : Optional[Any]							       =   self.num_labels
												lowercase_ : Optional[int]							       =   TFMobileBertForTokenClassification(config=A   )
												lowercase_ : List[Any]							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
												lowercase_ : Tuple							       =   model(A   )
												self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.seq_length, self.num_labels)   )




										def   A  (				self     :						Dict    ,						A     :						Dict    ,						A     :						Optional[Any]    ,						A     :						List[Any]    ,						A     :						Union[str, Any]    ,						A     :						List[str]    ,						A     :						int    ,						A     :						str   )  ->  Optional[Any]:
												lowercase_ : Optional[int]							       =   TFMobileBertForQuestionAnswering(config=A   )
												lowercase_ : str							       =   {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
												lowercase_ : Union[str, Any]							       =   model(A   )
												self.parent.assertEqual(result.start_logits.shape    ,						(self.batch_size, self.seq_length)   )
												self.parent.assertEqual(result.end_logits.shape    ,						(self.batch_size, self.seq_length)   )






										def   A  (				self     :						int   )  ->  Union[str, Any]:
												lowercase_ : Optional[int]							       =   self.prepare_config_and_inputs()
												(
												    (
												    lowercase_
												)       ,					(
												    lowercase_
												)       ,					(
												    lowercase_
												)       ,					(
												    lowercase_
												)       ,					(
												    lowercase_
												)       ,					(
												    lowercase_
												)       ,					(
												    lowercase_
												)       ,					
												) : Any							       =   config_and_inputs
												lowercase_ : List[str]							       =   {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
												return config, inputs_dict




					def   A  (				self     :						Tuple   )  ->  Tuple:
							lowercase_ : List[Any]							       =   TFMobileBertModelTest.TFMobileBertModelTester(self   )
							lowercase_ : List[str]							       =   ConfigTester(self    ,						config_class=A    ,						hidden_size=37   )




					def   A  (				self     :						int   )  ->  Optional[Any]:
							self.config_tester.run_common_tests()




					def   A  (				self     :						List[str]   )  ->  Tuple:
							lowercase_ : Any							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_model(*A   )




					def   A  (				self     :						List[Any]   )  ->  Union[str, Any]:
							lowercase_ : List[Any]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_masked_lm(*A   )




					def   A  (				self     :						List[str]   )  ->  List[Any]:
							lowercase_ : str							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_multiple_choice(*A   )




					def   A  (				self     :						str   )  ->  Any:
							lowercase_ : Union[str, Any]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*A   )




					def   A  (				self     :						Optional[int]   )  ->  int:
							lowercase_ : List[str]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_pretraining(*A   )




					def   A  (				self     :						str   )  ->  int:
							lowercase_ : str							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_question_answering(*A   )




					def   A  (				self     :						Optional[Any]   )  ->  List[str]:
							lowercase_ : List[Any]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_sequence_classification(*A   )




					def   A  (				self     :						str   )  ->  Any:
							lowercase_ : Union[str, Any]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_mobilebert_for_token_classification(*A   )






					@slow
					def   A  (				self     :						Optional[Any]   )  ->  Optional[Any]:
							# for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							for model_name in ["google/mobilebert-uncased"]:
									lowercase_ : Any							       =   TFMobileBertModel.from_pretrained(A   )
									self.assertIsNotNone(A   )





@require_tf
class 					_UpperCAmelCase       ( unittest.TestCase							):




					@slow
					def   A  (				self     :						Union[str, Any]   )  ->  List[Any]:
							lowercase_ : Any							       =   TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased'''   )
							lowercase_ : Any							       =   tf.constant([[0, 1, 2, 3, 4, 5]]   )
							lowercase_ : List[Any]							       =   model(A   )[0]

							lowercase_ : Optional[int]							       =   [1, 6, 3_05_22]
							self.assertEqual(output.shape    ,						A   )

							lowercase_ : Tuple							       =   tf.constant(
							    [
							        [
							            [-4.5919547, -9.248295, -9.645256],
							            [-6.7306175, -6.440284, -6.6052837],
							            [-7.2743506, -6.7847915, -6.024673],
							        ]
							    ]   )
							tf.debugging.assert_near(output[:, :3, :3]    ,						A    ,						atol=1e-4   )







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


set_seed(42)
a_  :				int						= """sshleifer/student_marian_en_ro_6_1"""
a_  :				str						= """sshleifer/tiny-mbart"""





@require_torch
class      snake_case    (    lowercase			):





  """simple docstring"""




  def   snake_case		(    self			,  UpperCamelCase=False			,  UpperCamelCase=None			,  UpperCamelCase=True			,  UpperCamelCase=True			,  UpperCamelCase=True			,  UpperCamelCase=True			,  ):







   """simple docstring"""
   lowerCamelCase_      =  self.run_trainer(
       eval_steps=1			,  max_len=12			,  model_name=UpperCamelCase			,  num_train_epochs=1			,  distributed=UpperCamelCase			,  extra_args_str=UpperCamelCase			,  predict_with_generate=UpperCamelCase			,  do_train=UpperCamelCase			,  do_eval=UpperCamelCase			,  do_predict=UpperCamelCase			,  )
   lowerCamelCase_      =  TrainerState.load_from_json(os.path.join(UpperCamelCase			,  "trainer_state.json"     )     ).log_history

   if not do_eval:
    return

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

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

    lowerCamelCase_      =  eval_metrics[-1]
    assert isinstance(last_step_stats["eval_bleu"]			,  UpperCamelCase     )
    assert not math.isnan(float(last_step_stats["eval_loss"]     )     ), "eval_loss must not be `nan`"






  @require_torch_non_multi_gpu
  def   snake_case		(    self     ):







   """simple docstring"""
   self.run_seqaseq_quick()






  @require_torch_multi_gpu
  def   snake_case		(    self     ):







   """simple docstring"""
   self.run_seqaseq_quick(distributed=UpperCamelCase     )






  @require_torch_multi_gpu
  def   snake_case		(    self     ):







   """simple docstring"""
   self.run_seqaseq_quick(distributed=UpperCamelCase     )






  @unittest.skip("Requires an update of the env running those tests"     )
  @require_torch_multi_gpu
  @require_fairscale
  def   snake_case		(    self     ):







   """simple docstring"""
   self.run_seqaseq_quick(distributed=UpperCamelCase			,  extra_args_str="--sharded_ddp simple"     )






  @unittest.skip("Requires an update of the env running those tests"     )
  @require_torch_multi_gpu
  @require_fairscale
  def   snake_case		(    self     ):







   """simple docstring"""
   self.run_seqaseq_quick(distributed=UpperCamelCase			,  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   snake_case		(    self     ):







   """simple docstring"""
   self.run_seqaseq_quick(distributed=UpperCamelCase			,  extra_args_str="--sharded_ddp zero_dp_2"			,  predict_with_generate=UpperCamelCase     )






  @unittest.skip("Requires an update of the env running those tests"     )
  @require_torch_multi_gpu
  @require_fairscale
  def   snake_case		(    self     ):







   """simple docstring"""
   self.run_seqaseq_quick(
       distributed=UpperCamelCase			,  extra_args_str="--sharded_ddp zero_dp_2 --fp16"			,  predict_with_generate=UpperCamelCase     )






  @require_apex
  @require_torch_gpu
  def   snake_case		(    self     ):







   """simple docstring"""
   # 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=UpperCamelCase			,  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=UpperCamelCase			,  extra_args_str="--fp16 --fp16_backend=apex"     )






  @parameterized.expand(["base", "low", "high", "mixed"]     )
  @require_torch_multi_gpu
  def   snake_case		(    self			,  UpperCamelCase     ):







   """simple docstring"""
   # as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
   lowerCamelCase_      =  {
       # 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},
   }

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






  @slow
  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.run_trainer(
       eval_steps=2			,  max_len=128			,  model_name=UpperCamelCase			,  learning_rate=3e-4			,  num_train_epochs=10			,  distributed=UpperCamelCase			,  )

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

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

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






  @slow
  @require_bitsandbytes
  def   snake_case		(    self     ):







   """simple docstring"""
   from transformers.training_args import OptimizerNames

   def train_and_return_metrics(UpperCamelCase     ) -> Tuple[int, float]:
    lowerCamelCase_      =  "--skip_memory_metrics 0"

    lowerCamelCase_      =  self.run_trainer(
        max_len=128			,  model_name=UpperCamelCase			,  learning_rate=3e-4			,  num_train_epochs=1			,  optim=UpperCamelCase			,  distributed=UpperCamelCase			,  extra_args_str=UpperCamelCase			,  do_eval=UpperCamelCase			,  do_predict=UpperCamelCase			,  n_gpus_to_use=1			,  )

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

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

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

   lowerCamelCase_      =  gpu_alloc_mem_orig - gpu_alloc_mem_bnb

   lowerCamelCase_      =  gpu_peak_mem_orig + gpu_alloc_mem_orig
   lowerCamelCase_      =  gpu_peak_mem_bnb + gpu_alloc_mem_bnb
   lowerCamelCase_      =  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
   lowerCamelCase_      =  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(
       UpperCamelCase			,  UpperCamelCase			,  "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(
       UpperCamelCase			,  UpperCamelCase			,  "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(
       UpperCamelCase			,  UpperCamelCase			,  f'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''     )







  def   snake_case		(    self			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase = 3e-3			,  UpperCamelCase = "adafactor"			,  UpperCamelCase = False			,  UpperCamelCase = None			,  UpperCamelCase = 0			,  UpperCamelCase = True			,  UpperCamelCase = True			,  UpperCamelCase = True			,  UpperCamelCase = True			,  UpperCamelCase = None			,  ):







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

   lowerCamelCase_      =  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(UpperCamelCase     )}
        '''.split()

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

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

   return output_dir



