






"""simple docstring"""







# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.utils import logging


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


if is_torch_available():
 import torch
 from torch import nn
 from torch.utils.data import Dataset

 from transformers import Trainer





 class      UpperCamelCase_       (				a_     ):



   def __init__(					self       ,						snake_case__ = 1_01 )	->	int:



          """simple docstring"""





          UpperCAmelCase					=	length



   def __len__(					self )	->	str:



          """simple docstring"""





          return self.length



   def __getitem__(					self       ,						snake_case__ )	->	int:



          """simple docstring"""





          return i





 class      UpperCamelCase_       :



   def __call__(					self       ,						snake_case__ )	->	Dict:



          """simple docstring"""





          return {"input_ids": torch.tensor(snake_case__ ), "labels": torch.tensor(snake_case__ )}





 class      UpperCamelCase_       (				nn.Module     ):



   def __init__(					self )	->	Dict:



          """simple docstring"""





          super().__init__()
          # Add some (unused) params otherwise DDP will complain.
          UpperCAmelCase					=	nn.Linear(1_20       ,						80 )



   def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__=None )	->	str:



          """simple docstring"""





          if labels is not None:
                 return torch.tensor(0.0       ,						device=input_ids.device ), input_ids
          else:
                 return input_ids






def 					_lowerCAmelCase				(					):







       '''simple docstring'''
       UpperCAmelCase					=	HfArgumentParser((TrainingArguments,)						)
       sys.argv += ["--output_dir", "./examples"]
       UpperCAmelCase					=	parser.parse_args_into_dataclasses()[0]

       logger.warning(
           F'''Process rank: {training_args.local_rank}, device: {training_args.device}, '''
           F'''tpu_num_cores: {training_args.tpu_num_cores}'''		,					)

       # Essentially, what we want to verify in the distributed case is
       # that we get all samples back, in the right order.
       # (this is crucial for prediction for instance)
       for dataset_length in [1001, 256, 15]:
              UpperCAmelCase					=	DummyDataset(lowerCAmelCase						)

              def compute_metrics(lowerCAmelCase						) -> Dict:
                     UpperCAmelCase					=	list(range(len(lowerCAmelCase						)						)						)
                     UpperCAmelCase					=	p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
                     return {"success": success}

              UpperCAmelCase					=	Trainer(
                  model=DummyModel()		,					args=lowerCAmelCase		,					data_collator=DummyDataCollator()		,					eval_dataset=lowerCAmelCase		,					compute_metrics=lowerCAmelCase		,					)
              UpperCAmelCase					=	trainer.evaluate()
              logger.info(lowerCAmelCase						)
              if metrics["eval_success"] is not True:
                     logger.error(lowerCAmelCase						)
                     exit(1						)

              UpperCAmelCase					=	trainer.predict(lowerCAmelCase						)
              logger.info(p.metrics						)
              if p.metrics["test_success"] is not True:
                     logger.error(p.metrics						)
                     exit(1						)

              UpperCAmelCase					=	2

              UpperCAmelCase					=	trainer.evaluate()
              logger.info(lowerCAmelCase						)
              if metrics["eval_success"] is not True:
                     logger.error(lowerCAmelCase						)
                     exit(1						)

              UpperCAmelCase					=	trainer.predict(lowerCAmelCase						)
              logger.info(p.metrics						)
              if p.metrics["test_success"] is not True:
                     logger.error(p.metrics						)
                     exit(1						)

              UpperCAmelCase					=	None

       logger.info("""🔥 All distributed tests successful"""						)






def 					_lowerCAmelCase				(					lowerCAmelCase						):







       '''simple docstring'''
       # For xla_spawn (TPUs)
       main()


if __name__ == "__main__":
 main()








"""simple docstring"""







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

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
 from .tokenization_mbartaa import MBartaaTokenizer
else:
 lowerCAmelCase_		:			int				   =					None


lowerCAmelCase_		:			Tuple				   =					logging.get_logger(__name__)

lowerCAmelCase_		:			List[Any]				   =					{'''vocab_file''': '''sentencepiece.bpe.model''', '''tokenizer_file''': '''tokenizer.json'''}

lowerCAmelCase_		:			List[Any]				   =					{
    '''vocab_file''': {
        '''facebook/mbart-large-50-one-to-many-mmt''': (
            '''https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'''
        ),
    },
    '''tokenizer_file''': {
        '''facebook/mbart-large-50-one-to-many-mmt''': (
            '''https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/tokenizer.json'''
        ),
    },
}

lowerCAmelCase_		:			int				   =					{
    '''facebook/mbart-large-50-one-to-many-mmt''': 1_0_2_4,
}

# fmt: off
lowerCAmelCase_		:			List[Any]				   =					['''ar_AR''', '''cs_CZ''', '''de_DE''', '''en_XX''', '''es_XX''', '''et_EE''', '''fi_FI''', '''fr_XX''', '''gu_IN''', '''hi_IN''', '''it_IT''', '''ja_XX''', '''kk_KZ''', '''ko_KR''', '''lt_LT''', '''lv_LV''', '''my_MM''', '''ne_NP''', '''nl_XX''', '''ro_RO''', '''ru_RU''', '''si_LK''', '''tr_TR''', '''vi_VN''', '''zh_CN''', '''af_ZA''', '''az_AZ''', '''bn_IN''', '''fa_IR''', '''he_IL''', '''hr_HR''', '''id_ID''', '''ka_GE''', '''km_KH''', '''mk_MK''', '''ml_IN''', '''mn_MN''', '''mr_IN''', '''pl_PL''', '''ps_AF''', '''pt_XX''', '''sv_SE''', '''sw_KE''', '''ta_IN''', '''te_IN''', '''th_TH''', '''tl_XX''', '''uk_UA''', '''ur_PK''', '''xh_ZA''', '''gl_ES''', '''sl_SI''']





class      UpperCamelCase_       (				a_     ):
  _A				:    Union[str, Any] =		VOCAB_FILES_NAMES
  _A				:    Optional[int] =		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
  _A				:    Tuple =		PRETRAINED_VOCAB_FILES_MAP
  _A				:    Dict =		['input_ids', 'attention_mask']
  _A				:    Any =		MBartaaTokenizer

  _A				:    List[int] =		[]
  _A				:    List[int] =		[]



  def __init__(					self       ,						snake_case__=None       ,						snake_case__=None       ,						snake_case__=None       ,						snake_case__=None       ,						snake_case__="</s>"       ,						snake_case__="</s>"       ,						snake_case__="<s>"       ,						snake_case__="<unk>"       ,						snake_case__="<pad>"       ,						snake_case__="<mask>"       ,						**snake_case__       ,						)	->	Optional[Any]:



         """simple docstring"""





         UpperCAmelCase					=	AddedToken(snake_case__       ,						lstrip=snake_case__       ,						rstrip=snake_case__ ) if isinstance(snake_case__       ,						snake_case__ ) else mask_token

         UpperCAmelCase					=	kwargs.get("""additional_special_tokens"""       ,						[] )
         kwargs["additional_special_tokens"] += [
             code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
         ]

         super().__init__(
             snake_case__       ,						src_lang=snake_case__       ,						tgt_lang=snake_case__       ,						tokenizer_file=snake_case__       ,						eos_token=snake_case__       ,						sep_token=snake_case__       ,						cls_token=snake_case__       ,						unk_token=snake_case__       ,						pad_token=snake_case__       ,						mask_token=snake_case__       ,						**snake_case__       ,						)

         UpperCAmelCase					=	vocab_file
         UpperCAmelCase					=	False if not self.vocab_file else True

         UpperCAmelCase					=	{
             lang_code: self.convert_tokens_to_ids(snake_case__ ) for lang_code in FAIRSEQ_LANGUAGE_CODES
         }

         UpperCAmelCase					=	src_lang if src_lang is not None else """en_XX"""
         UpperCAmelCase					=	tgt_lang
         UpperCAmelCase					=	self.lang_code_to_id[self._src_lang]
         self.set_src_lang_special_tokens(self._src_lang )



  @property
  def 						UpperCamelCase_ (					self )	->	str:



         """simple docstring"""





         return self._src_lang



  @src_lang.setter
  def 						UpperCamelCase_ (					self       ,						snake_case__ )	->	None:



         """simple docstring"""





         UpperCAmelCase					=	new_src_lang
         self.set_src_lang_special_tokens(self._src_lang )



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ = None )	->	List[int]:



         """simple docstring"""





         if token_ids_a is None:
                return self.prefix_tokens + token_ids_a + self.suffix_tokens
         # We don't expect to process pairs, but leave the pair logic for API consistency
         return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ = "en_XX"       ,						snake_case__ = None       ,						snake_case__ = "ro_RO"       ,						**snake_case__       ,						)	->	BatchEncoding:



         """simple docstring"""





         UpperCAmelCase					=	src_lang
         UpperCAmelCase					=	tgt_lang
         return super().prepare_seqaseq_batch(snake_case__       ,						snake_case__       ,						**snake_case__ )



  def 						UpperCamelCase_ (					self )	->	int:



         """simple docstring"""





         return self.set_src_lang_special_tokens(self.src_lang )



  def 						UpperCamelCase_ (					self )	->	Optional[int]:



         """simple docstring"""





         return self.set_tgt_lang_special_tokens(self.tgt_lang )



  def 						UpperCamelCase_ (					self       ,						snake_case__ )	->	None:



         """simple docstring"""





         UpperCAmelCase					=	self.convert_tokens_to_ids(snake_case__ )
         UpperCAmelCase					=	[self.cur_lang_code_id]
         UpperCAmelCase					=	[self.eos_token_id]

         UpperCAmelCase					=	self.convert_ids_to_tokens(self.prefix_tokens )
         UpperCAmelCase					=	self.convert_ids_to_tokens(self.suffix_tokens )

         UpperCAmelCase					=	processors.TemplateProcessing(
             single=prefix_tokens_str + ["""$A"""] + suffix_tokens_str       ,						pair=prefix_tokens_str + ["""$A""", """$B"""] + suffix_tokens_str       ,						special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str       ,						self.prefix_tokens + self.suffix_tokens ) )       ,						)



  def 						UpperCamelCase_ (					self       ,						snake_case__ )	->	None:



         """simple docstring"""





         UpperCAmelCase					=	self.convert_tokens_to_ids(snake_case__ )
         UpperCAmelCase					=	[self.cur_lang_code_id]
         UpperCAmelCase					=	[self.eos_token_id]

         UpperCAmelCase					=	self.convert_ids_to_tokens(self.prefix_tokens )
         UpperCAmelCase					=	self.convert_ids_to_tokens(self.suffix_tokens )

         UpperCAmelCase					=	processors.TemplateProcessing(
             single=prefix_tokens_str + ["""$A"""] + suffix_tokens_str       ,						pair=prefix_tokens_str + ["""$A""", """$B"""] + suffix_tokens_str       ,						special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str       ,						self.prefix_tokens + self.suffix_tokens ) )       ,						)



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__       ,						snake_case__       ,						snake_case__       ,						**snake_case__ )	->	Dict:



         """simple docstring"""





         if src_lang is None or tgt_lang is None:
                raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model""" )
         UpperCAmelCase					=	src_lang
         UpperCAmelCase					=	self(snake_case__       ,						add_special_tokens=snake_case__       ,						return_tensors=snake_case__       ,						**snake_case__ )
         UpperCAmelCase					=	self.convert_tokens_to_ids(snake_case__ )
         UpperCAmelCase					=	tgt_lang_id
         return inputs



  def 						UpperCamelCase_ (					self       ,						snake_case__       ,						snake_case__ = None )	->	Tuple[str]:



         """simple docstring"""





         if not self.can_save_slow_tokenizer:
                raise ValueError(
                    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
                    """tokenizer.""" )

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

         if os.path.abspath(self.vocab_file ) != os.path.abspath(snake_case__ ):
                copyfile(self.vocab_file       ,						snake_case__ )

         return (out_vocab_file,)
