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


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


set_seed(4_2)
UpperCAmelCase_           =				'''sshleifer/student_marian_en_ro_6_1'''
UpperCAmelCase_           =				'''sshleifer/tiny-mbart'''




@require_torch
class 			__magic_name__	(		__a     ):





   """simple docstring"""




   def 	lowerCAmelCase    (						self       :	Optional[int]		,	_lowercase       :	List[str]=False		,	_lowercase       :	str=None		,	_lowercase       :	List[Any]=True		,	_lowercase       :	List[Any]=True		,	_lowercase       :	Optional[Any]=True		,	_lowercase       :	Any=True		,	):



    """simple docstring"""



    _UpperCamelCase:  str    						=			self.run_trainer(
        eval_steps=1		,	max_len=12		,	model_name=_lowercase		,	num_train_epochs=1		,	distributed=_lowercase		,	extra_args_str=_lowercase		,	predict_with_generate=_lowercase		,	do_train=_lowercase		,	do_eval=_lowercase		,	do_predict=_lowercase		,	)
    _UpperCamelCase:  Optional[int]    						=			TrainerState.load_from_json(os.path.join(_lowercase		,	'''trainer_state.json'''				)				).log_history

    if not do_eval:
     return

    _UpperCamelCase:  Any    						=			[log for log in logs if '''eval_loss''' in log.keys()]

    _UpperCamelCase:  Optional[Any]    						=			eval_metrics[0]
    if predict_with_generate:
     assert "eval_bleu" in first_step_stats

     _UpperCamelCase:  str    						=			eval_metrics[-1]
     assert isinstance(last_step_stats['''eval_bleu''']		,	_lowercase				)
     assert not math.isnan(float(last_step_stats['''eval_loss''']				)				), "eval_loss must not be `nan`"



   @require_torch_non_multi_gpu
   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    self.run_seqaseq_quick()



   @require_torch_multi_gpu
   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    self.run_seqaseq_quick(distributed=_lowercase				)



   @require_torch_multi_gpu
   def 	lowerCAmelCase    (						self       :	Tuple				):



    """simple docstring"""



    self.run_seqaseq_quick(distributed=_lowercase				)



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



    """simple docstring"""



    self.run_seqaseq_quick(distributed=_lowercase		,	extra_args_str='''--sharded_ddp simple'''				)



   @unittest.skip('''Requires an update of the env running those tests'''				)
   @require_torch_multi_gpu
   @require_fairscale
   def 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    self.run_seqaseq_quick(distributed=_lowercase		,	extra_args_str='''--sharded_ddp simple --fp16'''				)



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



    """simple docstring"""



    self.run_seqaseq_quick(distributed=_lowercase		,	extra_args_str='''--sharded_ddp zero_dp_2'''		,	predict_with_generate=_lowercase				)



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



    """simple docstring"""



    self.run_seqaseq_quick(
        distributed=_lowercase		,	extra_args_str='''--sharded_ddp zero_dp_2 --fp16'''		,	predict_with_generate=_lowercase				)



   @require_apex
   @require_torch_gpu
   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    self.run_seqaseq_quick(distributed=_lowercase		,	extra_args_str='''--fp16 --fp16_backend=apex'''				)
    # test 2nd time - was getting eval_loss': nan'
    # to reproduce the problem set distributed=False
    self.run_seqaseq_quick(distributed=_lowercase		,	extra_args_str='''--fp16 --fp16_backend=apex'''				)



   @parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']				)
   @require_torch_multi_gpu
   def 	lowerCAmelCase    (						self       :	int		,	_lowercase       :	List[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			{
        # 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},
    }

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



   @slow
   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Union[str, Any]    						=			self.run_trainer(
        eval_steps=2		,	max_len=128		,	model_name=_lowercase		,	learning_rate=3E-4		,	num_train_epochs=10		,	distributed=_lowercase		,	)

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

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

    # test if do_predict saves generations and metrics
    _UpperCamelCase:  List[Any]    						=			os.listdir(_lowercase				)
    _UpperCamelCase:  List[Any]    						=			{os.path.basename(_lowercase				) for p in contents}
    assert "generated_predictions.txt" in contents
    assert "predict_results.json" in contents



   @slow
   @require_bitsandbytes
   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    from transformers.training_args import OptimizerNames

    def train_and_return_metrics(_lowercase       :	str				) -> Tuple[int, float]:
     _UpperCamelCase:  Union[str, Any]    						=			'''--skip_memory_metrics 0'''

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

     # Check metrics
     _UpperCamelCase:  str    						=			TrainerState.load_from_json(Path(_lowercase		,	'''trainer_state.json'''				)				).log_history
     _UpperCamelCase:  Dict    						=			int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**20				)
     _UpperCamelCase:  List[Any]    						=			int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**20				)

     _UpperCamelCase:  int    						=			logs[0]['''train_loss''']
     return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  Optional[int]    						=			train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value				)
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  str    						=			train_and_return_metrics(OptimizerNames.ADAMW_BNB.value				)

    _UpperCamelCase:  Optional[Any]    						=			gpu_alloc_mem_orig - gpu_alloc_mem_bnb

    _UpperCamelCase:  Tuple    						=			gpu_peak_mem_orig + gpu_alloc_mem_orig
    _UpperCamelCase:  Any    						=			gpu_peak_mem_bnb + gpu_alloc_mem_bnb
    _UpperCamelCase:  Dict    						=			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
    _UpperCamelCase:  str    						=			120

    # uncomment the following if this test starts failing - requires py38 for a new print feature
    # gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
    # print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
    # print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
    # print(f"{gpu_alloc_mem_diff=}MB")
    # print(f"{gpu_peak_mem_diff=}MB")
    # print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
    # print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

    self.assertGreater(
        _lowercase		,	_lowercase		,	'''should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'''
        f""" a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and"""
        f""" gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB"""		,	)

    self.assertGreater(
        _lowercase		,	_lowercase		,	'''should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'''
        f""" a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and"""
        f""" gpu_total_mem_bnb={gpu_total_mem_bnb}MB"""		,	)

    self.assertEqual(
        _lowercase		,	_lowercase		,	f"""loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}"""				)



   def 	lowerCAmelCase    (						self       :	List[Any]		,	_lowercase       :	int		,	_lowercase       :	str		,	_lowercase       :	int		,	_lowercase       :	float = 3E-3		,	_lowercase       :	str = "adafactor"		,	_lowercase       :	bool = False		,	_lowercase       :	str = None		,	_lowercase       :	int = 0		,	_lowercase       :	bool = True		,	_lowercase       :	bool = True		,	_lowercase       :	bool = True		,	_lowercase       :	bool = True		,	_lowercase       :	int = None		,	):



    """simple docstring"""



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

    _UpperCamelCase:  int    						=			f"""
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(_lowercase				)}
        """.split()

    _UpperCamelCase:  Tuple    						=			'''
            --do_predict
        '''.split()

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

    return output_dir
import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
  from transformers.models.bert.modeling_flax_bert import (
      FlaxBertForMaskedLM,
      FlaxBertForMultipleChoice,
      FlaxBertForNextSentencePrediction,
      FlaxBertForPreTraining,
      FlaxBertForQuestionAnswering,
      FlaxBertForSequenceClassification,
      FlaxBertForTokenClassification,
      FlaxBertModel,
  )




class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   def __init__(						self       :	Tuple		,	_lowercase       :	List[str]		,	_lowercase       :	Dict=13		,	_lowercase       :	Any=7		,	_lowercase       :	List[str]=True		,	_lowercase       :	Any=True		,	_lowercase       :	Optional[int]=True		,	_lowercase       :	int=True		,	_lowercase       :	str=99		,	_lowercase       :	Union[str, Any]=32		,	_lowercase       :	Any=5		,	_lowercase       :	Dict=4		,	_lowercase       :	Tuple=37		,	_lowercase       :	Optional[Any]="gelu"		,	_lowercase       :	Optional[int]=0.1		,	_lowercase       :	Any=0.1		,	_lowercase       :	Union[str, Any]=512		,	_lowercase       :	Dict=16		,	_lowercase       :	int=2		,	_lowercase       :	Union[str, Any]=0.02		,	_lowercase       :	Dict=4		,	):



    """simple docstring"""



    _UpperCamelCase:  List[str]    						=			parent
    _UpperCamelCase:  Optional[int]    						=			batch_size
    _UpperCamelCase:  Any    						=			seq_length
    _UpperCamelCase:  List[str]    						=			is_training
    _UpperCamelCase:  Union[str, Any]    						=			use_attention_mask
    _UpperCamelCase:  Union[str, Any]    						=			use_token_type_ids
    _UpperCamelCase:  str    						=			use_labels
    _UpperCamelCase:  List[str]    						=			vocab_size
    _UpperCamelCase:  Dict    						=			hidden_size
    _UpperCamelCase:  List[Any]    						=			num_hidden_layers
    _UpperCamelCase:  List[Any]    						=			num_attention_heads
    _UpperCamelCase:  Tuple    						=			intermediate_size
    _UpperCamelCase:  Optional[Any]    						=			hidden_act
    _UpperCamelCase:  Any    						=			hidden_dropout_prob
    _UpperCamelCase:  int    						=			attention_probs_dropout_prob
    _UpperCamelCase:  Optional[int]    						=			max_position_embeddings
    _UpperCamelCase:  List[Any]    						=			type_vocab_size
    _UpperCamelCase:  List[Any]    						=			type_sequence_label_size
    _UpperCamelCase:  List[Any]    						=			initializer_range
    _UpperCamelCase:  Optional[int]    						=			num_choices



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  int    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size				)

    _UpperCamelCase:  Optional[int]    						=			None
    if self.use_attention_mask:
     _UpperCamelCase:  List[str]    						=			random_attention_mask([self.batch_size, self.seq_length]				)

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

    _UpperCamelCase:  str    						=			BertConfig(
        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		,	is_decoder=_lowercase		,	initializer_range=self.initializer_range		,	)

    return config, input_ids, token_type_ids, attention_mask



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			self.prepare_config_and_inputs()
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  Any    						=			config_and_inputs
    _UpperCamelCase:  List[Any]    						=			{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
    return config, inputs_dict



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			self.prepare_config_and_inputs()
    _UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase						,				_UpperCamelCase:  str    						=			config_and_inputs

    _UpperCamelCase:  int    						=			True
    _UpperCamelCase:  Optional[Any]    						=			floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
    _UpperCamelCase:  Optional[Any]    						=			ids_tensor([self.batch_size, self.seq_length]		,	vocab_size=2				)

    return (
        config,
        input_ids,
        attention_mask,
        encoder_hidden_states,
        encoder_attention_mask,
    )




@require_flax
class 			__magic_name__	(		__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Dict	      = True

   lowerCAmelCase   :     List[Any]	      = (
       (
           FlaxBertModel,
           FlaxBertForPreTraining,
           FlaxBertForMaskedLM,
           FlaxBertForMultipleChoice,
           FlaxBertForQuestionAnswering,
           FlaxBertForNextSentencePrediction,
           FlaxBertForSequenceClassification,
           FlaxBertForTokenClassification,
           FlaxBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			FlaxBertModelTester(self				)



   @slow
   def 	lowerCAmelCase    (						self       :	Tuple				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			FlaxBertModel.from_pretrained('''bert-base-cased'''				)
    _UpperCamelCase:  int    						=			model(np.ones((1, 1)				)				)
    self.assertIsNotNone(_lowercase				)