

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


lowercase_     =       logging.get_logger(__name__)





class 			__lowerCAmelCase					( SCREAMING_SNAKE_CASE							):




  def __init__(    self							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = (None, None)							,       lowerCAmelCase = None							,       )  ->			List[Any]:
        '''simple docstring'''






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

        # 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   A__					(    lowerCAmelCase )  ->			GenerationConfig:
        '''simple docstring'''






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

        # str or Path
        _lowercase     	=Path(lowerCAmelCase ) if isinstance(lowerCAmelCase							,       lowerCAmelCase ) 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(lowerCAmelCase							,       lowerCAmelCase )
        return gen_config




  def   A__					(    self							,       lowerCAmelCase = None							,       lowerCAmelCase = None							,       lowerCAmelCase = "eval"							,       **lowerCAmelCase							,       )  ->			Dict[str, float]:
        '''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(lowerCAmelCase							,       ignore_keys=lowerCAmelCase							,       metric_key_prefix=lowerCAmelCase )




  def   A__					(    self							,       lowerCAmelCase							,       lowerCAmelCase = None							,       lowerCAmelCase = "test"							,       **lowerCAmelCase							,       )  ->			"PredictionOutput":
        '''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(lowerCAmelCase							,       ignore_keys=lowerCAmelCase							,       metric_key_prefix=lowerCAmelCase )




  def   A__					(    self							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase							,       lowerCAmelCase = None							,       **lowerCAmelCase							,       )  ->			Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
        '''simple docstring'''






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

        _lowercase     	='labels' in inputs
        _lowercase     	=self._prepare_inputs(lowerCAmelCase )

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

        if len(lowerCAmelCase ) == 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(**lowerCAmelCase							,       **lowerCAmelCase )

        # 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(lowerCAmelCase							,       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(lowerCAmelCase							,       gen_config.max_new_tokens + 1 )

        with torch.no_grad():
              if has_labels:
                    with self.compute_loss_context_manager():
                          _lowercase     	=model(**lowerCAmelCase )
                    if self.label_smoother is not None:
                          _lowercase     	=self.label_smoother(lowerCAmelCase							,       inputs['labels'] ).mean().detach()
                    else:
                          _lowercase     	=(outputs['loss'] if isinstance(lowerCAmelCase							,       lowerCAmelCase ) 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(lowerCAmelCase							,       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(lowerCAmelCase							,       gen_config.max_new_tokens + 1 )
        else:
              _lowercase     	=None

        return loss, generated_tokens, labels




  def   A__					(    self							,       lowerCAmelCase							,       lowerCAmelCase )  ->			Union[str, Any]:
        '''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





import warnings

from ...utils import logging
from .image_processing_layoutlmva import LayoutLMvaImageProcessor


lowercase_     =       logging.get_logger(__name__)





class 			__lowerCAmelCase					( SCREAMING_SNAKE_CASE							):




  def __init__(    self							,       *lowerCAmelCase							,       **lowerCAmelCase )  ->			None:
        '''simple docstring'''






        warnings.warn(
            'The class LayoutLMv3FeatureExtractor is deprecated and will be removed in version 5 of Transformers.'
            ' Please use LayoutLMv3ImageProcessor instead.'							,       lowerCAmelCase							,       )
        super().__init__(*lowerCAmelCase							,       **lowerCAmelCase )


