





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


lowerCAmelCase_        =			logging.get_logger(__name__)



class 				_lowerCAmelCase			(	UpperCAmelCase_							):


  '''simple docstring'''







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


    '''simple docstring'''





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







  @staticmethod
  def UpperCamelCase_       (					UpperCamelCase					:    Union[str, GenerationConfig]							):


    '''simple docstring'''





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

    # str or Path
    _snake_case      : int														=		Path(UpperCamelCase							) if isinstance(UpperCamelCase					, UpperCamelCase							) else gen_config_arg
    _snake_case      : List[Any]														=		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():
      _snake_case      : int														=		pretrained_model_name.name
      _snake_case      : List[str]														=		pretrained_model_name.parent
    # dir path
    elif pretrained_model_name.is_dir():
      pass
    # model id or URL
    else:
      _snake_case      : Any														=		gen_config_arg

    _snake_case      : Optional[int]														=		GenerationConfig.from_pretrained(UpperCamelCase					, UpperCamelCase							)
    return gen_config







  def UpperCamelCase_       (					self					:    int					, UpperCamelCase					:    Optional[Dataset] = None					, UpperCamelCase					:    Optional[List[str]] = None					, UpperCamelCase					:    str = "eval"					, **UpperCamelCase					:    str					, ):


    '''simple docstring'''





    _snake_case      : Optional[int]														=		gen_kwargs.copy()
    if gen_kwargs.get('max_length'							) is None and gen_kwargs.get('max_new_tokens'							) is None:
      _snake_case      : Tuple														=		self.args.generation_max_length
    _snake_case      : str														=		(
        gen_kwargs['num_beams'] if gen_kwargs.get('num_beams'							) is not None else self.args.generation_num_beams
    )
    _snake_case      : int														=		gen_kwargs

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







  def UpperCamelCase_       (					self					:    Dict					, UpperCamelCase					:    Dataset					, UpperCamelCase					:    Optional[List[str]] = None					, UpperCamelCase					:    str = "test"					, **UpperCamelCase					:    Tuple					, ):


    '''simple docstring'''





    _snake_case      : List[Any]														=		gen_kwargs.copy()
    if gen_kwargs.get('max_length'							) is None and gen_kwargs.get('max_new_tokens'							) is None:
      _snake_case      : List[str]														=		self.args.generation_max_length
    _snake_case      : Optional[Any]														=		(
        gen_kwargs['num_beams'] if gen_kwargs.get('num_beams'							) is not None else self.args.generation_num_beams
    )
    _snake_case      : Union[str, Any]														=		gen_kwargs

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







  def UpperCamelCase_       (					self					:    Union[str, Any]					, UpperCamelCase					:    nn.Module					, UpperCamelCase					:    Dict[str, Union[torch.Tensor, Any]]					, UpperCamelCase					:    bool					, UpperCamelCase					:    Optional[List[str]] = None					, **UpperCamelCase					:    Optional[int]					, ):


    '''simple docstring'''





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

    _snake_case      : Dict														=		'labels' in inputs
    _snake_case      : Tuple														=		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'							):
      _snake_case      : Union[str, Any]														=		self._gen_kwargs.copy()

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

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

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

    if has_labels:
      _snake_case      : Tuple														=		inputs['labels']
      if labels.shape[-1] < gen_config.max_length:
        _snake_case      : Dict														=		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:
        _snake_case      : Optional[int]														=		self._pad_tensors_to_max_len(UpperCamelCase					, gen_config.max_new_tokens + 1							)
    else:
      _snake_case      : List[Any]														=		None

    return loss, generated_tokens, labels







  def UpperCamelCase_       (					self					:    Any					, UpperCamelCase					:    str					, UpperCamelCase					:    Tuple							):


    '''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
      _snake_case      : Union[str, Any]														=		(
          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:
        _snake_case      : Dict														=		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'							)

    _snake_case      : List[str]														=		pad_token_id * torch.ones(
        (tensor.shape[0], max_length)					, dtype=tensor.dtype					, device=tensor.device							)
    _snake_case      : List[Any]														=		tensor
    return padded_tensor








'''simple docstring'''






from .hash_table import HashTable
from .number_theory.prime_numbers import is_prime, next_prime




class 		a__				(							a__  ):





							'''simple docstring'''






							def __init__(		self					,    *lowerCamelCase_					,    **lowerCamelCase_   )       ->							Optional[int]:
														super().__init__(*lowerCamelCase_					,    **lowerCamelCase_   )







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Optional[int]:
														lowerCAmelCase__          =	(
														    next_prime(value % self.size_table   )
														    if not is_prime(value % self.size_table   )
														    else value % self.size_table
														)  # gt = bigger than
														return next_prime_gt - (data % next_prime_gt)







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )       ->							Optional[int]:
														return (increment * self.__hash_function_a(lowerCamelCase_					,    lowerCamelCase_   )) % self.size_table







							def   __SCREAMING_SNAKE_CASE							(		self					,    lowerCamelCase_					,    lowerCamelCase_=None   )       ->							Any:
														lowerCAmelCase__          =	1
														lowerCAmelCase__          =	self.hash_function(lowerCamelCase_   )

														while self.values[new_key] is not None and self.values[new_key] != key:
																					lowerCAmelCase__          =	(
																					    self.__hash_double_function(lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_   )
																					    if self.balanced_factor() >= self.lim_charge
																					    else None
																					)
																					if new_key is None:
																												break
																					else:
																												i += 1

														return new_key