

"""simple docstring"""





import logging
import os

import quant_trainer
import torch
from torch.utils.data import DataLoader

from transformers import Trainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput


_a							=							logging.getLogger(__name__)

if is_torch_tpu_available(check_device=False):
   import torch_xla.core.xla_model as xm
   import torch_xla.debug.metrics as met




class  _lowerCAmelCase  (     lowercase       ):







   """simple docstring"""







   def __init__(      self							:      Tuple, *UpperCAmelCase__							:      Optional[Any], UpperCAmelCase__							:      Dict=None, UpperCAmelCase__							:      int=None, UpperCAmelCase__							:      Tuple=None, **UpperCAmelCase__							:      List[Any]     ):
      super().__init__(*UpperCAmelCase__, **UpperCAmelCase__     )
      __lowercase								=						eval_examples
      __lowercase								=						post_process_function
      __lowercase								=						quant_trainer_args
      __lowercase								=						1_2_8  # default number of calibration samples


   def  _lowercase			(      self							:      Tuple, UpperCAmelCase__							:      List[str]=None     ):
      if calib_dataset is None and self.calib_dataset is None:
         raise ValueError("Trainer: calibration requires an calib_dataset."     )
      __lowercase								=						calib_dataset if calib_dataset is not None else self.calib_dataset

      __lowercase								=						self._remove_unused_columns(UpperCAmelCase__, description="Calibration"     )

      return DataLoader(
          UpperCAmelCase__, batch_size=self.args.eval_batch_size, collate_fn=self.data_collator, drop_last=self.args.dataloader_drop_last, num_workers=self.args.dataloader_num_workers, pin_memory=self.args.dataloader_pin_memory, shuffle=UpperCAmelCase__, )


   def  _lowercase			(      self							:      Union[str, Any], UpperCAmelCase__							:      Optional[int]=None     ):
      __lowercase								=						self.train_dataset if calib_dataset is None else calib_dataset
      __lowercase								=						self.get_calib_dataloader(UpperCAmelCase__     )

      __lowercase								=						self.model
      quant_trainer.configure_model(UpperCAmelCase__, self.quant_trainer_args, calib=UpperCAmelCase__     )
      model.eval()
      quant_trainer.enable_calibration(UpperCAmelCase__     )

      logger.info("***** Running calibration *****"     )
      logger.info(F"""  Num examples = {self.calib_num}"""     )
      logger.info(F"""  Batch size = {calib_dataloader.batch_size}"""     )

      for step, inputs in enumerate(UpperCAmelCase__     ):
         # Prediction step
         __lowercase		,__lowercase		,__lowercase								=						self.prediction_step(UpperCAmelCase__, UpperCAmelCase__, prediction_loss_only=UpperCAmelCase__     )
         if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
            break

      quant_trainer.finish_calibration(UpperCAmelCase__, self.quant_trainer_args     )
      __lowercase								=						model


   def  _lowercase			(      self							:      List[Any], UpperCAmelCase__							:      Tuple=None, UpperCAmelCase__							:      List[Any]=None, UpperCAmelCase__							:      Dict=None, UpperCAmelCase__							:      str = "eval"     ):
      __lowercase								=						self.eval_dataset if eval_dataset is None else eval_dataset
      __lowercase								=						self.get_eval_dataloader(UpperCAmelCase__     )
      __lowercase								=						self.eval_examples if eval_examples is None else eval_examples

      # Temporarily disable metric computation, we will do it in the loop here.
      __lowercase								=						self.compute_metrics
      __lowercase								=						None
      __lowercase								=						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
      try:
         __lowercase								=						eval_loop(
             UpperCAmelCase__, description="Evaluation", prediction_loss_only=True if compute_metrics is None else None, ignore_keys=UpperCAmelCase__, )
      finally:
         __lowercase								=						compute_metrics

      if self.post_process_function is not None and self.compute_metrics is not None:
         __lowercase								=						self.post_process_function(UpperCAmelCase__, UpperCAmelCase__, output.predictions     )
         __lowercase								=						self.compute_metrics(UpperCAmelCase__     )

         # Prefix all keys with metric_key_prefix + '_'
         for key in list(metrics.keys()     ):
            if not key.startswith(F"""{metric_key_prefix}_"""     ):
               __lowercase								=						metrics.pop(UpperCAmelCase__     )

         self.log(UpperCAmelCase__     )
      else:
         __lowercase								=						{}

      if self.args.tpu_metrics_debug or self.args.debug:
         # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
         xm.master_print(met.metrics_report()     )

      __lowercase								=						self.callback_handler.on_evaluate(self.args, self.state, self.control, UpperCAmelCase__     )
      return metrics


   def  _lowercase			(      self							:      Union[str, Any], UpperCAmelCase__							:      Dict, UpperCAmelCase__							:      Tuple, UpperCAmelCase__							:      Optional[int]=None, UpperCAmelCase__							:      str = "test"     ):
      __lowercase								=						self.get_test_dataloader(UpperCAmelCase__     )

      # Temporarily disable metric computation, we will do it in the loop here.
      __lowercase								=						self.compute_metrics
      __lowercase								=						None
      __lowercase								=						self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
      try:
         __lowercase								=						eval_loop(
             UpperCAmelCase__, description="Prediction", prediction_loss_only=True if compute_metrics is None else None, ignore_keys=UpperCAmelCase__, )
      finally:
         __lowercase								=						compute_metrics

      if self.post_process_function is None or self.compute_metrics is None:
         return output

      __lowercase								=						self.post_process_function(UpperCAmelCase__, UpperCAmelCase__, output.predictions, "predict"     )
      __lowercase								=						self.compute_metrics(UpperCAmelCase__     )

      # Prefix all keys with metric_key_prefix + '_'
      for key in list(metrics.keys()     ):
         if not key.startswith(F"""{metric_key_prefix}_"""     ):
            __lowercase								=						metrics.pop(UpperCAmelCase__     )

      return PredictionOutput(predictions=predictions.predictions, label_ids=predictions.label_ids, metrics=UpperCAmelCase__     )



   def  _lowercase			(      self							:      List[str], UpperCAmelCase__							:      Optional[int]="./"     ):
      __lowercase								=						self.eval_dataset
      __lowercase								=						self.get_eval_dataloader(UpperCAmelCase__     )

      __lowercase								=						next(iter(UpperCAmelCase__     )     )

      # saving device - to make it consistent
      __lowercase								=						torch.device("cuda" if torch.cuda.is_available() else "cpu"     )

      # convert to tuple
      __lowercase								=						tuple(v.to(UpperCAmelCase__     ) for k, v in batch.items()     )

      logger.info("Converting model to be onnx compatible"     )
      from pytorch_quantization.nn import TensorQuantizer

      __lowercase								=						True

      __lowercase								=						self.model.to(UpperCAmelCase__     )

      model.eval()
      model.float()

      __lowercase								=						model.module if hasattr(UpperCAmelCase__, "module"     ) else model
      quant_trainer.configure_model(UpperCAmelCase__, self.quant_trainer_args     )

      __lowercase								=						os.path.join(UpperCAmelCase__, "model.onnx"     )
      logger.info(F"""exporting model to {output_model_file}"""     )

      __lowercase								=						{0: "batch_size", 1: "seq_len"}

      torch.onnx.export(
          UpperCAmelCase__, UpperCAmelCase__, UpperCAmelCase__, export_params=UpperCAmelCase__, opset_version=1_3, do_constant_folding=UpperCAmelCase__, input_names=["input_ids", "attention_mask", "token_type_ids"], output_names=["output_start_logits", "output_end_logits"], dynamic_axes={
              "input_ids": axes,
              "attention_mask": axes,
              "token_type_ids": axes,
              "output_start_logits": axes,
              "output_end_logits": axes,
          }, verbose=UpperCAmelCase__, )
      logger.info("onnx export finished"     )







"""simple docstring"""





from ...utils import is_note_seq_available, is_transformers_available, is_torch_available
from ...utils import OptionalDependencyNotAvailable


try:
   if not (is_transformers_available() and is_torch_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
   from .notes_encoder import SpectrogramNotesEncoder
   from .continous_encoder import SpectrogramContEncoder
   from .pipeline_spectrogram_diffusion import (
       SpectrogramContEncoder,
       SpectrogramDiffusionPipeline,
       TaFilmDecoder,
   )

try:
   if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
   from .midi_utils import MidiProcessor




