


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


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      lowerCamelCase_  (      __a ):






      def lowercase_				(   self							:       Dict		,		_A							:       int=False		,		_A							:       Optional[Any]=None		,		_A							:       Optional[Any]=True		,		_A							:       Dict=True		,		_A							:       Dict=True		,		_A							:       List[Any]=True		,		):


          '''simple docstring'''


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

          if not do_eval:
              return

          UpperCAmelCase__     :						Tuple						=							[log for log in logs if '''eval_loss''' in log.keys()]

          UpperCAmelCase__     :						List[str]						=							eval_metrics[0]
          if predict_with_generate:
              assert "eval_bleu" in first_step_stats

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






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


          '''simple docstring'''


          self.run_seqaseq_quick()






      @require_torch_multi_gpu
      def lowercase_				(   self							:       Optional[int]					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A					)






      @require_torch_multi_gpu
      def lowercase_				(   self							:       str					):


          '''simple docstring'''


          self.run_seqaseq_quick(distributed=_A					)






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


          '''simple docstring'''


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






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


          '''simple docstring'''


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






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


          '''simple docstring'''


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






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


          '''simple docstring'''


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






      @require_apex
      @require_torch_gpu
      def lowercase_				(   self							:       List[str]					):


          '''simple docstring'''


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






      @parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']					)
      @require_torch_multi_gpu
      def lowercase_				(   self							:       str		,		_A							:       Dict					):


          '''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__     :						int						=							experiments[experiment_id]
          UpperCAmelCase__     :						int						=							{'''distributed''': True, '''predict_with_generate''': False, '''do_eval''': False, '''do_predict''': False}
          UpperCAmelCase__     :						Optional[int]						=							'''Running training'''
          with CaptureStderr() as cl:
              self.run_seqaseq_quick(**_A		,		extra_args_str=data['''extra_args_str''']					)
          UpperCAmelCase__     :						List[Any]						=							len(re.findall(_A		,		cl.err					)					)
          self.assertEqual(_A		,		data['''n_matches''']					)






      @slow
      def lowercase_				(   self							:       int					):


          '''simple docstring'''


          UpperCAmelCase__     :						List[Any]						=							self.run_trainer(
              eval_steps=2		,		max_len=128		,		model_name=_A		,		learning_rate=3e-4		,		num_train_epochs=10		,		distributed=_A		,		)

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

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

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






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


          '''simple docstring'''


          from transformers.training_args import OptimizerNames

          def train_and_return_metrics(_A							:       str					) -> Tuple[int, float]:
              UpperCAmelCase__     :						Any						=							'''--skip_memory_metrics 0'''

              UpperCAmelCase__     :						List[Any]						=							self.run_trainer(
                  max_len=128		,		model_name=_A		,		learning_rate=3e-4		,		num_train_epochs=1		,		optim=_A		,		distributed=_A		,		extra_args_str=_A		,		do_eval=_A		,		do_predict=_A		,		n_gpus_to_use=1		,		)

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

              UpperCAmelCase__     :						Any						=							logs[0]['''train_loss''']
              return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

          UpperCAmelCase__		,		UpperCAmelCase__		,		UpperCAmelCase__     :						List[Any]						=							train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value					)
          UpperCAmelCase__		,		UpperCAmelCase__		,		UpperCAmelCase__     :						Optional[int]						=							train_and_return_metrics(OptimizerNames.ADAMW_BNB.value					)

          UpperCAmelCase__     :						Optional[Any]						=							gpu_alloc_mem_orig - gpu_alloc_mem_bnb

          UpperCAmelCase__     :						Union[str, Any]						=							gpu_peak_mem_orig + gpu_alloc_mem_orig
          UpperCAmelCase__     :						Optional[Any]						=							gpu_peak_mem_bnb + gpu_alloc_mem_bnb
          UpperCAmelCase__     :						List[str]						=							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__     :						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(
              _A		,		_A		,		'''should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'''
              f""" a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and"""
              f""" gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB"""		,		)

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

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






      def lowercase_				(   self							:       Optional[int]		,		_A							:       int		,		_A							:       str		,		_A							:       int		,		_A							:       float = 3e-3		,		_A							:       str = "adafactor"		,		_A							:       bool = False		,		_A							:       str = None		,		_A							:       int = 0		,		_A							:       bool = True		,		_A							:       bool = True		,		_A							:       bool = True		,		_A							:       bool = True		,		_A							:       int = None		,		):


          '''simple docstring'''


          UpperCAmelCase__     :						str						=							self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
          UpperCAmelCase__     :						List[Any]						=							self.get_auto_remove_tmp_dir()
          UpperCAmelCase__     :						Optional[int]						=							f"""
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(_A					)}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(_A					)}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        """.split()

          UpperCAmelCase__     :						List[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(_A					)}
        """.split()

          UpperCAmelCase__     :						Dict						=							'''
            --do_predict
        '''.split()

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

          return output_dir






import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__UpperCamelCase   :			Union[str, Any]          = logging.get_logger(__name__)

__UpperCamelCase   :			Optional[Any]          = "▁"

__UpperCamelCase   :			Any          = {"vocab_file": "sentencepiece.bpe.model"}

__UpperCamelCase   :			int          = {
    "vocab_file": {
        "facebook/xglm-564M": "https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model",
    }
}

__UpperCamelCase   :			Dict          = {
    "facebook/xglm-564M": 2_048,
}

class 						__lowerCAmelCase      (      __magic_name__					):
							UpperCamelCase__   	=    VOCAB_FILES_NAMES
							UpperCamelCase__   	=    PRETRAINED_VOCAB_FILES_MAP
							UpperCamelCase__   	=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							UpperCamelCase__   	=    ['''input_ids''', '''attention_mask''']
							def __init__(		self		:Optional[Any]						, __magic_name__		:List[str]						, __magic_name__		:Tuple="<s>"						, __magic_name__		:List[Any]="</s>"						, __magic_name__		:List[str]="</s>"						, __magic_name__		:Any="<s>"						, __magic_name__		:Tuple="<unk>"						, __magic_name__		:int="<pad>"						, __magic_name__		:Optional[Dict[str, Any]] = None						, **__magic_name__		:Optional[Any]						, ):







								'''simple docstring'''







								a									=				{} if sp_model_kwargs is None else sp_model_kwargs

								# Compatibility with the original tokenizer
								a									=				7
								a									=				[F'<madeupword{i}>' for i in range(self.num_madeup_words							)]

								a									=				kwargs.get("""additional_special_tokens"""						, []							)
								kwargs["additional_special_tokens"] += [
								    word for word in madeup_words if word not in kwargs["additional_special_tokens"]
								]

								super().__init__(
								    bos_token=__magic_name__						, eos_token=__magic_name__						, unk_token=__magic_name__						, sep_token=__magic_name__						, cls_token=__magic_name__						, pad_token=__magic_name__						, sp_model_kwargs=self.sp_model_kwargs						, **__magic_name__						, )

								a									=				spm.SentencePieceProcessor(**self.sp_model_kwargs							)
								self.sp_model.Load(str(__magic_name__							)							)
								a									=				vocab_file

								# Original fairseq vocab and spm vocab must be "aligned":
								# Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
								# -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
								# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
								# spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

								# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
								a									=				1

								# Mimic fairseq token-to-id alignment for the first 4 token
								a									=				{"""<s>""": 0, """<pad>""": 1, """</s>""": 2, """<unk>""": 3}

								a									=				len(self.sp_model							)
								a									=				{F'<madeupword{i}>': sp_size + i + self.fairseq_offset for i in range(self.num_madeup_words							)}
								self.fairseq_tokens_to_ids.update(__magic_name__							)

								a									=				{v: k for k, v in self.fairseq_tokens_to_ids.items()}
							def __getstate__(		self		:List[Any]							):







								'''simple docstring'''







								a									=				self.__dict__.copy()
								a									=				None
								a									=				self.sp_model.serialized_model_proto()
								return state
							def __setstate__(		self		:Any						, __magic_name__		:Any							):







								'''simple docstring'''







								a									=				d

								# for backward compatibility
								if not hasattr(self						, """sp_model_kwargs"""							):
									a									=				{}

								a									=				spm.SentencePieceProcessor(**self.sp_model_kwargs							)
								self.sp_model.LoadFromSerializedProto(self.sp_model_proto							)
							def  lowerCamelCase__				(		self		:Any						, __magic_name__		:List[int]						, __magic_name__		:Optional[List[int]] = None							):







								'''simple docstring'''







								if token_ids_a is None:
									return [self.sep_token_id] + token_ids_a
								a									=				[self.sep_token_id]
								return sep + token_ids_a + sep + sep + token_ids_a
							def  lowerCamelCase__				(		self		:Any						, __magic_name__		:List[int]						, __magic_name__		:Optional[List[int]] = None						, __magic_name__		:bool = False							):







								'''simple docstring'''







								if already_has_special_tokens:
									return super().get_special_tokens_mask(
									    token_ids_a=__magic_name__						, token_ids_a=__magic_name__						, already_has_special_tokens=__magic_name__							)

								if token_ids_a is None:
									return [1] + ([0] * len(__magic_name__							))
								return [1] + ([0] * len(__magic_name__							)) + [1, 1] + ([0] * len(__magic_name__							))
							def  lowerCamelCase__				(		self		:Optional[Any]						, __magic_name__		:List[int]						, __magic_name__		:Optional[List[int]] = None							):







								'''simple docstring'''







								a									=				[self.sep_token_id]

								if token_ids_a is None:
									return len(sep + token_ids_a							) * [0]
								return len(sep + token_ids_a + sep + sep + token_ids_a							) * [0]
							@property
							def  lowerCamelCase__				(		self		:str							):







								'''simple docstring'''







								return len(self.sp_model							) + self.fairseq_offset + self.num_madeup_words
							def  lowerCamelCase__				(		self		:Dict							):







								'''simple docstring'''







								a									=				{self.convert_ids_to_tokens(__magic_name__							): i for i in range(self.vocab_size							)}
								vocab.update(self.added_tokens_encoder							)
								return vocab
							def  lowerCamelCase__				(		self		:Any						, __magic_name__		:str							):







								'''simple docstring'''







								return self.sp_model.encode(__magic_name__						, out_type=__magic_name__							)
							def  lowerCamelCase__				(		self		:List[str]						, __magic_name__		:Any							):







								'''simple docstring'''







								if token in self.fairseq_tokens_to_ids:
									return self.fairseq_tokens_to_ids[token]
								a									=				self.sp_model.PieceToId(__magic_name__							)

								# Need to return unknown token if the SP model returned 0
								return spm_id + self.fairseq_offset if spm_id else self.unk_token_id
							def  lowerCamelCase__				(		self		:List[str]						, __magic_name__		:Dict							):







								'''simple docstring'''







								if index in self.fairseq_ids_to_tokens:
									return self.fairseq_ids_to_tokens[index]
								return self.sp_model.IdToPiece(index - self.fairseq_offset							)
							def  lowerCamelCase__				(		self		:List[str]						, __magic_name__		:List[str]							):







								'''simple docstring'''







								a									=				"""""".join(__magic_name__							).replace(__magic_name__						, """ """							).strip()
								return out_string





							def  lowerCamelCase__				(		self		:List[Any]						, __magic_name__		:str						, __magic_name__		:Optional[str] = None							):







								'''simple docstring'''







								if not os.path.isdir(__magic_name__							):
									logger.error(F'Vocabulary path ({save_directory}) should be a directory'							)
									return
								a									=				os.path.join(
								    __magic_name__						, (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]							)

								if os.path.abspath(self.vocab_file							) != os.path.abspath(__magic_name__							) and os.path.isfile(self.vocab_file							):
									copyfile(self.vocab_file						, __magic_name__							)
								elif not os.path.isfile(self.vocab_file							):
									with open(__magic_name__						, """wb"""							) as fi:
										a									=				self.sp_model.serialized_model_proto()
										fi.write(__magic_name__							)

								return (out_vocab_file,)

