
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


A__       :   int   						=					logging.get_logger(__name__)





class 						lowercase       (  __UpperCamelCase					):
					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:
											lowerCAmelCase__ :    Dict				      =			self.load_generation_config(self.args.generation_config    )
											lowerCAmelCase__ :    List[str]				      =			gen_config
					@staticmethod
					def lowercase_			(		SCREAMING_SNAKE_CASE__    ):
								"""simple docstring"""







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

								# str or Path
								lowerCAmelCase__ :    Optional[Any]				      =			Path(SCREAMING_SNAKE_CASE__    ) if isinstance(SCREAMING_SNAKE_CASE__ ,	SCREAMING_SNAKE_CASE__    ) 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[Any]				      =			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(SCREAMING_SNAKE_CASE__ ,	SCREAMING_SNAKE_CASE__    )
								return gen_config
					def lowercase_			(		self ,	SCREAMING_SNAKE_CASE__ = None ,	SCREAMING_SNAKE_CASE__ = None ,	SCREAMING_SNAKE_CASE__ = "eval" ,	**SCREAMING_SNAKE_CASE__ ,	):
								"""simple docstring"""







								lowerCAmelCase__ :    Optional[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__ :    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(SCREAMING_SNAKE_CASE__ ,	ignore_keys=SCREAMING_SNAKE_CASE__ ,	metric_key_prefix=SCREAMING_SNAKE_CASE__    )
					def lowercase_			(		self ,	SCREAMING_SNAKE_CASE__ ,	SCREAMING_SNAKE_CASE__ = None ,	SCREAMING_SNAKE_CASE__ = "test" ,	**SCREAMING_SNAKE_CASE__ ,	):
								"""simple docstring"""







								lowerCAmelCase__ :    List[Any]				      =			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__ :    Union[str, Any]				      =			(
								    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(SCREAMING_SNAKE_CASE__ ,	ignore_keys=SCREAMING_SNAKE_CASE__ ,	metric_key_prefix=SCREAMING_SNAKE_CASE__    )
					def lowercase_			(		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__    )

								lowerCAmelCase__ :    Tuple				      =			'''labels''' in inputs
								lowerCAmelCase__ :    Dict				      =			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'''    ):
											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__ :    Tuple				      =			self.model.config.max_length
								lowerCAmelCase__ :    Tuple				      =			(
								    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''    ) is not None else self.model.config.num_beams
								)
								lowerCAmelCase__ :    str				      =			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(**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:
											lowerCAmelCase__ :    Union[str, Any]				      =			False

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

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

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

								return loss, generated_tokens, labels
					def lowercase_			(		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
											lowerCAmelCase__ :    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:
														lowerCAmelCase__ :    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[Any]				      =			pad_token_id * torch.ones(
								    (tensor.shape[0], max_length) ,	dtype=tensor.dtype ,	device=tensor.device    )
								lowerCAmelCase__ :    str				      =			tensor
								return padded_tensor





import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import Mock, patch

import diffusers.utils.hub_utils





class 						lowercase       (  unittest.TestCase					):
					@patch('''diffusers.utils.hub_utils.get_full_repo_name'''    )
					def lowercase_			(		self ,	SCREAMING_SNAKE_CASE__    ):
								"""simple docstring"""







								lowerCAmelCase__ :    str				      =			'''full_repo_name'''
								with TemporaryDirectory() as tmpdir:
											# Dummy args values
											lowerCAmelCase__ :    Optional[int]				      =			Mock()
											lowerCAmelCase__ :    Optional[int]				      =			tmpdir
											lowerCAmelCase__ :    Optional[Any]				      =			0
											lowerCAmelCase__ :    Optional[int]				      =			'''hub_token'''
											lowerCAmelCase__ :    Tuple				      =			'''dataset_name'''
											lowerCAmelCase__ :    int				      =			0.01
											lowerCAmelCase__ :    Optional[int]				      =			100000
											lowerCAmelCase__ :    Any				      =			10000
											lowerCAmelCase__ :    Union[str, Any]				      =			0.01
											lowerCAmelCase__ :    int				      =			0.02
											lowerCAmelCase__ :    Union[str, Any]				      =			0.03
											lowerCAmelCase__ :    str				      =			0.0_005
											lowerCAmelCase__ :    Dict				      =			0.000_001
											lowerCAmelCase__ :    Tuple				      =			1
											lowerCAmelCase__ :    Optional[int]				      =			10
											lowerCAmelCase__ :    int				      =			0.001
											lowerCAmelCase__ :    Union[str, Any]				      =			0.1
											lowerCAmelCase__ :    Tuple				      =			0.2
											lowerCAmelCase__ :    Optional[int]				      =			True

											# Model card mush be rendered and saved
											diffusers.utils.hub_utils.create_model_card(SCREAMING_SNAKE_CASE__ ,	model_name='''model_name'''    )
											self.assertTrue((Path(SCREAMING_SNAKE_CASE__    ) / '''README.md''').is_file()    )



