





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       :				List[str]           =			logging.get_logger(__name__)







class 		_lowerCAmelCase(		_a):

  """simple docstring"""







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

  @staticmethod
  def 						SCREAMING_SNAKE_CASE__			(   UpperCAmelCase      )->   GenerationConfig:

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

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

     __A    =  GenerationConfig.from_pretrained(UpperCAmelCase    ,							UpperCAmelCase      )
     return gen_config

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

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

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

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

  def 						SCREAMING_SNAKE_CASE__			(   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      )

     __A    =  '''labels''' in inputs
     __A    =  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'''      ):
        __A    =  self._gen_kwargs.copy()

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

     # Retrieves GenerationConfig from model.generation_config
     __A    =  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:
        __A    =  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:
        __A    =  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():
              __A    =  model(**UpperCAmelCase      )
           if self.label_smoother is not None:
              __A    =  self.label_smoother(UpperCAmelCase    ,							inputs['''labels''']      ).mean().detach()
           else:
              __A    =  (outputs['''loss'''] if isinstance(UpperCAmelCase    ,							UpperCAmelCase      ) else outputs[0]).mean().detach()
        else:
           __A    =  None

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

     if has_labels:
        __A    =  inputs['''labels''']
        if labels.shape[-1] < gen_config.max_length:
           __A    =  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:
           __A    =  self._pad_tensors_to_max_len(UpperCAmelCase    ,							gen_config.max_new_tokens + 1      )
     else:
        __A    =  None

     return loss, generated_tokens, labels





  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   Tuple:
     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
        __A    =  (
            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:
           __A    =  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'''      )

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








class 		_lowerCAmelCase:

  """simple docstring"""







  def __init__(   self      )->   None:
     # numbermap keys are the number and its values are lists of indexes sorted
     # in ascending order
     __A    =  {}
     # indexmap keys are an index and it's values are the number at that index
     __A    =  {}

  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   list[int]:
     if isinstance(UpperCAmelCase    ,							(range, str)      ):
        __A    =  list(UpperCAmelCase      )
     elif not isinstance(UpperCAmelCase    ,							UpperCAmelCase      ):
        raise TypeError(
            '''binary_search_delete() only accepts either a list, range or str'''      )

     __A    =  0
     __A    =  len(UpperCAmelCase      ) - 1

     while low <= high:
        __A    =  (low + high) // 2
        if array[mid] == item:
           array.pop(UpperCAmelCase      )
           return array
        elif array[mid] < item:
           __A    =  mid + 1
        else:
           __A    =  mid - 1
     raise ValueError(
         '''Either the item is not in the array or the array was unsorted'''      )

  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   list[int]:
     if isinstance(UpperCAmelCase    ,							(range, str)      ):
        __A    =  list(UpperCAmelCase      )
     elif not isinstance(UpperCAmelCase    ,							UpperCAmelCase      ):
        raise TypeError(
            '''binary_search_insert() only accepts either a list, range or str'''      )

     __A    =  0
     __A    =  len(UpperCAmelCase      ) - 1

     while low <= high:
        __A    =  (low + high) // 2
        if array[mid] == index:
           # If the item already exists in the array,
           # insert it after the existing item
           array.insert(mid + 1    ,							UpperCAmelCase      )
           return array
        elif array[mid] < index:
           __A    =  mid + 1
        else:
           __A    =  mid - 1

        # If the item doesn't exist in the array, insert it at the appropriate position
     array.insert(UpperCAmelCase    ,							UpperCAmelCase      )
     return array

  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase    ,							UpperCAmelCase      )->   None:
     # Remove previous index
     if index in self.indexmap:
        __A    =  self.indexmap[index]
        if len(self.numbermap[n]      ) == 1:
           del self.numbermap[n]
        else:
           __A    =  self.binary_search_delete(self.numbermap[n]    ,							UpperCAmelCase      )

        # Set new index
     __A    =  number

     # Number not seen before or empty so insert number value
     if number not in self.numbermap:
        __A    =  [index]

     # Here we need to perform a binary search insertion in order to insert
     # The item in the correct place
     else:
        __A    =  self.binary_search_insert(
            self.numbermap[number]    ,							UpperCAmelCase      )





  def 						SCREAMING_SNAKE_CASE__			(   self    ,							UpperCAmelCase      )->   int:
     return self.numbermap.get(UpperCAmelCase    ,							[-1]      )[0]


if __name__ == "__main__":
    import doctest

    doctest.testmod()

