



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_			      =  logging.get_logger(__name__)





class 		__SCREAMING_SNAKE_CASE      (				lowercase__      ):






				"""simple docstring"""

				def __init__(				self						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = (None, None)						,	SCREAMING_SNAKE_CASE__ = None						,	):






					"""simple docstring"""

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

					# 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					:							str					     =   self.load_generation_config(self.args.generation_config      )
						_snake_case					:							Union[str, Any]					     =   gen_config







				@staticmethod
				def 		__lowerCamelCase(				SCREAMING_SNAKE_CASE__      ):






					"""simple docstring"""

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

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

					_snake_case					:							int					     =   GenerationConfig.from_pretrained(SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      )
					return gen_config







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = "eval"						,	**SCREAMING_SNAKE_CASE__						,	):






					"""simple docstring"""

					_snake_case					:							Optional[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					:							str					     =   (
					    gen_kwargs["""num_beams"""] if gen_kwargs.get("""num_beams"""      ) is not None else self.args.generation_num_beams
					)
					_snake_case					:							Tuple					     =   gen_kwargs

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







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__ = None						,	SCREAMING_SNAKE_CASE__ = "test"						,	**SCREAMING_SNAKE_CASE__						,	):






					"""simple docstring"""

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

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







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__ = None						,	**SCREAMING_SNAKE_CASE__						,	):






					"""simple docstring"""

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

					_snake_case					:							int					     =   """labels""" in inputs
					_snake_case					:							int					     =   self._prepare_inputs(SCREAMING_SNAKE_CASE__      )

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

					if len(SCREAMING_SNAKE_CASE__      ) == 0 and hasattr(self						,	"""_gen_kwargs"""      ):
						_snake_case					:							Optional[Any]					     =   self._gen_kwargs.copy()

					if gen_kwargs.get("""max_length"""      ) is None and gen_kwargs.get("""max_new_tokens"""      ) is None:
						_snake_case					:							Union[str, Any]					     =   self.model.config.max_length
					_snake_case					:							Optional[int]					     =   (
					    gen_kwargs["""num_beams"""] if gen_kwargs.get("""num_beams"""      ) is not None else self.model.config.num_beams
					)
					_snake_case					:							Union[str, Any]					     =   True if is_deepspeed_zeroa_enabled() else False
					_snake_case					:							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					:							Union[str, Any]					     =   self.model.generate(**SCREAMING_SNAKE_CASE__						,	**SCREAMING_SNAKE_CASE__      )

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

					# Retrieves GenerationConfig from model.generation_config
					_snake_case					:							Union[str, 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					:							Optional[int]					     =   self._pad_tensors_to_max_len(SCREAMING_SNAKE_CASE__						,	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					:							Optional[Any]					     =   self._pad_tensors_to_max_len(SCREAMING_SNAKE_CASE__						,	gen_config.max_new_tokens + 1      )

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

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

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

					return loss, generated_tokens, labels







				def 		__lowerCamelCase(				self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__      ):






					"""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					:							int					     =   (
						    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					:							List[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"""      )

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










from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase_			      =  {'''configuration_qdqbert''': ['''QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''QDQBertConfig''']}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						UpperCAmelCase_			      =  [
						    '''QDQBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''QDQBertForMaskedLM''',
						    '''QDQBertForMultipleChoice''',
						    '''QDQBertForNextSentencePrediction''',
						    '''QDQBertForQuestionAnswering''',
						    '''QDQBertForSequenceClassification''',
						    '''QDQBertForTokenClassification''',
						    '''QDQBertLayer''',
						    '''QDQBertLMHeadModel''',
						    '''QDQBertModel''',
						    '''QDQBertPreTrainedModel''',
						    '''load_tf_weights_in_qdqbert''',
						]


if TYPE_CHECKING:
						from .configuration_qdqbert import QDQBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, QDQBertConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_qdqbert import (
												    QDQBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
												    QDQBertForMaskedLM,
												    QDQBertForMultipleChoice,
												    QDQBertForNextSentencePrediction,
												    QDQBertForQuestionAnswering,
												    QDQBertForSequenceClassification,
												    QDQBertForTokenClassification,
												    QDQBertLayer,
												    QDQBertLMHeadModel,
												    QDQBertModel,
												    QDQBertPreTrainedModel,
												    load_tf_weights_in_qdqbert,
												)


else:
						import sys

						UpperCAmelCase_			      =  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





