





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     SCREAMING_SNAKE_CASE       (__snake_case				):


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





     """simple docstring"""







     super().__init__(*UpperCamelCase__      ,					**UpperCamelCase__	)

     if config is None:
       assert isinstance(self.model      ,					UpperCamelCase__	), (
           "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
           f""" {self.model.__class__}"""
       )
       lowercase__           =  self.model.config
     else:
       lowercase__           =  config

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

       lowercase__           =  label_smoothed_nll_loss


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict      ,					a      :    int	)->       Dict:





     """simple docstring"""







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

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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]      ,					a      :    Any	)->       Union[str, Any]:





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   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 	SCREAMING_SNAKE_CASE_							(   self      :    str      ,					a      :    Optional[Any]      ,					a      :    str      ,					a      :    List[str]	)->       Dict:





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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]      ,					a      :    Any      ,					a      :    Optional[Any]	)->       Any:





     """simple docstring"""







     lowercase__           =  inputs.pop('labels'	)
     lowercase__           =  self._compute_loss(UpperCamelCase__      ,					UpperCamelCase__      ,					UpperCamelCase__	)
     return loss


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





     """simple docstring"""







     lowercase__           =  self._prepare_inputs(UpperCamelCase__	)

     lowercase__           =  {
         "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:
       lowercase__           =  self.model.generate(
           inputs['input_ids']      ,					attention_mask=inputs['attention_mask']      ,					**UpperCamelCase__      ,					)
       # in case the batch is shorter than max length, the output should be padded
       if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
         lowercase__           =  self._pad_tensors_to_max_len(UpperCamelCase__      ,					gen_kwargs['max_length']	)

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

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

     lowercase__           =  generated_tokens if self.args.predict_with_generate else logits

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

     return (loss, logits, labels)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[Any]      ,					a      :    Dict      ,					a      :    Optional[int]	)->       List[Any]:





     """simple docstring"""







     lowercase__           =  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}"""	)

     lowercase__           =  pad_token_id * torch.ones(
         (tensor.shape[0], max_length)      ,					dtype=tensor.dtype      ,					device=tensor.device	)
     lowercase__           =  tensor
     return padded_tensor









import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      import jax.numpy as jnp

      from transformers.models.distilbert.modeling_flax_distilbert import (
          FlaxDistilBertForMaskedLM,
          FlaxDistilBertForMultipleChoice,
          FlaxDistilBertForQuestionAnswering,
          FlaxDistilBertForSequenceClassification,
          FlaxDistilBertForTokenClassification,
          FlaxDistilBertModel,
      )
class     SCREAMING_SNAKE_CASE       (unittest.TestCase				):


   def __init__(   self      :    List[str]      ,					a      :    Dict      ,					a      :    Optional[int]=13      ,					a      :    int=7      ,					a      :    List[str]=True      ,					a      :    Any=True      ,					a      :    Dict=True      ,					a      :    List[Any]=True      ,					a      :    List[str]=99      ,					a      :    Dict=32      ,					a      :    List[str]=5      ,					a      :    Tuple=4      ,					a      :    Optional[int]=37      ,					a      :    Union[str, Any]="gelu"      ,					a      :    Optional[Any]=0.1      ,					a      :    Optional[int]=0.1      ,					a      :    Optional[Any]=512      ,					a      :    Dict=16      ,					a      :    Any=2      ,					a      :    Tuple=0.02      ,					a      :    Optional[Any]=4      ,					)->       int:





     """simple docstring"""







     lowercase__           =  parent
     lowercase__           =  batch_size
     lowercase__           =  seq_length
     lowercase__           =  is_training
     lowercase__           =  use_attention_mask
     lowercase__           =  use_token_type_ids
     lowercase__           =  use_labels
     lowercase__           =  vocab_size
     lowercase__           =  hidden_size
     lowercase__           =  num_hidden_layers
     lowercase__           =  num_attention_heads
     lowercase__           =  intermediate_size
     lowercase__           =  hidden_act
     lowercase__           =  hidden_dropout_prob
     lowercase__           =  attention_probs_dropout_prob
     lowercase__           =  max_position_embeddings
     lowercase__           =  type_vocab_size
     lowercase__           =  type_sequence_label_size
     lowercase__           =  initializer_range
     lowercase__           =  num_choices


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Union[str, Any]	)->       Union[str, Any]:





     """simple docstring"""







     lowercase__           =  ids_tensor([self.batch_size, self.seq_length]      ,					self.vocab_size	)

     lowercase__           =  None
     if self.use_attention_mask:
       lowercase__           =  random_attention_mask([self.batch_size, self.seq_length]	)

     lowercase__           =  DistilBertConfig(
         vocab_size=self.vocab_size      ,					dim=self.hidden_size      ,					n_layers=self.num_hidden_layers      ,					n_heads=self.num_attention_heads      ,					hidden_dim=self.intermediate_size      ,					hidden_act=self.hidden_act      ,					dropout=self.hidden_dropout_prob      ,					attention_dropout=self.attention_probs_dropout_prob      ,					max_position_embeddings=self.max_position_embeddings      ,					initializer_range=self.initializer_range      ,					tie_weights_=a      ,					)

     return config, input_ids, attention_mask


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]	)->       Any:





     """simple docstring"""







     lowercase__           =  self.prepare_config_and_inputs()
     lowercase__      ,   lowercase__      ,   lowercase__           =  config_and_inputs
     lowercase__           =  {'input_ids': input_ids, 'attention_mask': attention_mask}
     return config, inputs_dict






@require_flax
class     SCREAMING_SNAKE_CASE       (UpperCAmelCase      ,     unittest.TestCase				):
   _UpperCamelCase   :		List[Any]			 =	(
       (
           FlaxDistilBertModel,
           FlaxDistilBertForMaskedLM,
           FlaxDistilBertForMultipleChoice,
           FlaxDistilBertForQuestionAnswering,
           FlaxDistilBertForSequenceClassification,
           FlaxDistilBertForTokenClassification,
           FlaxDistilBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict	)->       List[str]:





     """simple docstring"""







     lowercase__           =  FlaxDistilBertModelTester(self	)


   @slow
   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]	)->       int:





     """simple docstring"""







     for model_class_name in self.all_model_classes:
       lowercase__           =  model_class_name.from_pretrained('distilbert-base-uncased'	)
       lowercase__           =  model(np.ones((1, 1)	)	)
       self.assertIsNotNone(a	)






@require_flax
class     SCREAMING_SNAKE_CASE       (unittest.TestCase				):


   @slow
   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]	)->       List[str]:





     """simple docstring"""







     lowercase__           =  FlaxDistilBertModel.from_pretrained('distilbert-base-uncased'	)
     lowercase__           =  np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]	)
     lowercase__           =  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
     lowercase__           =  model(a      ,					attention_mask=a	)[0]
     lowercase__           =  (1, 11, 768)
     self.assertEqual(output.shape      ,					a	)
     lowercase__           =  np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]	)

     self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]      ,					a      ,					atol=1E-4	)	)


