"""simple docstring"""




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

class        __SCREAMING_SNAKE_CASE     (   _lowerCAmelCase							):







				def __init__(    self    :List[Any]      ,__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      ,)      ->					Dict:




									"""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:
														lowerCamelCase__		:   Tuple  				=					self.load_generation_config(self.args.generation_config  )
														lowerCamelCase__		:   int  				=					gen_config







				@staticmethod
				def   lowercase_		(    __UpperCAmelCase    :Union[str, GenerationConfig]  )      ->					GenerationConfig:




									"""simple docstring"""



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

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

									lowerCamelCase__		:   Optional[int]  				=					GenerationConfig.from_pretrained(__UpperCAmelCase      ,__UpperCAmelCase  )
									return gen_config







				def   lowercase_		(    self    :List[Any]      ,__UpperCAmelCase    :Optional[Dataset] = None      ,__UpperCAmelCase    :Optional[List[str]] = None      ,__UpperCAmelCase    :str = "eval"      ,**__UpperCAmelCase    :Dict      ,)      ->					Dict[str, float]:




									"""simple docstring"""



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

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







				def   lowercase_		(    self    :int      ,__UpperCAmelCase    :Dataset      ,__UpperCAmelCase    :Optional[List[str]] = None      ,__UpperCAmelCase    :str = "test"      ,**__UpperCAmelCase    :List[str]      ,)      ->					"PredictionOutput":




									"""simple docstring"""



									lowerCamelCase__		:   Any  				=					gen_kwargs.copy()
									if gen_kwargs.get('''max_length'''  ) is None and gen_kwargs.get('''max_new_tokens'''  ) is None:
														lowerCamelCase__		:   Any  				=					self.args.generation_max_length
									lowerCamelCase__		:   Tuple  				=					(
									    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''  ) is not None else self.args.generation_num_beams
									)
									lowerCamelCase__		:   Dict  				=					gen_kwargs

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







				def   lowercase_		(    self    :List[str]      ,__UpperCAmelCase    :nn.Module      ,__UpperCAmelCase    :Dict[str, Union[torch.Tensor, Any]]      ,__UpperCAmelCase    :bool      ,__UpperCAmelCase    :Optional[List[str]] = None      ,**__UpperCAmelCase    :int      ,)      ->					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(
														    __UpperCAmelCase      ,__UpperCAmelCase      ,prediction_loss_only=__UpperCAmelCase      ,ignore_keys=__UpperCAmelCase  )

									lowerCamelCase__		:   Dict  				=					'''labels''' in inputs
									lowerCamelCase__		:   List[str]  				=					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'''  ):
														lowerCamelCase__		:   List[str]  				=					self._gen_kwargs.copy()

									if gen_kwargs.get('''max_length'''  ) is None and gen_kwargs.get('''max_new_tokens'''  ) is None:
														lowerCamelCase__		:   Union[str, Any]  				=					self.model.config.max_length
									lowerCamelCase__		:   List[Any]  				=					(
									    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''  ) is not None else self.model.config.num_beams
									)
									lowerCamelCase__		:   Optional[Any]  				=					True if is_deepspeed_zeroa_enabled() else False
									lowerCamelCase__		:   str  				=					(
									    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__		:   Union[str, Any]  				=					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:
														lowerCamelCase__		:   List[str]  				=					False

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

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

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

									return loss, generated_tokens, labels







				def   lowercase_		(    self    :Tuple      ,__UpperCAmelCase    :Optional[Any]      ,__UpperCAmelCase    :Optional[Any]  )      ->					Optional[int]:




									"""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
														lowerCamelCase__		:   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:
																			lowerCamelCase__		:   str  				=					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__		:   int  				=					pad_token_id * torch.ones(
									    (tensor.shape[0], max_length)      ,dtype=tensor.dtype      ,device=tensor.device  )
									lowerCamelCase__		:   Optional[Any]  				=					tensor
									return padded_tensor







"""simple docstring"""




import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
					import torch

					from transformers import GPTaLMHeadModel

@require_torch
@require_sentencepiece
@require_tokenizers
class        __SCREAMING_SNAKE_CASE     (   unittest.TestCase							):







				@slow
				@unittest.skip('''Model is not available.'''  )
				def   lowercase_		(    self    :Union[str, Any]  )      ->					Optional[int]:




									"""simple docstring"""



									lowerCamelCase__		:   Any  				=					'''nvidia/megatron-gpt2-345m/'''
									if "MYDIR" in os.environ:
														lowerCamelCase__		:   str  				=					os.path.join(os.environ['''MYDIR''']      ,__UpperCAmelCase  )
									lowerCamelCase__		:   Dict  				=					GPTaLMHeadModel.from_pretrained(__UpperCAmelCase  )
									model.to(__UpperCAmelCase  )
									model.half()

									lowerCamelCase__		:   Dict  				=					torch.tensor(
									    [[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      ,device=__UpperCAmelCase      ,dtype=torch.long      ,)

									with torch.no_grad():
														lowerCamelCase__		:   Tuple  				=					model(__UpperCAmelCase  ).logits

									lowerCamelCase__		:   List[str]  				=					torch.Size((1, 9, 5_02_57)  )
									self.assertEqual(output.shape      ,__UpperCAmelCase  )

									lowerCamelCase__		:   Tuple  				=					torch.tensor(
									    [
									        4.9_414,
									        -0.2_920,
									        -1.2_148,
									        -4.0_273,
									        -0.5_161,
									        -5.2_109,
									        -1.2_412,
									        -1.8_301,
									        -1.7_734,
									        -4.7_148,
									        -0.2_317,
									        -1.0_811,
									        -2.1_777,
									        0.4_141,
									        -3.7_969,
									        -4.0_586,
									        -2.5_332,
									        -3.3_809,
									        4.3_867,
									    ]      ,device=__UpperCAmelCase      ,dtype=torch.half      ,)

									for i in range(19  ):
														lowerCamelCase__				,						lowerCamelCase__		:   Dict  				=					8 * i // 17, 27_92 * i  # along the diagonal
														lowerCamelCase__				,						lowerCamelCase__		:   Optional[int]  				=					output[0, r, c], expected_diag[i]
														lowerCamelCase__		:   Dict  				=					F"""row={r} col={c} computed={computed} expected={expected}"""
														self.assertAlmostEqual(__UpperCAmelCase      ,__UpperCAmelCase      ,delta=1E-4      ,msg=__UpperCAmelCase  )






