



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


__UpperCAmelCase			  =			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_						(    lowerCAmelCase__				):
       def __init__(     self,     *SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_=None,     SCREAMING_SNAKE_CASE_=None,     SCREAMING_SNAKE_CASE_=None,     **SCREAMING_SNAKE_CASE_			)	->						Optional[int]:
          super().__init__(*__lowerCAmelCase,     **__lowerCAmelCase			)
          UpperCamelCase       :    Union[str, Any]		      =      eval_examples
          UpperCamelCase       :    List[str]		      =      post_process_function
          UpperCamelCase       :    Union[str, Any]		      =      quant_trainer_args
          UpperCamelCase       :    Union[str, Any]		      =      128  # default number of calibration samples
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_=None			)	->						Optional[int]:


          if calib_dataset is None and self.calib_dataset is None:
             raise ValueError('Trainer: calibration requires an calib_dataset.'			)
          UpperCamelCase       :    Dict		      =      calib_dataset if calib_dataset is not None else self.calib_dataset

          UpperCamelCase       :    Optional[int]		      =      self._remove_unused_columns(__lowerCAmelCase,     description='Calibration'			)

          return DataLoader(
              __lowerCAmelCase,     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=__lowerCAmelCase,     )
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_=None			)	->						Dict:
          UpperCamelCase       :    int		      =      self.train_dataset if calib_dataset is None else calib_dataset
          UpperCamelCase       :    int		      =      self.get_calib_dataloader(__lowerCAmelCase			)

          UpperCamelCase       :    List[Any]		      =      self.model
          quant_trainer.configure_model(__lowerCAmelCase,     self.quant_trainer_args,     calib=__lowerCAmelCase			)
          model.eval()
          quant_trainer.enable_calibration(__lowerCAmelCase			)

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

          quant_trainer.finish_calibration(__lowerCAmelCase,     self.quant_trainer_args			)
          UpperCamelCase       :    int		      =      model
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_=None,     SCREAMING_SNAKE_CASE_=None,     SCREAMING_SNAKE_CASE_=None,     SCREAMING_SNAKE_CASE_ = "eval"			)	->						Optional[Any]:
          UpperCamelCase       :    Optional[Any]		      =      self.eval_dataset if eval_dataset is None else eval_dataset
          UpperCamelCase       :    int		      =      self.get_eval_dataloader(__lowerCAmelCase			)
          UpperCamelCase       :    Dict		      =      self.eval_examples if eval_examples is None else eval_examples

          # Temporarily disable metric computation, we will do it in the loop here.
          UpperCamelCase       :    str		      =      self.compute_metrics
          UpperCamelCase       :    int		      =      None
          UpperCamelCase       :    Tuple		      =      self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             UpperCamelCase       :    Any		      =      eval_loop(
                 __lowerCAmelCase,     description='Evaluation',     prediction_loss_only=True if compute_metrics is None else None,     ignore_keys=__lowerCAmelCase,     )
          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(__lowerCAmelCase,     __lowerCAmelCase,     output.predictions			)
             UpperCamelCase       :    List[Any]		      =      self.compute_metrics(__lowerCAmelCase			)

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

             self.log(__lowerCAmelCase			)
          else:
             UpperCamelCase       :    Optional[int]		      =      {}

          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,     __lowerCAmelCase			)
          return metrics
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_=None,     SCREAMING_SNAKE_CASE_ = "test"			)	->						Optional[Any]:
          UpperCamelCase       :    List[str]		      =      self.get_test_dataloader(__lowerCAmelCase			)

          # Temporarily disable metric computation, we will do it in the loop here.
          UpperCamelCase       :    Dict		      =      self.compute_metrics
          UpperCamelCase       :    str		      =      None
          UpperCamelCase       :    Tuple		      =      self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
          try:
             UpperCamelCase       :    str		      =      eval_loop(
                 __lowerCAmelCase,     description='Prediction',     prediction_loss_only=True if compute_metrics is None else None,     ignore_keys=__lowerCAmelCase,     )
          finally:
             UpperCamelCase       :    List[str]		      =      compute_metrics

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

          UpperCamelCase       :    Optional[int]		      =      self.post_process_function(__lowerCAmelCase,     __lowerCAmelCase,     output.predictions,     'predict'			)
          UpperCamelCase       :    int		      =      self.compute_metrics(__lowerCAmelCase			)

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

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






       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_="./"			)	->						Optional[Any]:
          UpperCamelCase       :    Any		      =      self.eval_dataset
          UpperCamelCase       :    Dict		      =      self.get_eval_dataloader(__lowerCAmelCase			)

          UpperCamelCase       :    Union[str, Any]		      =      next(iter(__lowerCAmelCase			)			)

          # saving device - to make it consistent
          UpperCamelCase       :    Optional[Any]		      =      torch.device('cuda' if torch.cuda.is_available() else 'cpu'			)

          # convert to tuple
          UpperCamelCase       :    Dict		      =      tuple(v.to(__lowerCAmelCase			) for k, v in batch.items()			)

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

          UpperCamelCase       :    Optional[int]		      =      True

          UpperCamelCase       :    Tuple		      =      self.model.to(__lowerCAmelCase			)

          model.eval()
          model.float()

          UpperCamelCase       :    int		      =      model.module if hasattr(__lowerCAmelCase,     'module'			) else model
          quant_trainer.configure_model(__lowerCAmelCase,     self.quant_trainer_args			)

          UpperCamelCase       :    str		      =      os.path.join(__lowerCAmelCase,     'model.onnx'			)
          logger.info(F"""exporting model to {output_model_file}"""			)

          UpperCamelCase       :    int		      =      {0: 'batch_size', 1: 'seq_len'}

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




'''simple docstring'''




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


__lowercase        =    '''Create a default config file for Accelerate with only a few flags set.'''






def 							snake_case__				(     _A:   Dict="no" ,  _A:   str = default_json_config_file ,  _A:   bool = False					)       ->      List[str]:

		'''simple docstring'''


		lowerCAmelCase     =    Path(_A					)
		path.parent.mkdir(parents=_A ,  exist_ok=_A					)
		if path.exists():
				print(
				    f"Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."					)
				return False
		lowerCAmelCase     =    mixed_precision.lower()
		if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
				raise ValueError(
				    f"`mixed_precision` should be one of 'no', 'fp16', 'bf16', or 'fp8'. Received {mixed_precision}"					)
		lowerCAmelCase     =    {
		    """compute_environment""": """LOCAL_MACHINE""",
		    """mixed_precision""": mixed_precision,
		}
		if torch.cuda.is_available():
				lowerCAmelCase     =    torch.cuda.device_count()
				lowerCAmelCase     =    num_gpus
				lowerCAmelCase     =    False
				if num_gpus > 1:
						lowerCAmelCase     =    """MULTI_GPU"""
				else:
						lowerCAmelCase     =    """NO"""
		elif is_xpu_available() and use_xpu:
				lowerCAmelCase     =    torch.xpu.device_count()
				lowerCAmelCase     =    num_xpus
				lowerCAmelCase     =    False
				if num_xpus > 1:
						lowerCAmelCase     =    """MULTI_XPU"""
				else:
						lowerCAmelCase     =    """NO"""
		elif is_npu_available():
				lowerCAmelCase     =    torch.npu.device_count()
				lowerCAmelCase     =    num_npus
				lowerCAmelCase     =    False
				if num_npus > 1:
						lowerCAmelCase     =    """MULTI_NPU"""
				else:
						lowerCAmelCase     =    """NO"""
		else:
				lowerCAmelCase     =    0
				lowerCAmelCase     =    True
				lowerCAmelCase     =    1
				lowerCAmelCase     =    """NO"""
		lowerCAmelCase     =    ClusterConfig(**_A					)
		config.to_json_file(_A					)
		return path






def 							snake_case__				(     _A:   str ,  _A:   Any					)       ->      Optional[Any]:

		'''simple docstring'''


		lowerCAmelCase     =    parser.add_parser("""default""" ,  parents=_A ,  help=_A ,  formatter_class=_A					)
		parser.add_argument(
		    """--config_file""" ,  default=_A ,  help=(
		        """The path to use to store the config file. Will default to a file named default_config.yaml in the cache """
		        """location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have """
		        """such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed """
		        """with 'huggingface'."""
		    ) ,  dest="""save_location""" ,  )

		parser.add_argument(
		    """--mixed_precision""" ,  choices=["""no""", """fp16""", """bf16"""] ,  type=_A ,  help="""Whether or not to use mixed precision training. """
		    """Choose between FP16 and BF16 (bfloat16) training. """
		    """BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.""" ,  default="""no""" ,  )
		parser.set_defaults(func=_A					)
		return parser






def 							snake_case__				(     _A:   Optional[int]					)       ->      Optional[int]:

		'''simple docstring'''


		lowerCAmelCase     =    write_basic_config(args.mixed_precision ,  args.save_location					)
		if config_file:
				print(f"accelerate configuration saved at {config_file}"					)




