

"""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:
    _UpperCamelCase		    =						None


_UpperCamelCase		    =						logging.get_logger(__name__)

_UpperCamelCase		    =						{"""vocab_file""": """sentencepiece.bpe.model""", """tokenizer_file""": """tokenizer.json"""}

_UpperCamelCase		    =						{
    """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"""
        ),
    },
}

_UpperCamelCase		    =						{
    """facebook/mbart-large-50-one-to-many-mmt""": 1_0_2_4,
}

# fmt: off
_UpperCamelCase		    =						["""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  __a				(	__magic_name__					):
       __UpperCamelCase				:			Optional[int]          =	VOCAB_FILES_NAMES
       __UpperCamelCase				:			int          =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       __UpperCamelCase				:			Optional[int]          =	PRETRAINED_VOCAB_FILES_MAP
       __UpperCamelCase				:			List[Any]          =	['input_ids', 'attention_mask']
       __UpperCamelCase				:			List[str]          =	MBartaaTokenizer

       __UpperCamelCase				:			List[int]          =	[]
       __UpperCamelCase				:			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 ,				):





        """simple docstring"""

        lowerCAmelCase__      :							Union[str, Any]    =							AddedToken(snake_case ,				lstrip=snake_case ,				rstrip=snake_case      ) if isinstance(snake_case ,				snake_case      ) else mask_token

        lowerCAmelCase__      :							Tuple    =							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 ,				)

        lowerCAmelCase__      :							Tuple    =							vocab_file
        lowerCAmelCase__      :							Tuple    =							False if not self.vocab_file else True

        lowerCAmelCase__      :							str    =							{
            lang_code: self.convert_tokens_to_ids(snake_case      ) for lang_code in FAIRSEQ_LANGUAGE_CODES
        }

        lowerCAmelCase__      :							Optional[Any]    =							src_lang if src_lang is not None else "en_XX"
        lowerCAmelCase__      :							Dict    =							tgt_lang
        lowerCAmelCase__      :							Dict    =							self.lang_code_to_id[self._src_lang]
        self.set_src_lang_special_tokens(self._src_lang      )






       @property
       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        return self._src_lang






       @src_lang.setter
       def  SCREAMING_SNAKE_CASE_  (      self ,				snake_case      ):





        """simple docstring"""

        lowerCAmelCase__      :							Optional[Any]    =							new_src_lang
        self.set_src_lang_special_tokens(self._src_lang      )






       def  SCREAMING_SNAKE_CASE_  (      self ,				snake_case ,				snake_case = None      ):





        """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  SCREAMING_SNAKE_CASE_  (      self ,				snake_case ,				snake_case = "en_XX" ,				snake_case = None ,				snake_case = "ro_RO" ,				**snake_case ,				):





        """simple docstring"""

        lowerCAmelCase__      :							int    =							src_lang
        lowerCAmelCase__      :							Dict    =							tgt_lang
        return super().prepare_seqaseq_batch(snake_case ,				snake_case ,				**snake_case      )






       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        return self.set_src_lang_special_tokens(self.src_lang      )






       def  SCREAMING_SNAKE_CASE_  (      self      ):





        """simple docstring"""

        return self.set_tgt_lang_special_tokens(self.tgt_lang      )






       def  SCREAMING_SNAKE_CASE_  (      self ,				snake_case      ):





        """simple docstring"""

        lowerCAmelCase__      :							Any    =							self.convert_tokens_to_ids(snake_case      )
        lowerCAmelCase__      :							Dict    =							[self.cur_lang_code_id]
        lowerCAmelCase__      :							List[str]    =							[self.eos_token_id]

        lowerCAmelCase__      :							List[str]    =							self.convert_ids_to_tokens(self.prefix_tokens      )
        lowerCAmelCase__      :							Dict    =							self.convert_ids_to_tokens(self.suffix_tokens      )

        lowerCAmelCase__      :							Dict    =							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  SCREAMING_SNAKE_CASE_  (      self ,				snake_case      ):





        """simple docstring"""

        lowerCAmelCase__      :							Any    =							self.convert_tokens_to_ids(snake_case      )
        lowerCAmelCase__      :							List[Any]    =							[self.cur_lang_code_id]
        lowerCAmelCase__      :							int    =							[self.eos_token_id]

        lowerCAmelCase__      :							Optional[int]    =							self.convert_ids_to_tokens(self.prefix_tokens      )
        lowerCAmelCase__      :							Dict    =							self.convert_ids_to_tokens(self.suffix_tokens      )

        lowerCAmelCase__      :							List[Any]    =							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  SCREAMING_SNAKE_CASE_  (      self ,				snake_case ,				snake_case ,				snake_case ,				snake_case ,				**snake_case      ):





        """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"      )
        lowerCAmelCase__      :							Union[str, Any]    =							src_lang
        lowerCAmelCase__      :							Tuple    =							self(snake_case ,				add_special_tokens=snake_case ,				return_tensors=snake_case ,				**snake_case      )
        lowerCAmelCase__      :							Union[str, Any]    =							self.convert_tokens_to_ids(snake_case      )
        lowerCAmelCase__      :							Optional[Any]    =							tgt_lang_id
        return inputs






       def  SCREAMING_SNAKE_CASE_  (      self ,				snake_case ,				snake_case = None      ):





        """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
        lowerCAmelCase__      :							str    =							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,)









"""simple docstring"""

import logging
import os
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import numpy as np
import torch
from utils_hans import HansDataset, InputFeatures, hans_processors, hans_tasks_num_labels

import transformers
from transformers import (
    AutoConfig,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import is_main_process


_UpperCamelCase		    =						logging.getLogger(__name__)






@dataclass
class  __a				:





       """simple docstring"""


       __UpperCamelCase				:			str          =	field(
           metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}					)
       __UpperCamelCase				:			Optional[str]          =	field(
           default=__magic_name__					,  metadata={'help': 'Pretrained config name or path if not the same as model_name'}					)
       __UpperCamelCase				:			Optional[str]          =	field(
           default=__magic_name__					,  metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'}					)
       __UpperCamelCase				:			Optional[str]          =	field(
           default=__magic_name__					,  metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}					,  )






@dataclass
class  __a				:





       """simple docstring"""


       __UpperCamelCase				:			str          =	field(
           metadata={'help': 'The name of the task to train selected in the list: ' + ', '.join(hans_processors.keys()					)}					)
       __UpperCamelCase				:			str          =	field(
           metadata={'help': 'The input data dir. Should contain the .tsv files (or other data files) for the task.'}					)
       __UpperCamelCase				:			int          =	field(
           default=128					,  metadata={
               'help': (
                   'The maximum total input sequence length after tokenization. Sequences longer '
                   'than this will be truncated, sequences shorter will be padded.'
               )
           }					,  )
       __UpperCamelCase				:			bool          =	field(
           default=__magic_name__					,  metadata={'help': 'Overwrite the cached training and evaluation sets'}					)




def 			SCREAMING_SNAKE_CASE       (   lowercase__      ) ->	Dict[str, torch.Tensor]:
 lowerCAmelCase__      :							List[Any]    =							default_data_collator(lowercase__      )
 lowerCAmelCase__      :							Union[str, Any]    =							batch.pop("pairID"				,				lowercase__      )
 return batch



def 			SCREAMING_SNAKE_CASE       (   ) ->	int:
 # See all possible arguments in src/transformers/training_args.py
 # or by passing the --help flag to this script.
 # We now keep distinct sets of args, for a cleaner separation of concerns.

 lowerCAmelCase__      :							str    =							HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)      )
 lowerCAmelCase__ ,	lowerCAmelCase__ ,	lowerCAmelCase__      :							List[Any]    =							parser.parse_args_into_dataclasses()

 if (
     os.path.exists(training_args.output_dir      )
     and os.listdir(training_args.output_dir      )
     and training_args.do_train
     and not training_args.overwrite_output_dir
 ):
  raise ValueError(
      F"""Output directory ({training_args.output_dir}) already exists and is not empty. Use"""
      " --overwrite_output_dir to overcome."      )

 # Setup logging
 logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"				,				datefmt="%m/%d/%Y %H:%M:%S"				,				level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN				,				)
 logger.warning(
     "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s"				,				training_args.local_rank				,				training_args.device				,				training_args.n_gpu				,				bool(training_args.local_rank != -1      )				,				training_args.fpaa				,				)
 # Set the verbosity to info of the Transformers logger (on main process only):
 if is_main_process(training_args.local_rank      ):
  transformers.utils.logging.set_verbosity_info()
  transformers.utils.logging.enable_default_handler()
  transformers.utils.logging.enable_explicit_format()
 logger.info("Training/evaluation parameters %s"				,				lowercase__      )

 # Set seed
 set_seed(training_args.seed      )

 try:
  lowerCAmelCase__      :							Tuple    =							hans_tasks_num_labels[data_args.task_name]
 except KeyError:
  raise ValueError("Task not found: %s" % (data_args.task_name)      )

 # Load pretrained model and tokenizer
 #
 # Distributed training:
 # The .from_pretrained methods guarantee that only one local process can concurrently
 # download model & vocab.

 lowerCAmelCase__      :							int    =							AutoConfig.from_pretrained(
     model_args.config_name if model_args.config_name else model_args.model_name_or_path				,				num_labels=lowercase__				,				finetuning_task=data_args.task_name				,				cache_dir=model_args.cache_dir				,				)
 lowerCAmelCase__      :							Optional[int]    =							AutoTokenizer.from_pretrained(
     model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path				,				cache_dir=model_args.cache_dir				,				)
 lowerCAmelCase__      :							Union[str, Any]    =							AutoModelForSequenceClassification.from_pretrained(
     model_args.model_name_or_path				,				from_tf=bool(".ckpt" in model_args.model_name_or_path      )				,				config=lowercase__				,				cache_dir=model_args.cache_dir				,				)

 # Get datasets
 lowerCAmelCase__      :							Dict    =							(
     HansDataset(
         data_dir=data_args.data_dir				,				tokenizer=lowercase__				,				task=data_args.task_name				,				max_seq_length=data_args.max_seq_length				,				overwrite_cache=data_args.overwrite_cache				,				)
     if training_args.do_train
     else None
 )
 lowerCAmelCase__      :							Union[str, Any]    =							(
     HansDataset(
         data_dir=data_args.data_dir				,				tokenizer=lowercase__				,				task=data_args.task_name				,				max_seq_length=data_args.max_seq_length				,				overwrite_cache=data_args.overwrite_cache				,				evaluate=lowercase__				,				)
     if training_args.do_eval
     else None
 )

 # Initialize our Trainer
 lowerCAmelCase__      :							Any    =							Trainer(
     model=lowercase__				,				args=lowercase__				,				train_dataset=lowercase__				,				eval_dataset=lowercase__				,				data_collator=lowercase__				,				)

 # Training
 if training_args.do_train:
  trainer.train(
      model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path      ) else None      )
  trainer.save_model()
  # For convenience, we also re-save the tokenizer to the same directory,
  # so that you can share your model easily on huggingface.co/models =)
  if trainer.is_world_master():
   tokenizer.save_pretrained(training_args.output_dir      )

    # Evaluation
 if training_args.do_eval:
  logger.info("*** Evaluate ***"      )

  lowerCAmelCase__      :							Tuple    =							trainer.predict(lowercase__      )
  lowerCAmelCase__      :							Optional[Any]    =							output.predictions
  lowerCAmelCase__      :							Dict    =							np.argmax(lowercase__				,				axis=1      )

  lowerCAmelCase__      :							int    =							[ex.pairID for ex in eval_dataset]
  lowerCAmelCase__      :							Tuple    =							os.path.join(training_args.output_dir				,				"hans_predictions.txt"      )
  lowerCAmelCase__      :							Tuple    =							eval_dataset.get_labels()
  if trainer.is_world_master():
   with open(lowercase__				,				"w"      ) as writer:
    writer.write("pairID,gold_label\n"      )
    for pid, pred in zip(lowercase__				,				lowercase__      ):
     writer.write("ex" + str(lowercase__      ) + "," + label_list[int(lowercase__      )] + "\n"      )

  trainer._log(output.metrics      )



def 			SCREAMING_SNAKE_CASE       (   lowercase__      ) ->	Optional[Any]:
 # For xla_spawn (TPUs)
 main()


if __name__ == "__main__":
    main()






