
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
      from fairscale.optim import OSS


lowercase__          =   logging.get_logger(__name__)

lowercase__          =   {
    "linear": get_linear_schedule_with_warmup,
    "cosine": get_cosine_schedule_with_warmup,
    "cosine_w_restarts": get_cosine_with_hard_restarts_schedule_with_warmup,
    "polynomial": get_polynomial_decay_schedule_with_warmup,
    "constant": get_constant_schedule,
    "constant_w_warmup": get_constant_schedule_with_warmup,
}



class 		UpperCamelCase__  (	__lowerCAmelCase   ):
   def __init__(self			: Dict ,    snake_case_			: str=None ,    snake_case_			: List[str]=None ,    *snake_case_			: Union[str, Any] ,    **snake_case_			: Tuple       ):
    super().__init__(*lowerCamelCase__ ,    **lowerCamelCase__       )

    if config is None:
     assert isinstance(self.model ,    lowerCamelCase__       ), (
         "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
         f" {self.model.__class__}"
     )
     __a			:  Union[str, Any]		    =  self.model.config
    else:
     __a			:  List[Any]		    =  config

    __a			:  Dict		    =  data_args
    __a			:  str		    =  self.config.tgt_vocab_size if isinstance(self.config ,    lowerCamelCase__       ) else self.config.vocab_size

    if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
     assert self.config.pad_token_id is not None, (
         "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
         " calculation or doing label smoothing."
     )

    if self.config.pad_token_id is None and self.config.eos_token_id is not None:
     logger.warning(
         f"The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for"
         ''' padding..'''       )

    if self.args.label_smoothing == 0:
     __a			:  Optional[Any]		    =  torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id       )
    else:
     # dynamically import label_smoothed_nll_loss
     from utils import label_smoothed_nll_loss

     __a			:  Optional[Any]		    =  label_smoothed_nll_loss
   def 				lowerCAmelCase      (self			: List[Any] ,    snake_case_			: int       ):






    if self.optimizer is None:
     __a			:  List[str]		    =  ['''bias''', '''LayerNorm.weight''']
     __a			:  Optional[Any]		    =  [
         {
             '''params''': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay       )],
             '''weight_decay''': self.args.weight_decay,
         },
         {
             '''params''': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay       )],
             '''weight_decay''': 0.0,
         },
     ]
     __a			:  Optional[int]		    =  Adafactor if self.args.adafactor else AdamW
     if self.args.adafactor:
      __a			:  int		    =  Adafactor
      __a			:  Union[str, Any]		    =  {'''scale_parameter''': False, '''relative_step''': False}
     else:
      __a			:  Tuple		    =  AdamW
      __a			:  List[Any]		    =  {
          '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
          '''eps''': self.args.adam_epsilon,
      }
     __a			:  str		    =  self.args.learning_rate
     if self.sharded_ddp:
      __a			:  str		    =  OSS(
          params=lowerCamelCase__ ,    optim=lowerCamelCase__ ,    **lowerCamelCase__ ,    )
     else:
      __a			:  int		    =  optimizer_cls(lowerCamelCase__ ,    **lowerCamelCase__       )

    if self.lr_scheduler is None:
     __a			:  Optional[int]		    =  self._get_lr_scheduler(lowerCamelCase__       )
    else:  # ignoring --lr_scheduler
     logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''       )
   def 				lowerCAmelCase      (self			: List[Any] ,    snake_case_			: Union[str, Any]       ):
    __a			:  Union[str, Any]		    =  arg_to_scheduler[self.args.lr_scheduler]
    if self.args.lr_scheduler == "constant":
     __a			:  str		    =  schedule_func(self.optimizer       )
    elif self.args.lr_scheduler == "constant_w_warmup":
     __a			:  str		    =  schedule_func(self.optimizer ,    num_warmup_steps=self.args.warmup_steps       )
    else:
     __a			:  str		    =  schedule_func(
         self.optimizer ,    num_warmup_steps=self.args.warmup_steps ,    num_training_steps=lowerCamelCase__       )
    return scheduler
   def 				lowerCAmelCase      (self			: Any       ):






    if isinstance(self.train_dataset ,    torch.utils.data.IterableDataset       ):
     return None
    elif is_torch_tpu_available():
     return get_tpu_sampler(self.train_dataset       )
    else:
     if self.args.sortish_sampler:
      self.train_dataset.make_sortish_sampler(
          self.args.per_device_train_batch_size ,    distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED) ,    )

     return (
         RandomSampler(self.train_dataset       )
         if self.args.local_rank == -1
         else DistributedSampler(self.train_dataset       )
     )
   def 				lowerCAmelCase      (self			: Optional[Any] ,    snake_case_			: str ,    snake_case_			: List[str] ,    snake_case_			: int       ):






    if self.args.label_smoothing == 0:
     if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
      # force training to ignore pad token
      __a			:  Any		    =  model(**lowerCamelCase__ ,    use_cache=lowerCamelCase__       )[0]
      __a			:  int		    =  self.loss_fn(logits.view(-1 ,    logits.shape[-1]       ) ,    labels.view(-1       )       )
     else:
      # compute usual loss via models
      __a			:  List[Any]		    =  model(**lowerCamelCase__ ,    labels=lowerCamelCase__ ,    use_cache=lowerCamelCase__       )[:2]
    else:
     # compute label smoothed loss
     __a			:  Optional[Any]		    =  model(**lowerCamelCase__ ,    use_cache=lowerCamelCase__       )[0]
     __a			:  List[Any]		    =  torch.nn.functional.log_softmax(lowerCamelCase__ ,    dim=-1       )
     __a			:  Dict		    =  self.loss_fn(lowerCamelCase__ ,    lowerCamelCase__ ,    self.args.label_smoothing ,    ignore_index=self.config.pad_token_id       )
    return loss, logits
   def 				lowerCAmelCase      (self			: Optional[int] ,    snake_case_			: Optional[Any] ,    snake_case_			: List[str]       ):
    __a			:  Tuple		    =  inputs.pop('''labels'''       )
    __a			:  Tuple		    =  self._compute_loss(lowerCamelCase__ ,    lowerCamelCase__ ,    lowerCamelCase__       )
    return loss
   def 				lowerCAmelCase      (self			: List[Any] ,    snake_case_			: nn.Module ,    snake_case_			: Dict[str, Union[torch.Tensor, Any]] ,    snake_case_			: bool ,    snake_case_			: Optional[List[str]] = None ,    ):
    __a			:  Dict		    =  self._prepare_inputs(lowerCamelCase__       )

    __a			:  Dict		    =  {
        '''max_length''': self.data_args.val_max_target_length
        if self.data_args is not None
        else self.config.max_length,
        '''num_beams''': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
    }

    if self.args.predict_with_generate and not self.args.prediction_loss_only:
     __a			:  Tuple		    =  self.model.generate(
         inputs['''input_ids'''] ,    attention_mask=inputs['''attention_mask'''] ,    **lowerCamelCase__ ,    )
     # in case the batch is shorter than max length, the output should be padded
     if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
      __a			:  Any		    =  self._pad_tensors_to_max_len(lowerCamelCase__ ,    gen_kwargs['''max_length''']       )

    __a			:  Optional[Any]		    =  inputs.pop('''labels'''       )
    with torch.no_grad():
     # compute loss on predict data
     __a			:  str		    =  self._compute_loss(lowerCamelCase__ ,    lowerCamelCase__ ,    lowerCamelCase__       )

    __a			:  Union[str, Any]		    =  loss.mean().detach()
    if self.args.prediction_loss_only:
     return (loss, None, None)

    __a			:  Tuple		    =  generated_tokens if self.args.predict_with_generate else logits

    if labels.shape[-1] < gen_kwargs["max_length"]:
     __a			:  int		    =  self._pad_tensors_to_max_len(lowerCamelCase__ ,    gen_kwargs['''max_length''']       )

    return (loss, logits, labels)







   def 				lowerCAmelCase      (self			: int ,    snake_case_			: Tuple ,    snake_case_			: str       ):
    __a			:  Optional[int]		    =  self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

    if pad_token_id is None:
     raise ValueError(
         '''Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'''
         f" padded to `max_length`={max_length}"       )

    __a			:  List[str]		    =  pad_token_id * torch.ones(
        (tensor.shape[0], max_length) ,    dtype=tensor.dtype ,    device=tensor.device       )
    __a			:  Optional[Any]		    =  tensor
    return padded_tensor


import os
import unittest

from transformers.models.phobert.tokenization_phobert import VOCAB_FILES_NAMES, PhobertTokenizer

from ...test_tokenization_common import TokenizerTesterMixin



class 		UpperCamelCase__  (	__lowercase   ,unittest.TestCase   ):
   _SCREAMING_SNAKE_CASE				:							Union[str, Any]     	=				PhobertTokenizer
   _SCREAMING_SNAKE_CASE				:							int     	=				False
   def 				lowerCAmelCase      (self			: Optional[int]       ):
    super().setUp()

    # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
    __a			:  Optional[int]		    =  ['''T@@''', '''i''', '''I''', '''R@@''', '''r''', '''e@@''']
    __a			:  Tuple		    =  dict(zip(snake_case_ ,    range(len(snake_case_       )       )       )       )
    __a			:  int		    =  ['''#version: 0.2''', '''l à</w>''']
    __a			:  List[Any]		    =  {'''unk_token''': '''<unk>'''}

    __a			:  Any		    =  os.path.join(self.tmpdirname ,    VOCAB_FILES_NAMES['''vocab_file''']       )
    __a			:  List[Any]		    =  os.path.join(self.tmpdirname ,    VOCAB_FILES_NAMES['''merges_file''']       )

    with open(self.vocab_file ,    '''w''' ,    encoding='''utf-8'''       ) as fp:
     for token in vocab_tokens:
      fp.write(f"{token} {vocab_tokens[token]}\n"       )
    with open(self.merges_file ,    '''w''' ,    encoding='''utf-8'''       ) as fp:
     fp.write('''\n'''.join(snake_case_       )       )
   def 				lowerCAmelCase      (self			: Union[str, Any] ,    **snake_case_			: List[str]       ):
    kwargs.update(self.special_tokens_map       )
    return PhobertTokenizer.from_pretrained(self.tmpdirname ,    **snake_case_       )
   def 				lowerCAmelCase      (self			: Any ,    snake_case_			: Dict       ):
    __a			:  Union[str, Any]		    =  '''Tôi là VinAI Research'''
    __a			:  int		    =  '''T<unk> i <unk> <unk> <unk> <unk> <unk> <unk> I Re<unk> e<unk> <unk> <unk> <unk>'''
    return input_text, output_text







   def 				lowerCAmelCase      (self			: Optional[Any]       ):
    __a			:  Optional[int]		    =  PhobertTokenizer(self.vocab_file ,    self.merges_file ,    **self.special_tokens_map       )
    __a			:  Any		    =  '''Tôi là VinAI Research'''
    __a			:  Union[str, Any]		    =  '''T@@ ô@@ i l@@ à V@@ i@@ n@@ A@@ I R@@ e@@ s@@ e@@ a@@ r@@ c@@ h'''.split()
    __a			:  List[str]		    =  tokenizer.tokenize(snake_case_       )
    print(snake_case_       )
    self.assertListEqual(snake_case_ ,    snake_case_       )

    __a			:  str		    =  tokens + [tokenizer.unk_token]

    __a			:  str		    =  [4, 3, 5, 3, 3, 3, 3, 3, 3, 6, 7, 9, 3, 9, 3, 3, 3, 3, 3]
    self.assertListEqual(tokenizer.convert_tokens_to_ids(snake_case_       ) ,    snake_case_       )
