




"""simple docstring"""





import logging
import os
import sys
import threading
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

from tqdm import auto as tqdm_lib


_lowercase	      =    threading.Lock()
_lowercase	      =    None

_lowercase	      =    {
    """debug""": logging.DEBUG,
    """info""": logging.INFO,
    """warning""": logging.WARNING,
    """error""": logging.ERROR,
    """critical""": logging.CRITICAL,
}

_lowercase	      =    logging.WARNING

_lowercase	      =    True

def   lowerCAmelCase							( ):
						'''simple docstring'''


						snake_case__       :					Dict										=					os.getenv('''DIFFUSERS_VERBOSITY'''      ,			lowerCamelCase		)
						if env_level_str:
												if env_level_str in log_levels:
																		return log_levels[env_level_str]
												else:
																		logging.getLogger().warning(
																		    f'''Unknown option DIFFUSERS_VERBOSITY={env_level_str}, '''
																		    f'''has to be one of: { ', '.join(log_levels.keys()		) }'''		)
						return _default_log_level

def   lowerCAmelCase							( ):
						'''simple docstring'''


						return __name__.split('''.'''		)[0]

def   lowerCAmelCase							( ):
						'''simple docstring'''


						return logging.getLogger(_get_library_name()		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						global _default_handler

						with _lock:
												if _default_handler:
																		# This library has already configured the library root logger.
																		return
												snake_case__       :					Optional[Any]										=					logging.StreamHandler()  # Set sys.stderr as stream.
												snake_case__       :					List[Any]										=					sys.stderr.flush

												# Apply our default configuration to the library root logger.
												snake_case__       :					List[Any]										=					_get_library_root_logger()
												library_root_logger.addHandler(_default_handler		)
												library_root_logger.setLevel(_get_default_logging_level()		)
												snake_case__       :					List[Any]										=					False

def   lowerCAmelCase							( ):
						'''simple docstring'''


						global _default_handler

						with _lock:
												if not _default_handler:
																		return

												snake_case__       :					Tuple										=					_get_library_root_logger()
												library_root_logger.removeHandler(_default_handler		)
												library_root_logger.setLevel(logging.NOTSET		)
												snake_case__       :					List[str]										=					None

def   lowerCAmelCase							( ):
						'''simple docstring'''


						return log_levels

def   lowerCAmelCase							( lowerCamelCase = None		):
						'''simple docstring'''


						if name is None:
												snake_case__       :					Optional[int]										=					_get_library_name()

						_configure_library_root_logger()
						return logging.getLogger(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						_configure_library_root_logger()
						return _get_library_root_logger().getEffectiveLevel()

def   lowerCAmelCase							( lowerCamelCase		):
						'''simple docstring'''


						_configure_library_root_logger()
						_get_library_root_logger().setLevel(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						return set_verbosity(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						return set_verbosity(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						return set_verbosity(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						return set_verbosity(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						_configure_library_root_logger()

						assert _default_handler is not None
						_get_library_root_logger().removeHandler(_default_handler		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						_configure_library_root_logger()

						assert _default_handler is not None
						_get_library_root_logger().addHandler(_default_handler		)

def   lowerCAmelCase							( lowerCamelCase		):
						'''simple docstring'''


						_configure_library_root_logger()

						assert handler is not None
						_get_library_root_logger().addHandler(lowerCamelCase		)

def   lowerCAmelCase							( lowerCamelCase		):
						'''simple docstring'''


						_configure_library_root_logger()

						assert handler is not None and handler not in _get_library_root_logger().handlers
						_get_library_root_logger().removeHandler(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						_configure_library_root_logger()
						snake_case__       :					str										=					False

def   lowerCAmelCase							( ):
						'''simple docstring'''


						_configure_library_root_logger()
						snake_case__       :					int										=					True

def   lowerCAmelCase							( ):
						'''simple docstring'''


						snake_case__       :					str										=					_get_library_root_logger().handlers

						for handler in handlers:
												snake_case__       :					int										=					logging.Formatter('''[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s'''		)
												handler.setFormatter(lowerCamelCase		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						snake_case__       :					Optional[Any]										=					_get_library_root_logger().handlers

						for handler in handlers:
												handler.setFormatter(lowerCamelCase		)

def   lowerCAmelCase							( self      ,			*lowerCamelCase      ,			**lowerCamelCase		):
						'''simple docstring'''


						snake_case__       :					Optional[Any]										=					os.getenv('''DIFFUSERS_NO_ADVISORY_WARNINGS'''      ,			lowerCamelCase		)
						if no_advisory_warnings:
												return
						self.warning(*lowerCamelCase      ,			**lowerCamelCase		)


_lowercase	      =    warning_advice

class       __lowercase       :


				'''simple docstring'''


				def __init__(   self   :     List[str]      ,*SCREAMING_SNAKE_CASE_   :     int      ,**SCREAMING_SNAKE_CASE_   :     Optional[Any]      ):  # pylint: disable=unused-argument
										snake_case__       :					List[str]										=					args[0] if args else None

				def __iter__(   self   :     Any      ):
										return iter(self._iterator      )

				def __getattr__(   self   :     Any      ,SCREAMING_SNAKE_CASE_   :     Dict      ):

										def empty_fn(*SCREAMING_SNAKE_CASE_   :     Optional[int]      ,**SCREAMING_SNAKE_CASE_   :     List[Any]      ):  # pylint: disable=unused-argument
																return

										return empty_fn

				def __enter__(   self   :     Dict      ):
										return self





				def __exit__(   self   :     List[Any]      ,SCREAMING_SNAKE_CASE_   :     List[Any]      ,SCREAMING_SNAKE_CASE_   :     Any      ,SCREAMING_SNAKE_CASE_   :     Union[str, Any]      ):
										return







class       __lowercase       :


				'''simple docstring'''


				def __call__(   self   :     Dict      ,*SCREAMING_SNAKE_CASE_   :     List[str]      ,**SCREAMING_SNAKE_CASE_   :     str      ):
										if _tqdm_active:
																return tqdm_lib.tqdm(*SCREAMING_SNAKE_CASE_      ,**SCREAMING_SNAKE_CASE_      )
										else:
																return EmptyTqdm(*SCREAMING_SNAKE_CASE_      ,**SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Optional[Any]      ,*SCREAMING_SNAKE_CASE_   :     Union[str, Any]      ,**SCREAMING_SNAKE_CASE_   :     List[str]      ):
										snake_case__       :					int										=					None
										if _tqdm_active:
																return tqdm_lib.tqdm.set_lock(*SCREAMING_SNAKE_CASE_      ,**SCREAMING_SNAKE_CASE_      )





				def       _lowerCamelCase					(   self   :     Dict      ):
										if _tqdm_active:
																return tqdm_lib.tqdm.get_lock()







_lowercase	      =    _tqdm_cls()

def   lowerCAmelCase							( ):
						'''simple docstring'''


						global _tqdm_active
						return bool(_tqdm_active		)

def   lowerCAmelCase							( ):
						'''simple docstring'''


						global _tqdm_active
						snake_case__       :					List[str]										=					True

def   lowerCAmelCase							( ):
						'''simple docstring'''


						global _tqdm_active
						snake_case__       :					int										=					False








"""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


_lowercase	      =    logging.get_logger(__name__)

class       __lowercase       (			snake_case_	):


				'''simple docstring'''


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

				@staticmethod
				def       _lowerCamelCase					(   SCREAMING_SNAKE_CASE_   :     Union[str, GenerationConfig]      ):

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

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

										snake_case__       :					Optional[Any]										=					GenerationConfig.from_pretrained(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      )
										return gen_config

				def       _lowerCamelCase					(   self   :     List[str]      ,SCREAMING_SNAKE_CASE_   :     Optional[Dataset] = None      ,SCREAMING_SNAKE_CASE_   :     Optional[List[str]] = None      ,SCREAMING_SNAKE_CASE_   :     str = "eval"      ,**SCREAMING_SNAKE_CASE_   :     int      ,):
										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__       :					Optional[int]										=					self.args.generation_max_length
										snake_case__       :					Optional[Any]										=					(
										    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''      ) is not None else self.args.generation_num_beams
										)
										snake_case__       :					Optional[Any]										=					gen_kwargs

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

				def       _lowerCamelCase					(   self   :     List[Any]      ,SCREAMING_SNAKE_CASE_   :     Dataset      ,SCREAMING_SNAKE_CASE_   :     Optional[List[str]] = None      ,SCREAMING_SNAKE_CASE_   :     str = "test"      ,**SCREAMING_SNAKE_CASE_   :     int      ,):
										snake_case__       :					Union[str, Any]										=					gen_kwargs.copy()
										if gen_kwargs.get('''max_length'''      ) is None and gen_kwargs.get('''max_new_tokens'''      ) is None:
																snake_case__       :					Any										=					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__       :					List[str]										=					gen_kwargs

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

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

										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__       :					List[Any]										=					'''labels''' in inputs
										snake_case__       :					Tuple										=					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__       :					Any										=					self._gen_kwargs.copy()

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

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

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

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

										return loss, generated_tokens, labels





				def       _lowerCamelCase					(   self   :     Dict      ,SCREAMING_SNAKE_CASE_   :     Dict      ,SCREAMING_SNAKE_CASE_   :     Dict      ):
										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__       :					List[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:
																						snake_case__       :					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__       :					str										=					pad_token_id * torch.ones(
										    (tensor.shape[0], max_length)      ,dtype=tensor.dtype      ,device=tensor.device      )
										snake_case__       :					str										=					tensor
										return padded_tensor


