






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


_UpperCAmelCase				:	Dict    						=    logging.get_logger(__name__)


class 		lowerCAmelCase_							(     UpperCamelCase   ):





     def __init__(  self       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = (None, None)       ,   _UpperCAmelCase = None       ,   )		->							List[Any]:
        super().__init__(
            model=_UpperCAmelCase       ,   args=_UpperCAmelCase       ,   data_collator=_UpperCAmelCase       ,   train_dataset=_UpperCAmelCase       ,   eval_dataset=_UpperCAmelCase       ,   tokenizer=_UpperCAmelCase       ,   model_init=_UpperCAmelCase       ,   compute_metrics=_UpperCAmelCase       ,   callbacks=_UpperCAmelCase       ,   optimizers=_UpperCAmelCase       ,   preprocess_logits_for_metrics=_UpperCAmelCase       ,   )

        # 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:
           lowerCamelCase											=self.load_generation_config(self.args.generation_config     )
           lowerCamelCase											=gen_config





     @staticmethod
     def 	lowerCamelCase       (  _UpperCAmelCase     )		->							GenerationConfig:

        # GenerationConfig provided, nothing to do
        if isinstance(_UpperCAmelCase       ,   _UpperCAmelCase     ):
           return deepcopy(_UpperCAmelCase     )

        # str or Path
        lowerCamelCase											=Path(_UpperCAmelCase     ) if isinstance(_UpperCAmelCase       ,   _UpperCAmelCase     ) else gen_config_arg
        lowerCamelCase											=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():
           lowerCamelCase											=pretrained_model_name.name
           lowerCamelCase											=pretrained_model_name.parent
        # dir path
        elif pretrained_model_name.is_dir():
           pass
        # model id or URL
        else:
           lowerCamelCase											=gen_config_arg

        lowerCamelCase											=GenerationConfig.from_pretrained(_UpperCAmelCase       ,   _UpperCAmelCase     )
        return gen_config





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = "eval"       ,   **_UpperCAmelCase       ,   )		->							Dict[str, float]:
        lowerCamelCase											=gen_kwargs.copy()
        if gen_kwargs.get("""max_length"""     ) is None and gen_kwargs.get("""max_new_tokens"""     ) is None:
           lowerCamelCase											=self.args.generation_max_length
        lowerCamelCase											=(
            gen_kwargs["""num_beams"""] if gen_kwargs.get("""num_beams"""     ) is not None else self.args.generation_num_beams
        )
        lowerCamelCase											=gen_kwargs

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





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase       ,   _UpperCAmelCase = None       ,   _UpperCAmelCase = "test"       ,   **_UpperCAmelCase       ,   )		->							"PredictionOutput":
        lowerCamelCase											=gen_kwargs.copy()
        if gen_kwargs.get("""max_length"""     ) is None and gen_kwargs.get("""max_new_tokens"""     ) is None:
           lowerCamelCase											=self.args.generation_max_length
        lowerCamelCase											=(
            gen_kwargs["""num_beams"""] if gen_kwargs.get("""num_beams"""     ) is not None else self.args.generation_num_beams
        )
        lowerCamelCase											=gen_kwargs

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





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase       ,   _UpperCAmelCase = None       ,   **_UpperCAmelCase       ,   )		->							Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:

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

        lowerCamelCase											="""labels""" in inputs
        lowerCamelCase											=self._prepare_inputs(_UpperCAmelCase     )

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

        if len(_UpperCAmelCase     ) == 0 and hasattr(self       ,   """_gen_kwargs"""     ):
           lowerCamelCase											=self._gen_kwargs.copy()

        if gen_kwargs.get("""max_length"""     ) is None and gen_kwargs.get("""max_new_tokens"""     ) is None:
           lowerCamelCase											=self.model.config.max_length
        lowerCamelCase											=(
            gen_kwargs["""num_beams"""] if gen_kwargs.get("""num_beams"""     ) is not None else self.model.config.num_beams
        )
        lowerCamelCase											=True if is_deepspeed_zeroa_enabled() else False
        lowerCamelCase											=(
            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
        ):
           lowerCamelCase											={k: v for k, v in inputs.items() if k != """decoder_input_ids"""}
        lowerCamelCase											=self.model.generate(**_UpperCAmelCase       ,   **_UpperCAmelCase     )

        # 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:
           lowerCamelCase											=False

        # Retrieves GenerationConfig from model.generation_config
        lowerCamelCase											=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:
           lowerCamelCase											=self._pad_tensors_to_max_len(_UpperCAmelCase       ,   gen_config.max_length     )
        elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
           lowerCamelCase											=self._pad_tensors_to_max_len(_UpperCAmelCase       ,   gen_config.max_new_tokens + 1     )

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

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

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

        return loss, generated_tokens, labels





     def 	lowerCamelCase       (  self       ,   _UpperCAmelCase       ,   _UpperCAmelCase     )		->							int:
        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
           lowerCamelCase											=(
               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:
              lowerCamelCase											=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"""     )

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







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


_UpperCAmelCase				:	List[str]    						=    {
    'configuration_unispeech_sat': ['UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'UniSpeechSatConfig'],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase				:	Optional[Any]    						=    [
       'UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'UniSpeechSatForAudioFrameClassification',
       'UniSpeechSatForCTC',
       'UniSpeechSatForPreTraining',
       'UniSpeechSatForSequenceClassification',
       'UniSpeechSatForXVector',
       'UniSpeechSatModel',
       'UniSpeechSatPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_unispeech_sat import UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechSatConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_unispeech_sat import (
          UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST,
          UniSpeechSatForAudioFrameClassification,
          UniSpeechSatForCTC,
          UniSpeechSatForPreTraining,
          UniSpeechSatForSequenceClassification,
          UniSpeechSatForXVector,
          UniSpeechSatModel,
          UniSpeechSatPreTrainedModel,
      )

else:
   import sys

   _UpperCAmelCase				:	Optional[Any]    						=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)