





from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import Dataset

from .deepspeed import is_deepspeed_zeroa_enabled
from .generation.configuration_utils import GenerationConfig
from .trainer import Trainer
from .utils import logging


if TYPE_CHECKING:
  from .data.data_collator import DataCollator
  from .modeling_utils import PreTrainedModel
  from .tokenization_utils_base import PreTrainedTokenizerBase
  from .trainer_callback import TrainerCallback
  from .trainer_utils import EvalPrediction, PredictionOutput
  from .training_args import TrainingArguments


_snake_case				= logging.get_logger(__name__)





class 					lowerCAmelCase							(   lowercase_					):







  def __init__(			self  :Tuple    ,							_lowercase  :Union["PreTrainedModel", nn.Module] = None    ,							_lowercase  :"TrainingArguments" = None    ,							_lowercase  :Optional["DataCollator"] = None    ,							_lowercase  :Optional[Dataset] = None    ,							_lowercase  :Optional[Union[Dataset, Dict[str, Dataset]]] = None    ,							_lowercase  :Optional["PreTrainedTokenizerBase"] = None    ,							_lowercase  :Optional[Callable[[], "PreTrainedModel"]] = None    ,							_lowercase  :Optional[Callable[["EvalPrediction"], Dict]] = None    ,							_lowercase  :Optional[List["TrainerCallback"]] = None    ,							_lowercase  :Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)    ,							_lowercase  :Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None    ,							):






     '''simple docstring'''





     super().__init__(
         model=_lowercase    ,							args=_lowercase    ,							data_collator=_lowercase    ,							train_dataset=_lowercase    ,							eval_dataset=_lowercase    ,							tokenizer=_lowercase    ,							model_init=_lowercase    ,							compute_metrics=_lowercase    ,							callbacks=_lowercase    ,							optimizers=_lowercase    ,							preprocess_logits_for_metrics=_lowercase    ,							)

     # Override self.model.generation_config if a GenerationConfig is specified in args.
     # Priority: args.generation_config > model.generation_config > default GenerationConfig.
     if self.args.generation_config is not None:
        lowercase__    =  self.load_generation_config(self.args.generation_config					)
        lowercase__    =  gen_config







  @staticmethod
  def   UpperCAmelCase     (			_lowercase  :Union[str, GenerationConfig]					):






     '''simple docstring'''





     if isinstance(_lowercase    ,							_lowercase					):
        return deepcopy(_lowercase					)

     # str or Path
     lowercase__    =  Path(_lowercase					) if isinstance(_lowercase    ,							_lowercase					) else gen_config_arg
     lowercase__    =  None

     # Figuring if it is path pointing to a file, pointing to a directory or else a model id or URL
     # This step is required in order to determine config_file_name
     if pretrained_model_name.is_file():
        lowercase__    =  pretrained_model_name.name
        lowercase__    =  pretrained_model_name.parent
     # dir path
     elif pretrained_model_name.is_dir():
        pass
     # model id or URL
     else:
        lowercase__    =  gen_config_arg

     lowercase__    =  GenerationConfig.from_pretrained(_lowercase    ,							_lowercase					)
     return gen_config







  def   UpperCAmelCase     (			self  :int    ,							_lowercase  :Optional[Dataset] = None    ,							_lowercase  :Optional[List[str]] = None    ,							_lowercase  :str = "eval"    ,							**_lowercase  :Any    ,							):






     '''simple docstring'''





     lowercase__    =  gen_kwargs.copy()
     if gen_kwargs.get("max_length"					) is None and gen_kwargs.get("max_new_tokens"					) is None:
        lowercase__    =  self.args.generation_max_length
     lowercase__    =  (
         gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"					) is not None else self.args.generation_num_beams
     )
     lowercase__    =  gen_kwargs

     return super().evaluate(_lowercase    ,							ignore_keys=_lowercase    ,							metric_key_prefix=_lowercase					)







  def   UpperCAmelCase     (			self  :int    ,							_lowercase  :Dataset    ,							_lowercase  :Optional[List[str]] = None    ,							_lowercase  :str = "test"    ,							**_lowercase  :int    ,							):






     '''simple docstring'''





     lowercase__    =  gen_kwargs.copy()
     if gen_kwargs.get("max_length"					) is None and gen_kwargs.get("max_new_tokens"					) is None:
        lowercase__    =  self.args.generation_max_length
     lowercase__    =  (
         gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"					) is not None else self.args.generation_num_beams
     )
     lowercase__    =  gen_kwargs

     return super().predict(_lowercase    ,							ignore_keys=_lowercase    ,							metric_key_prefix=_lowercase					)







  def   UpperCAmelCase     (			self  :Any    ,							_lowercase  :nn.Module    ,							_lowercase  :Dict[str, Union[torch.Tensor, Any]]    ,							_lowercase  :bool    ,							_lowercase  :Optional[List[str]] = None    ,							**_lowercase  :Tuple    ,							):






     '''simple docstring'''





     if not self.args.predict_with_generate or prediction_loss_only:
        return super().prediction_step(
            _lowercase    ,							_lowercase    ,							prediction_loss_only=_lowercase    ,							ignore_keys=_lowercase					)

     lowercase__    =  "labels" in inputs
     lowercase__    =  self._prepare_inputs(_lowercase					)

     # XXX: adapt synced_gpus for fairscale as well
     # Priority (handled in generate):
     # gen_kwargs > model.generation_config > default GenerationConfig()

     if len(_lowercase					) == 0 and hasattr(self    ,							"_gen_kwargs"					):
        lowercase__    =  self._gen_kwargs.copy()

     if gen_kwargs.get("max_length"					) is None and gen_kwargs.get("max_new_tokens"					) is None:
        lowercase__    =  self.model.config.max_length
     lowercase__    =  (
         gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"					) is not None else self.model.config.num_beams
     )
     lowercase__    =  True if is_deepspeed_zeroa_enabled() else False
     lowercase__    =  (
         gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus"					) is not None else default_synced_gpus
     )

     # If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate
     # (otherwise, it would continue generating from the padded `decoder_input_ids`)
     if (
         "labels" in inputs
         and "decoder_input_ids" in inputs
         and inputs["labels"].shape == inputs["decoder_input_ids"].shape
     ):
        lowercase__    =  {k: v for k, v in inputs.items() if k != "decoder_input_ids"}
     lowercase__    =  self.model.generate(**_lowercase    ,							**_lowercase					)

     # Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
     # TODO: remove this hack when the legacy code that initializes generation_config from a model config is
     # removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183
     if self.model.generation_config._from_model_config:
        lowercase__    =  False

     # Retrieves GenerationConfig from model.generation_config
     lowercase__    =  self.model.generation_config
     # in case the batch is shorter than max length, the output should be padded
     if generated_tokens.shape[-1] < gen_config.max_length:
        lowercase__    =  self._pad_tensors_to_max_len(_lowercase    ,							gen_config.max_length					)
     elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
        lowercase__    =  self._pad_tensors_to_max_len(_lowercase    ,							gen_config.max_new_tokens + 1					)

     with torch.no_grad():
        if has_labels:
           with self.compute_loss_context_manager():
              lowercase__    =  model(**_lowercase					)
           if self.label_smoother is not None:
              lowercase__    =  self.label_smoother(_lowercase    ,							inputs["labels"]					).mean().detach()
           else:
              lowercase__    =  (outputs["loss"] if isinstance(_lowercase    ,							_lowercase					) else outputs[0]).mean().detach()
        else:
           lowercase__    =  None

     if self.args.prediction_loss_only:
        return loss, None, None

     if has_labels:
        lowercase__    =  inputs["labels"]
        if labels.shape[-1] < gen_config.max_length:
           lowercase__    =  self._pad_tensors_to_max_len(_lowercase    ,							gen_config.max_length					)
        elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
           lowercase__    =  self._pad_tensors_to_max_len(_lowercase    ,							gen_config.max_new_tokens + 1					)
     else:
        lowercase__    =  None

     return loss, generated_tokens, labels







  def   UpperCAmelCase     (			self  :Dict    ,							_lowercase  :str    ,							_lowercase  :List[str]					):






     '''simple docstring'''





     if self.tokenizer is not None and hasattr(self.tokenizer    ,							"pad_token_id"					):
        # If PAD token is not defined at least EOS token has to be defined
        lowercase__    =  (
            self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
        )
     else:
        if self.model.config.pad_token_id is not None:
           lowercase__    =  self.model.config.pad_token_id
        else:
           raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors"					)

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













from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_snake_case				= {"""configuration_dinat""": ["""DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DinatConfig"""]}


try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _snake_case				= [
      """DINAT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DinatForImageClassification""",
      """DinatModel""",
      """DinatPreTrainedModel""",
      """DinatBackbone""",
  ]

if TYPE_CHECKING:
  from .configuration_dinat import DINAT_PRETRAINED_CONFIG_ARCHIVE_MAP, DinatConfig

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_dinat import (
        DINAT_PRETRAINED_MODEL_ARCHIVE_LIST,
        DinatBackbone,
        DinatForImageClassification,
        DinatModel,
        DinatPreTrainedModel,
    )

else:
  import sys

  _snake_case				= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






