



'''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_							: Optional[int]						 =	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 							lowercase       (  _lowerCamelCase				):


 """simple docstring"""


 def __init__(							self  ,*a_  ,a_=None  ,a_=None  ,a_=None  ,**a_							)		->   Optional[Any]:
        super().__init__(*a_  ,**a_							)
        _UpperCAmelCase  :						List[str]			     = eval_examples
        _UpperCAmelCase  :						Optional[int]			     = post_process_function
        _UpperCAmelCase  :						List[Any]			     = quant_trainer_args
        _UpperCAmelCase  :						Dict			     = 128  # default number of calibration samples







 def _snake_case (							self  ,a_=None							)		->   Union[str, Any]:
        if calib_dataset is None and self.calib_dataset is None:
               raise ValueError("""Trainer: calibration requires an calib_dataset."""							)
        _UpperCAmelCase  :						Tuple			     = calib_dataset if calib_dataset is not None else self.calib_dataset

        _UpperCAmelCase  :						Union[str, Any]			     = self._remove_unused_columns(a_  ,description="""Calibration"""							)

        return DataLoader(
            a_  ,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=a_  ,)







 def _snake_case (							self  ,a_=None							)		->   Union[str, Any]:
        _UpperCAmelCase  :						List[str]			     = self.train_dataset if calib_dataset is None else calib_dataset
        _UpperCAmelCase  :						Any			     = self.get_calib_dataloader(a_							)

        _UpperCAmelCase  :						Tuple			     = self.model
        quant_trainer.configure_model(a_  ,self.quant_trainer_args  ,calib=a_							)
        model.eval()
        quant_trainer.enable_calibration(a_							)

        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(a_							):
               # Prediction step
               _UpperCAmelCase     ,_UpperCAmelCase     ,_UpperCAmelCase  :						int			     = self.prediction_step(a_  ,a_  ,prediction_loss_only=a_							)
               if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                      break

        quant_trainer.finish_calibration(a_  ,self.quant_trainer_args							)
        _UpperCAmelCase  :						Union[str, Any]			     = model







 def _snake_case (							self  ,a_=None  ,a_=None  ,a_=None  ,a_ = "eval"							)		->   Optional[int]:
        _UpperCAmelCase  :						List[Any]			     = self.eval_dataset if eval_dataset is None else eval_dataset
        _UpperCAmelCase  :						Optional[int]			     = self.get_eval_dataloader(a_							)
        _UpperCAmelCase  :						List[str]			     = self.eval_examples if eval_examples is None else eval_examples

        # Temporarily disable metric computation, we will do it in the loop here.
        _UpperCAmelCase  :						Tuple			     = self.compute_metrics
        _UpperCAmelCase  :						str			     = None
        _UpperCAmelCase  :						Union[str, Any]			     = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
               _UpperCAmelCase  :						Union[str, Any]			     = eval_loop(
                   a_  ,description="""Evaluation"""  ,prediction_loss_only=True if compute_metrics is None else None  ,ignore_keys=a_  ,)
        finally:
               _UpperCAmelCase  :						Tuple			     = compute_metrics

        if self.post_process_function is not None and self.compute_metrics is not None:
               _UpperCAmelCase  :						Dict			     = self.post_process_function(a_  ,a_  ,output.predictions							)
               _UpperCAmelCase  :						List[str]			     = self.compute_metrics(a_							)

               # Prefix all keys with metric_key_prefix + '_'
               for key in list(metrics.keys()							):
                      if not key.startswith(f'''{metric_key_prefix}_'''							):
                             _UpperCAmelCase  :						Tuple			     = metrics.pop(a_							)

               self.log(a_							)
        else:
               _UpperCAmelCase  :						List[str]			     = {}

        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()							)

        _UpperCAmelCase  :						Optional[int]			     = self.callback_handler.on_evaluate(self.args  ,self.state  ,self.control  ,a_							)
        return metrics







 def _snake_case (							self  ,a_  ,a_  ,a_=None  ,a_ = "test"							)		->   List[str]:
        _UpperCAmelCase  :						Tuple			     = self.get_test_dataloader(a_							)

        # Temporarily disable metric computation, we will do it in the loop here.
        _UpperCAmelCase  :						Any			     = self.compute_metrics
        _UpperCAmelCase  :						Dict			     = None
        _UpperCAmelCase  :						Union[str, Any]			     = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
               _UpperCAmelCase  :						Optional[int]			     = eval_loop(
                   a_  ,description="""Prediction"""  ,prediction_loss_only=True if compute_metrics is None else None  ,ignore_keys=a_  ,)
        finally:
               _UpperCAmelCase  :						Dict			     = compute_metrics

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

        _UpperCAmelCase  :						List[str]			     = self.post_process_function(a_  ,a_  ,output.predictions  ,"""predict"""							)
        _UpperCAmelCase  :						Dict			     = self.compute_metrics(a_							)

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()							):
               if not key.startswith(f'''{metric_key_prefix}_'''							):
                      _UpperCAmelCase  :						List[str]			     = metrics.pop(a_							)

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







 def _snake_case (							self  ,a_="./"							)		->   Union[str, Any]:
        _UpperCAmelCase  :						List[Any]			     = self.eval_dataset
        _UpperCAmelCase  :						Tuple			     = self.get_eval_dataloader(a_							)

        _UpperCAmelCase  :						Tuple			     = next(iter(a_							)							)

        # saving device - to make it consistent
        _UpperCAmelCase  :						List[str]			     = torch.device("""cuda""" if torch.cuda.is_available() else """cpu"""							)

        # convert to tuple
        _UpperCAmelCase  :						List[str]			     = tuple(v.to(a_							) for k, v in batch.items()							)

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

        _UpperCAmelCase  :						Tuple			     = True

        _UpperCAmelCase  :						str			     = self.model.to(a_							)

        model.eval()
        model.float()

        _UpperCAmelCase  :						Any			     = model.module if hasattr(a_  ,"""module"""							) else model
        quant_trainer.configure_model(a_  ,self.quant_trainer_args							)

        _UpperCAmelCase  :						Any			     = os.path.join(a_  ,"""model.onnx"""							)
        logger.info(f'''exporting model to {output_model_file}'''							)

        _UpperCAmelCase  :						Optional[int]			     = {0: """batch_size""", 1: """seq_len"""}

        torch.onnx.export(
            a_  ,a_  ,a_  ,export_params=a_  ,opset_version=13  ,do_constant_folding=a_  ,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=a_  ,)
        logger.info("""onnx export finished"""							)








'''simple docstring'''






import shutil
import tempfile
import unittest

from transformers import ClapFeatureExtractor, ClapProcessor, RobertaTokenizer, RobertaTokenizerFast
from transformers.testing_utils import require_sentencepiece, require_torchaudio

from .test_feature_extraction_clap import floats_list



@require_torchaudio
@require_sentencepiece
class 							lowercase       (  unittest.TestCase				):


 """simple docstring"""


 def _snake_case (							self							)		->   Optional[Any]:
        _UpperCAmelCase  :						str			     = """laion/clap-htsat-unfused"""
        _UpperCAmelCase  :						int			     = tempfile.mkdtemp()







 def _snake_case (							self  ,**a_							)		->   str:
        return RobertaTokenizer.from_pretrained(self.checkpoint  ,**a_							)







 def _snake_case (							self  ,**a_							)		->   Tuple:
        return ClapFeatureExtractor.from_pretrained(self.checkpoint  ,**a_							)







 def _snake_case (							self							)		->   int:
        shutil.rmtree(self.tmpdirname							)







 def _snake_case (							self							)		->   Optional[int]:
        _UpperCAmelCase  :						str			     = self.get_tokenizer()
        _UpperCAmelCase  :						Any			     = self.get_feature_extractor()

        _UpperCAmelCase  :						int			     = ClapProcessor(tokenizer=a_  ,feature_extractor=a_							)

        processor.save_pretrained(self.tmpdirname							)
        _UpperCAmelCase  :						List[Any]			     = ClapProcessor.from_pretrained(self.tmpdirname							)

        self.assertEqual(processor.tokenizer.get_vocab()  ,tokenizer.get_vocab()							)
        self.assertIsInstance(processor.tokenizer  ,a_							)

        self.assertEqual(processor.feature_extractor.to_json_string()  ,feature_extractor.to_json_string()							)
        self.assertIsInstance(processor.feature_extractor  ,a_							)







 def _snake_case (							self							)		->   List[Any]:
        _UpperCAmelCase  :						int			     = ClapProcessor(tokenizer=self.get_tokenizer()  ,feature_extractor=self.get_feature_extractor()							)
        processor.save_pretrained(self.tmpdirname							)

        _UpperCAmelCase  :						List[str]			     = self.get_tokenizer(bos_token="""(BOS)"""  ,eos_token="""(EOS)"""							)
        _UpperCAmelCase  :						List[Any]			     = self.get_feature_extractor(do_normalize=a_  ,padding_value=1.0							)

        _UpperCAmelCase  :						Optional[Any]			     = ClapProcessor.from_pretrained(
            self.tmpdirname  ,bos_token="""(BOS)"""  ,eos_token="""(EOS)"""  ,do_normalize=a_  ,padding_value=1.0							)

        self.assertEqual(processor.tokenizer.get_vocab()  ,tokenizer_add_kwargs.get_vocab()							)
        self.assertIsInstance(processor.tokenizer  ,a_							)

        self.assertEqual(processor.feature_extractor.to_json_string()  ,feature_extractor_add_kwargs.to_json_string()							)
        self.assertIsInstance(processor.feature_extractor  ,a_							)







 def _snake_case (							self							)		->   str:
        _UpperCAmelCase  :						Tuple			     = self.get_feature_extractor()
        _UpperCAmelCase  :						Dict			     = self.get_tokenizer()

        _UpperCAmelCase  :						str			     = ClapProcessor(tokenizer=a_  ,feature_extractor=a_							)

        _UpperCAmelCase  :						Tuple			     = floats_list((3, 1_000)							)

        _UpperCAmelCase  :						int			     = feature_extractor(a_  ,return_tensors="""np"""							)
        _UpperCAmelCase  :						Union[str, Any]			     = processor(audios=a_  ,return_tensors="""np"""							)

        for key in input_feat_extract.keys():
               self.assertAlmostEqual(input_feat_extract[key].sum()  ,input_processor[key].sum()  ,delta=1E-2							)







 def _snake_case (							self							)		->   Tuple:
        _UpperCAmelCase  :						List[Any]			     = self.get_feature_extractor()
        _UpperCAmelCase  :						Any			     = self.get_tokenizer()

        _UpperCAmelCase  :						Optional[int]			     = ClapProcessor(tokenizer=a_  ,feature_extractor=a_							)

        _UpperCAmelCase  :						Union[str, Any]			     = """This is a test string"""

        _UpperCAmelCase  :						Optional[Any]			     = processor(text=a_							)

        _UpperCAmelCase  :						Any			     = tokenizer(a_							)

        for key in encoded_tok.keys():
               self.assertListEqual(encoded_tok[key]  ,encoded_processor[key]							)







 def _snake_case (							self							)		->   List[Any]:
        _UpperCAmelCase  :						str			     = self.get_feature_extractor()
        _UpperCAmelCase  :						List[str]			     = self.get_tokenizer()

        _UpperCAmelCase  :						Any			     = ClapProcessor(tokenizer=a_  ,feature_extractor=a_							)

        _UpperCAmelCase  :						List[str]			     = [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

        _UpperCAmelCase  :						Dict			     = processor.batch_decode(a_							)
        _UpperCAmelCase  :						Any			     = tokenizer.batch_decode(a_							)

        self.assertListEqual(a_  ,a_							)







 def _snake_case (							self							)		->   Dict:
        _UpperCAmelCase  :						List[str]			     = self.get_feature_extractor()
        _UpperCAmelCase  :						int			     = self.get_tokenizer()

        _UpperCAmelCase  :						Dict			     = ClapProcessor(tokenizer=a_  ,feature_extractor=a_							)

        self.assertListEqual(
            processor.model_input_names[2:]  ,feature_extractor.model_input_names  ,msg="""`processor` and `feature_extractor` model input names do not match"""  ,)



