

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


lowerCamelCase 	=   logging.get_logger(__name__)







class 		__magic_name__     (				lowerCamelCase__				):

   '''simple docstring'''




   def __init__(							self,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = None,       lowercase_ = (None, None),       lowercase_ = None,       )				->       List[Any]:






     """simple docstring"""






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

     # 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:
       a__    =self.load_generation_config(self.args.generation_config			)
       a__    =gen_config



   @staticmethod
   def 					_UpperCAmelCase					(							lowercase_			)				->       GenerationConfig:






     """simple docstring"""






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

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

     a__    =GenerationConfig.from_pretrained(lowercase_,       lowercase_			)
     return gen_config



   def 					_UpperCAmelCase					(							self,       lowercase_ = None,       lowercase_ = None,       lowercase_ = "eval",       **lowercase_,       )				->       Dict[str, float]:






     """simple docstring"""






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

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



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_ = None,       lowercase_ = "test",       **lowercase_,       )				->       "PredictionOutput":






     """simple docstring"""






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

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



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_,       lowercase_ = None,       **lowercase_,       )				->       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(
           lowercase_,       lowercase_,       prediction_loss_only=lowercase_,       ignore_keys=lowercase_			)

     a__    ='''labels''' in inputs
     a__    =self._prepare_inputs(lowercase_			)

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

     if len(lowercase_			) == 0 and hasattr(self,       '''_gen_kwargs'''			):
       a__    =self._gen_kwargs.copy()

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

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

     # Retrieves GenerationConfig from model.generation_config
     a__    =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:
       a__    =self._pad_tensors_to_max_len(lowercase_,       gen_config.max_length			)
     elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
       a__    =self._pad_tensors_to_max_len(lowercase_,       gen_config.max_new_tokens + 1			)

     with torch.no_grad():
       if has_labels:
         with self.compute_loss_context_manager():
           a__    =model(**lowercase_			)
         if self.label_smoother is not None:
           a__    =self.label_smoother(lowercase_,       inputs['''labels''']			).mean().detach()
         else:
           a__    =(outputs['''loss'''] if isinstance(lowercase_,       lowercase_			) else outputs[0]).mean().detach()
       else:
         a__    =None

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

     if has_labels:
       a__    =inputs['''labels''']
       if labels.shape[-1] < gen_config.max_length:
         a__    =self._pad_tensors_to_max_len(lowercase_,       gen_config.max_length			)
       elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
         a__    =self._pad_tensors_to_max_len(lowercase_,       gen_config.max_new_tokens + 1			)
     else:
       a__    =None

     return loss, generated_tokens, labels



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_			)				->       List[Any]:






     """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
       a__    =(
           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:
         a__    =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'''			)

     a__    =pad_token_id * torch.ones(
         (tensor.shape[0], max_length),       dtype=tensor.dtype,       device=tensor.device			)
     a__    =tensor
     return padded_tensor



import unittest

import numpy as np
import PIL
import torch

from diffusers.image_processor import VaeImageProcessor







class 		__magic_name__     (				unittest.TestCase				):

   '''simple docstring'''




   @property
   def 					_UpperCAmelCase					(							self			)				->       Tuple:






     """simple docstring"""






     a__    =1
     a__    =3
     a__    =8
     a__    =8

     a__    =torch.rand((batch_size, num_channels, height, width)			)

     return sample



   def 					_UpperCAmelCase					(							self,       lowercase_			)				->       int:






     """simple docstring"""






     if isinstance(image[0],       PIL.Image.Image			):
       return np.stack([np.array(lowercase_			) for i in image],       axis=0			)
     elif isinstance(lowercase_,       torch.Tensor			):
       return image.cpu().numpy().transpose(0,       2,       3,       1			)
     return image



   def 					_UpperCAmelCase					(							self			)				->       int:






     """simple docstring"""






     a__    =VaeImageProcessor(do_resize=lowercase_,       do_normalize=lowercase_			)

     a__    =self.dummy_sample
     a__    =self.to_np(lowercase_			)

     for output_type in ["pt", "np", "pil"]:
       a__    =image_processor.postprocess(
           image_processor.preprocess(lowercase_			),       output_type=lowercase_,       )
       a__    =self.to_np(lowercase_			)
       a__    =(input_np * 255).round() if output_type == '''pil''' else input_np
       assert (
           np.abs(in_np - out_np			).max() < 1E-6
       ), F"""decoded output does not match input for output_type {output_type}"""



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =VaeImageProcessor(do_resize=lowercase_,       do_normalize=lowercase_			)
     a__    =self.dummy_sample.cpu().numpy().transpose(0,       2,       3,       1			)

     for output_type in ["pt", "np", "pil"]:
       a__    =image_processor.postprocess(image_processor.preprocess(lowercase_			),       output_type=lowercase_			)

       a__    =self.to_np(lowercase_			)
       a__    =(input_np * 255).round() if output_type == '''pil''' else input_np
       assert (
           np.abs(in_np - out_np			).max() < 1E-6
       ), F"""decoded output does not match input for output_type {output_type}"""



   def 					_UpperCAmelCase					(							self			)				->       Optional[Any]:






     """simple docstring"""






     a__    =VaeImageProcessor(do_resize=lowercase_,       do_normalize=lowercase_			)

     a__    =self.dummy_sample.cpu().numpy().transpose(0,       2,       3,       1			)
     a__    =image_processor.numpy_to_pil(lowercase_			)

     for output_type in ["pt", "np", "pil"]:
       a__    =image_processor.postprocess(image_processor.preprocess(lowercase_			),       output_type=lowercase_			)
       for i, o in zip(lowercase_,       lowercase_			):
         a__    =np.array(lowercase_			)
         a__    =self.to_np(lowercase_			) if output_type == '''pil''' else (self.to_np(lowercase_			) * 255).round()
         assert (
             np.abs(in_np - out_np			).max() < 1E-6
         ), F"""decoded output does not match input for output_type {output_type}"""



   def 					_UpperCAmelCase					(							self			)				->       Union[str, Any]:






     """simple docstring"""






     a__    =VaeImageProcessor(do_resize=lowercase_,       do_normalize=lowercase_			)

     a__    =self.dummy_sample
     a__    =input_pt_ad.squeeze(0			)

     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )
     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )

     a__    =self.to_np(self.dummy_sample			)
     a__    =input_np_ad.squeeze(0			)

     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )
     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )

     assert np.abs(out_pt_ad - out_pt_ad			).max() < 1E-6
     assert np.abs(out_np_ad - out_np_ad			).max() < 1E-6



   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     a__    =VaeImageProcessor(do_resize=lowercase_,       do_normalize=lowercase_			)

     a__    =self.dummy_sample
     a__    =list(lowercase_			)

     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )

     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )

     a__    =self.to_np(self.dummy_sample			)
     list(lowercase_			)

     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )

     a__    =image_processor.postprocess(
         image_processor.preprocess(lowercase_			),       output_type='''np''',       )

     assert np.abs(out_pt_ad - out_pt_list			).max() < 1E-6
     assert np.abs(out_np_ad - out_np_list			).max() < 1E-6
