
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


__SCREAMING_SNAKE_CASE			:   List[Any]         =     logging.get_logger(__name__)


class 		__magic_name__   (						snake_case    ):





		def __init__(							self    :	str						,  lowerCamelCase__    :	Union["PreTrainedModel", nn.Module] = None						,  lowerCamelCase__    :	"TrainingArguments" = None						,  lowerCamelCase__    :	Optional["DataCollator"] = None						,  lowerCamelCase__    :	Optional[Dataset] = None						,  lowerCamelCase__    :	Optional[Union[Dataset, Dict[str, Dataset]]] = None						,  lowerCamelCase__    :	Optional["PreTrainedTokenizerBase"] = None						,  lowerCamelCase__    :	Optional[Callable[[], "PreTrainedModel"]] = None						,  lowerCamelCase__    :	Optional[Callable[["EvalPrediction"], Dict]] = None						,  lowerCamelCase__    :	Optional[List["TrainerCallback"]] = None						,  lowerCamelCase__    :	Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)						,  lowerCamelCase__    :	Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None						,  ):
							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:
												lowerCAmelCase		:					int  	=     self.load_generation_config(self.args.generation_config				)
												lowerCAmelCase		:					Any  	=     gen_config





		@staticmethod
		def       _A							(							lowerCamelCase__    :	Union[str, GenerationConfig]				):

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

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

							lowerCAmelCase		:					int  	=     GenerationConfig.from_pretrained(lowerCamelCase__						,  lowerCamelCase__				)
							return gen_config





		def       _A							(							self    :	List[Any]						,  lowerCamelCase__    :	Optional[Dataset] = None						,  lowerCamelCase__    :	Optional[List[str]] = None						,  lowerCamelCase__    :	str = "eval"						,  **lowerCamelCase__    :	Union[str, Any]						,  ):
							lowerCAmelCase		:					Any  	=     gen_kwargs.copy()
							if gen_kwargs.get('''max_length'''				) is None and gen_kwargs.get('''max_new_tokens'''				) is None:
												lowerCAmelCase		:					Optional[Any]  	=     self.args.generation_max_length
							lowerCAmelCase		:					int  	=     (
							    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''				) is not None else self.args.generation_num_beams
							)
							lowerCAmelCase		:					int  	=     gen_kwargs

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





		def       _A							(							self    :	List[str]						,  lowerCamelCase__    :	Dataset						,  lowerCamelCase__    :	Optional[List[str]] = None						,  lowerCamelCase__    :	str = "test"						,  **lowerCamelCase__    :	int						,  ):
							lowerCAmelCase		:					Union[str, Any]  	=     gen_kwargs.copy()
							if gen_kwargs.get('''max_length'''				) is None and gen_kwargs.get('''max_new_tokens'''				) is None:
												lowerCAmelCase		:					Dict  	=     self.args.generation_max_length
							lowerCAmelCase		:					str  	=     (
							    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''				) is not None else self.args.generation_num_beams
							)
							lowerCAmelCase		:					List[Any]  	=     gen_kwargs

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





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

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

							lowerCAmelCase		:					Optional[Any]  	=     '''labels''' in inputs
							lowerCAmelCase		:					int  	=     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'''				):
												lowerCAmelCase		:					Any  	=     self._gen_kwargs.copy()

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

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

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

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

							return loss, generated_tokens, labels





		def       _A							(							self    :	Optional[int]						,  lowerCamelCase__    :	Dict						,  lowerCamelCase__    :	List[Any]				):
							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		:					str  	=     (
												    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		:					Union[str, Any]  	=     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		:					Optional[int]  	=     pad_token_id * torch.ones(
							    (tensor.shape[0], max_length)						,  dtype=tensor.dtype						,  device=tensor.device				)
							lowerCAmelCase		:					List[str]  	=     tensor
							return padded_tensor


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__SCREAMING_SNAKE_CASE			:   List[str]         =     {'configuration_vilt': ['VILT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ViltConfig']}

try:
		if not is_vision_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE			:   Any         =     ['ViltFeatureExtractor']
		__SCREAMING_SNAKE_CASE			:   Optional[int]         =     ['ViltImageProcessor']
		__SCREAMING_SNAKE_CASE			:   str         =     ['ViltProcessor']

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__SCREAMING_SNAKE_CASE			:   Dict         =     [
		    'VILT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'ViltForImageAndTextRetrieval',
		    'ViltForImagesAndTextClassification',
		    'ViltForTokenClassification',
		    'ViltForMaskedLM',
		    'ViltForQuestionAnswering',
		    'ViltLayer',
		    'ViltModel',
		    'ViltPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_vilt import VILT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViltConfig

		try:
				if not is_vision_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .feature_extraction_vilt import ViltFeatureExtractor
				from .image_processing_vilt import ViltImageProcessor
				from .processing_vilt import ViltProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_vilt import (
				    VILT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    ViltForImageAndTextRetrieval,
				    ViltForImagesAndTextClassification,
				    ViltForMaskedLM,
				    ViltForQuestionAnswering,
				    ViltForTokenClassification,
				    ViltLayer,
				    ViltModel,
				    ViltPreTrainedModel,
				)


else:
		import sys

		__SCREAMING_SNAKE_CASE			:   Optional[Any]         =     _LazyModule(__name__, globals()['__file__'], _import_structure)
