


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


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

__UpperCamelCase       :       List[str]						       =				{
    "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       __magic_name__	(     __lowerCAmelCase):


     def __init__(      self							:     int							,	lowerCamelCase__							:     Any=None							,	lowerCamelCase__							:     str=None							,	*lowerCamelCase__							:     Union[str, Any]							,	**lowerCamelCase__							:     List[str]			)  ->      Optional[int]:





      '''simple docstring'''






      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__}"
       )
       UpperCamelCase__					:    int          =   self.model.config
      else:
       UpperCamelCase__					:    Optional[Any]          =   config

      UpperCamelCase__					:    str          =   data_args
      UpperCamelCase__					:    Any          =   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:
       UpperCamelCase__					:    str          =   torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id			)
      else:
       # dynamically import label_smoothed_nll_loss
       from utils import label_smoothed_nll_loss

       UpperCamelCase__					:    Optional[Any]          =   label_smoothed_nll_loss


     def      UpperCAmelCase__	(      self							:     Dict							,	lowerCamelCase__							:     int			)  ->      Tuple:





      '''simple docstring'''






      if self.optimizer is None:
       UpperCamelCase__					:    Optional[int]          =   ['''bias''', '''LayerNorm.weight''']
       UpperCamelCase__					:    List[str]          =   [
           {
               '''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,
           },
       ]
       UpperCamelCase__					:    str          =   Adafactor if self.args.adafactor else AdamW
       if self.args.adafactor:
        UpperCamelCase__					:    Dict          =   Adafactor
        UpperCamelCase__					:    Dict          =   {'''scale_parameter''': False, '''relative_step''': False}
       else:
        UpperCamelCase__					:    Optional[Any]          =   AdamW
        UpperCamelCase__					:    int          =   {
            '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
            '''eps''': self.args.adam_epsilon,
        }
       UpperCamelCase__					:    List[Any]          =   self.args.learning_rate
       if self.sharded_ddp:
        UpperCamelCase__					:    List[str]          =   OSS(
            params=lowerCamelCase__							,	optim=lowerCamelCase__							,	**lowerCamelCase__							,	)
       else:
        UpperCamelCase__					:    Optional[int]          =   optimizer_cls(lowerCamelCase__							,	**lowerCamelCase__			)

      if self.lr_scheduler is None:
       UpperCamelCase__					:    Optional[Any]          =   self._get_lr_scheduler(lowerCamelCase__			)
      else:  # ignoring --lr_scheduler
       logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''			)


     def      UpperCAmelCase__	(      self							:     List[str]							,	lowerCamelCase__							:     int			)  ->      Tuple:





      '''simple docstring'''






      UpperCamelCase__					:    Any          =   arg_to_scheduler[self.args.lr_scheduler]
      if self.args.lr_scheduler == "constant":
       UpperCamelCase__					:    int          =   schedule_func(self.optimizer			)
      elif self.args.lr_scheduler == "constant_w_warmup":
       UpperCamelCase__					:    Optional[int]          =   schedule_func(self.optimizer							,	num_warmup_steps=self.args.warmup_steps			)
      else:
       UpperCamelCase__					:    int          =   schedule_func(
           self.optimizer							,	num_warmup_steps=self.args.warmup_steps							,	num_training_steps=lowerCamelCase__			)
      return scheduler


     def      UpperCAmelCase__	(      self							:     List[str]			)  ->      Optional[torch.utils.data.Sampler]:





      '''simple docstring'''






      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      UpperCAmelCase__	(      self							:     List[str]							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     List[Any]			)  ->      Optional[int]:





      '''simple docstring'''






      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
        UpperCamelCase__					:    Optional[Any]          =   model(**lowerCamelCase__							,	use_cache=lowerCamelCase__			)[0]
        UpperCamelCase__					:    Tuple          =   self.loss_fn(logits.view(-1							,	logits.shape[-1]			)							,	labels.view(-1			)			)
       else:
        # compute usual loss via models
        UpperCamelCase__       ,				UpperCamelCase__					:    List[str]          =   model(**lowerCamelCase__							,	labels=lowerCamelCase__							,	use_cache=lowerCamelCase__			)[:2]
      else:
       # compute label smoothed loss
       UpperCamelCase__					:    Union[str, Any]          =   model(**lowerCamelCase__							,	use_cache=lowerCamelCase__			)[0]
       UpperCamelCase__					:    List[str]          =   torch.nn.functional.log_softmax(lowerCamelCase__							,	dim=-1			)
       UpperCamelCase__       ,				UpperCamelCase__					:    Dict          =   self.loss_fn(lowerCamelCase__							,	lowerCamelCase__							,	self.args.label_smoothing							,	ignore_index=self.config.pad_token_id			)
      return loss, logits


     def      UpperCAmelCase__	(      self							:     List[Any]							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Union[str, Any]			)  ->      Any:





      '''simple docstring'''






      UpperCamelCase__					:    int          =   inputs.pop('''labels'''			)
      UpperCamelCase__       ,				UpperCamelCase__					:    Dict          =   self._compute_loss(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)
      return loss


     def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     nn.Module							,	lowerCamelCase__							:     Dict[str, Union[torch.Tensor, Any]]							,	lowerCamelCase__							:     bool							,	lowerCamelCase__							:     Optional[List[str]] = None							,	)  ->      Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:





      '''simple docstring'''






      UpperCamelCase__					:    str          =   self._prepare_inputs(lowerCamelCase__			)

      UpperCamelCase__					:    str          =   {
          '''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:
       UpperCamelCase__					:    List[str]          =   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"]:
        UpperCamelCase__					:    Tuple          =   self._pad_tensors_to_max_len(lowerCamelCase__							,	gen_kwargs['''max_length''']			)

      UpperCamelCase__					:    Tuple          =   inputs.pop('''labels'''			)
      with torch.no_grad():
       # compute loss on predict data
       UpperCamelCase__       ,				UpperCamelCase__					:    Tuple          =   self._compute_loss(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)

      UpperCamelCase__					:    Dict          =   loss.mean().detach()
      if self.args.prediction_loss_only:
       return (loss, None, None)

      UpperCamelCase__					:    int          =   generated_tokens if self.args.predict_with_generate else logits

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

      return (loss, logits, labels)


     def      UpperCAmelCase__	(      self							:     str							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     Tuple			)  ->      Optional[int]:





      '''simple docstring'''






      UpperCamelCase__					:    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}"			)

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






import unicodedata
from dataclasses import dataclass
from typing import Optional, Union

import numpy as np

from transformers.data.data_collator import DataCollatorMixin
from transformers.file_utils import PaddingStrategy
from transformers.tokenization_utils_base import PreTrainedTokenizerBase

def 			_a  (				SCREAMING_SNAKE_CASE	:				Optional[Any]    , SCREAMING_SNAKE_CASE	:				Any    , SCREAMING_SNAKE_CASE	:				List[str]    , SCREAMING_SNAKE_CASE	:				int     ):







 """simple docstring"""



 if isinstance(SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE     ):
  UpperCamelCase__					:    Dict          =   np.full((len(SCREAMING_SNAKE_CASE     ), sequence_length, 2)    , SCREAMING_SNAKE_CASE     )
 else:
  UpperCamelCase__					:    Dict          =   np.full((len(SCREAMING_SNAKE_CASE     ), sequence_length)    , SCREAMING_SNAKE_CASE     )

 for i, tensor in enumerate(SCREAMING_SNAKE_CASE     ):
  if padding_side == "right":
   if isinstance(SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE     ):
    UpperCamelCase__					:    str          =   tensor[:sequence_length]
   else:
    UpperCamelCase__					:    Union[str, Any]          =   tensor[:sequence_length]
  else:
   if isinstance(SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE     ):
    UpperCamelCase__					:    str          =   tensor[:sequence_length]
   else:
    UpperCamelCase__					:    str          =   tensor[:sequence_length]

 return out_tensor.tolist()

def 			_a  (				SCREAMING_SNAKE_CASE	:				List[Any]     ):







 """simple docstring"""



 UpperCamelCase__					:    Union[str, Any]          =   ord(SCREAMING_SNAKE_CASE     )
 if (cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126):
  return True
 UpperCamelCase__					:    Any          =   unicodedata.category(SCREAMING_SNAKE_CASE     )
 if cat.startswith('''P'''     ):
  return True
 return False






@dataclass
class       __magic_name__	(     __lowerCAmelCase):
     A:		PreTrainedTokenizerBase
     A:		Union[bool, str, PaddingStrategy]			=      True
     A:		Optional[int]			=      None
     A:		Optional[int]			=      None
     A:		int			=      -1_0_0
     A:		str			=      "pt"


     def      UpperCAmelCase__	(      self							:     int							,	lowerCamelCase__							:     List[Any]			)  ->      Any:





      '''simple docstring'''






      import torch

      UpperCamelCase__					:    Tuple          =   '''label''' if '''label''' in features[0].keys() else '''labels'''
      UpperCamelCase__					:    Optional[Any]          =   [feature[label_name] for feature in features] if label_name in features[0].keys() else None
      UpperCamelCase__					:    Optional[int]          =   self.tokenizer.pad(
          lowerCamelCase__							,	padding=self.padding							,	max_length=self.max_length							,	pad_to_multiple_of=self.pad_to_multiple_of							,	return_tensors='''pt''' if labels is None else None							,	)

      if labels is None:
       return batch

      UpperCamelCase__					:    Dict          =   torch.tensor(batch['''entity_ids''']			).shape[1]
      UpperCamelCase__					:    Union[str, Any]          =   self.tokenizer.padding_side
      if padding_side == "right":
       UpperCamelCase__					:    Optional[int]          =   [
           list(lowerCamelCase__			) + [self.label_pad_token_id] * (sequence_length - len(lowerCamelCase__			)) for label in labels
       ]
      else:
       UpperCamelCase__					:    int          =   [
           [self.label_pad_token_id] * (sequence_length - len(lowerCamelCase__			)) + list(lowerCamelCase__			) for label in labels
       ]

      UpperCamelCase__					:    Union[str, Any]          =   [feature['''ner_tags'''] for feature in features]
      UpperCamelCase__					:    List[str]          =   padding_tensor(lowerCamelCase__							,	-1							,	lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    Union[str, Any]          =   [feature['''original_entity_spans'''] for feature in features]
      UpperCamelCase__					:    Any          =   padding_tensor(lowerCamelCase__							,	(-1, -1)							,	lowerCamelCase__							,	lowerCamelCase__			)
      UpperCamelCase__					:    str          =   {k: torch.tensor(lowerCamelCase__							,	dtype=torch.intaa			) for k, v in batch.items()}

      return batch


