



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


__lowerCAmelCase		     =  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       __magic_name__		(       _UpperCamelCase  ):







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







							def 				__lowercase     (				self   :			List[Any]      ,_UpperCAmelCase   :			List[Any]=None						):
									if calib_dataset is None and self.calib_dataset is None:
											raise ValueError('Trainer: calibration requires an calib_dataset.'						)
									_a	: List[Any]								=				calib_dataset if calib_dataset is not None else self.calib_dataset

									_a	: Dict								=				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   :			str      ,_UpperCAmelCase   :			int=None						):
									_a	: Optional[Any]								=				self.train_dataset if calib_dataset is None else calib_dataset
									_a	: Tuple								=				self.get_calib_dataloader(_UpperCAmelCase						)

									_a	: Optional[int]								=				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
											_a			,   _a			,   _a	: List[Any]								=				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						)
									_a	: Union[str, Any]								=				model







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

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

									if self.post_process_function is not None and self.compute_metrics is not None:
											_a	: Tuple								=				self.post_process_function(_UpperCAmelCase      ,_UpperCAmelCase      ,output.predictions						)
											_a	: Optional[Any]								=				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}_"""						):
															_a	: List[str]								=				metrics.pop(_UpperCAmelCase						)

											self.log(_UpperCAmelCase						)
									else:
											_a	: Union[str, Any]								=				{}

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

									_a	: List[str]								=				self.callback_handler.on_evaluate(self.args      ,self.state      ,self.control      ,_UpperCAmelCase						)
									return metrics







							def 				__lowercase     (				self   :			Any      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			str      ,_UpperCAmelCase   :			int=None      ,_UpperCAmelCase   :			str = "test"						):
									_a	: Any								=				self.get_test_dataloader(_UpperCAmelCase						)

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

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

									_a	: str								=				self.post_process_function(_UpperCAmelCase      ,_UpperCAmelCase      ,output.predictions      ,'predict'						)
									_a	: Optional[Any]								=				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}_"""						):
													_a	: List[str]								=				metrics.pop(_UpperCAmelCase						)

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







							def 				__lowercase     (				self   :			int      ,_UpperCAmelCase   :			Dict="./"						):
									_a	: Optional[Any]								=				self.eval_dataset
									_a	: List[Any]								=				self.get_eval_dataloader(_UpperCAmelCase						)

									_a	: Optional[Any]								=				next(iter(_UpperCAmelCase						)						)

									# saving device - to make it consistent
									_a	: str								=				torch.device('cuda' if torch.cuda.is_available() else 'cpu'						)

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

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

									_a	: Tuple								=				True

									_a	: Union[str, Any]								=				self.model.to(_UpperCAmelCase						)

									model.eval()
									model.float()

									_a	: Union[str, Any]								=				model.module if hasattr(_UpperCAmelCase      ,'module'						) else model
									quant_trainer.configure_model(_UpperCAmelCase      ,self.quant_trainer_args						)

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

									_a	: Any								=				{0: 'batch_size', 1: 'seq_len'}

									torch.onnx.export(
									    _UpperCAmelCase      ,_UpperCAmelCase      ,_UpperCAmelCase      ,export_params=_UpperCAmelCase      ,opset_version=13      ,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'''

# Lint as: python3
import os
import re
import urllib.parse
from pathlib import Path
from typing import Callable, List, Optional, Union
from zipfile import ZipFile

from ..utils.file_utils import cached_path, hf_github_url
from ..utils.logging import get_logger
from ..utils.version import Version


__lowerCAmelCase		     =  get_logger(__name__)

class       __magic_name__		:
							lowerCAmelCase			:       Optional[Any]  							=   'dummy_data'
							lowerCAmelCase			:       List[str]  							=   'datasets'
							lowerCAmelCase			:       str  							=   False







							def __init__(				self   :			Tuple      ,_UpperCAmelCase   :			str      ,_UpperCAmelCase   :			str      ,_UpperCAmelCase   :			Union[Version, str]      ,_UpperCAmelCase   :			Optional[str] = None      ,_UpperCAmelCase   :			bool = False      ,_UpperCAmelCase   :			bool = True      ,_UpperCAmelCase   :			Optional[List[Callable]] = None      ,):
									_a	: int								=				0
									_a	: str								=				dataset_name
									_a	: Optional[int]								=				cache_dir
									_a	: int								=				use_local_dummy_data
									_a	: Optional[int]								=				config
									# download_callbacks take a single url as input
									_a	: List[Callable]								=				download_callbacks or []
									# if False, it doesn't load existing files and it returns the paths of the dummy files relative
									# to the dummy_data zip file root
									_a	: str								=				load_existing_dummy_data

									# TODO(PVP, QL) might need to make this more general
									_a	: Optional[int]								=				str(_UpperCAmelCase						)
									# to be downloaded
									_a	: Union[str, Any]								=				None
									_a	: Dict								=				None







							@property
							def 				__lowercase     (				self   :			Any						):
									if self._dummy_file is None:
											_a	: Dict								=				self.download_dummy_data()
									return self._dummy_file







							@property
							def 				__lowercase     (				self   :			Optional[int]						):
									if self.config is not None:
											# structure is dummy / config_name / version_name
											return os.path.join('dummy'      ,self.config.name      ,self.version_name						)
									# structure is dummy / version_name
									return os.path.join('dummy'      ,self.version_name						)







							@property
							def 				__lowercase     (				self   :			Any						):
									return os.path.join(self.dummy_data_folder      ,'dummy_data.zip'						)







							def 				__lowercase     (				self   :			List[str]						):
									_a	: Dict								=				(
									    self.local_path_to_dummy_data if self.use_local_dummy_data is True else self.github_path_to_dummy_data
									)

									_a	: Dict								=				cached_path(
									    _UpperCAmelCase      ,cache_dir=self.cache_dir      ,extract_compressed_file=_UpperCAmelCase      ,force_extract=_UpperCAmelCase						)

									return os.path.join(_UpperCAmelCase      ,self.dummy_file_name						)







							@property
							def 				__lowercase     (				self   :			List[str]						):
									return os.path.join(self.datasets_scripts_dir      ,self.dataset_name      ,self.dummy_zip_file						)







							@property
							def 				__lowercase     (				self   :			Optional[Any]						):
									if self._bucket_url is None:
											_a	: Any								=				hf_github_url(self.dataset_name      ,self.dummy_zip_file.replace(os.sep      ,'/'						)						)
									return self._bucket_url







							@property
							def 				__lowercase     (				self   :			int						):
									# return full path if its a dir
									if os.path.isdir(self.dummy_file						):
											return self.dummy_file
									# else cut off path to file -> example `xsum`.
									return "/".join(self.dummy_file.replace(os.sep      ,'/'						).split('/'						)[:-1]						)







							def 				__lowercase     (				self   :			Dict      ,_UpperCAmelCase   :			str      ,*_UpperCAmelCase   :			int						):
									if self.load_existing_dummy_data:
											# dummy data is downloaded and tested
											_a	: List[Any]								=				self.dummy_file
									else:
											# dummy data cannot be downloaded and only the path to dummy file is returned
											_a	: str								=				self.dummy_file_name

									# special case when data_url is a dict
									if isinstance(_UpperCAmelCase      ,_UpperCAmelCase						):
											return self.create_dummy_data_dict(_UpperCAmelCase      ,_UpperCAmelCase						)
									elif isinstance(_UpperCAmelCase      ,(list, tuple)						):
											return self.create_dummy_data_list(_UpperCAmelCase      ,_UpperCAmelCase						)
									else:
											return self.create_dummy_data_single(_UpperCAmelCase      ,_UpperCAmelCase						)







							def 				__lowercase     (				self   :			List[str]      ,_UpperCAmelCase   :			Optional[int]      ,*_UpperCAmelCase   :			Optional[int]						):
									return self.download_and_extract(_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Tuple      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			Optional[Any]						):
									return self.download_and_extract(_UpperCAmelCase						)







							def 				__lowercase     (				self   :			List[str]      ,_UpperCAmelCase   :			List[str]      ,*_UpperCAmelCase   :			int      ,**_UpperCAmelCase   :			Dict						):
									return path







							def 				__lowercase     (				self   :			Optional[Any]						):
									return {}







							def 				__lowercase     (				self   :			Optional[Any]      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			str						):
									_a	: int								=				{}
									for key, single_urls in data_url.items():
											for download_callback in self.download_callbacks:
													if isinstance(_UpperCAmelCase      ,_UpperCAmelCase						):
															for single_url in single_urls:
																	download_callback(_UpperCAmelCase						)
													else:
															_a	: Union[str, Any]								=				single_urls
															download_callback(_UpperCAmelCase						)
            # we force the name of each key to be the last file / folder name of the url path
            # if the url has arguments, we need to encode them with urllib.parse.quote_plus
											if isinstance(_UpperCAmelCase      ,_UpperCAmelCase						):
													_a	: Optional[int]								=				[os.path.join(_UpperCAmelCase      ,urllib.parse.quote_plus(Path(_UpperCAmelCase						).name						)						) for x in single_urls]
											else:
													_a	: Tuple								=				single_urls
													_a	: Optional[Any]								=				os.path.join(_UpperCAmelCase      ,urllib.parse.quote_plus(Path(_UpperCAmelCase						).name						)						)
											_a	: Union[str, Any]								=				value

									# make sure that values are unique
									if all(isinstance(_UpperCAmelCase      ,_UpperCAmelCase						) for i in dummy_data_dict.values()						) and len(set(dummy_data_dict.values()						)						) < len(
									    dummy_data_dict.values()						):
											# append key to value to make its name unique
											_a	: Dict								=				{key: value + key for key, value in dummy_data_dict.items()}

									return dummy_data_dict







							def 				__lowercase     (				self   :			List[Any]      ,_UpperCAmelCase   :			List[str]      ,_UpperCAmelCase   :			List[str]						):
									_a	: str								=				[]
									# trick: if there are many shards named like `data.txt-000001-of-00300`, only use the first one
									_a	: Tuple								=				all(bool(re.findall('[0-9]{3,}-of-[0-9]{3,}'      ,_UpperCAmelCase						)						) for url in data_url						)
									_a	: str								=				all(
									    url.startswith('https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/pubmed'						) for url in data_url						)
									if data_url and (is_tf_records or is_pubmed_records):
											_a	: str								=				[data_url[0]] * len(_UpperCAmelCase						)
									for single_url in data_url:
											for download_callback in self.download_callbacks:
													download_callback(_UpperCAmelCase						)
											# we force the name of each key to be the last file / folder name of the url path
											# if the url has arguments, we need to encode them with urllib.parse.quote_plus
											_a	: List[Any]								=				os.path.join(_UpperCAmelCase      ,urllib.parse.quote_plus(single_url.split('/'						)[-1]						)						)
											dummy_data_list.append(_UpperCAmelCase						)
									return dummy_data_list







							def 				__lowercase     (				self   :			List[Any]      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			List[str]						):
									for download_callback in self.download_callbacks:
											download_callback(_UpperCAmelCase						)
									# we force the name of each key to be the last file / folder name of the url path
									# if the url has arguments, we need to encode them with urllib.parse.quote_plus
									_a	: List[Any]								=				os.path.join(_UpperCAmelCase      ,urllib.parse.quote_plus(data_url.split('/'						)[-1]						)						)
									if os.path.exists(_UpperCAmelCase						) or not self.load_existing_dummy_data:
											return value
									else:
											# Backward compatibility, maybe deprecate at one point.
											# For many datasets with single url calls to dl_manager.download_and_extract,
											# the dummy_data.zip file is actually the zipped downloaded file
											# while now we expected the dummy_data.zip file to be a directory containing
											# the downloaded file.
											return path_to_dummy_data







							def 				__lowercase     (				self   :			Tuple						):
									pass







							def 				__lowercase     (				self   :			Optional[int]						):
									pass







							def 				__lowercase     (				self   :			List[str]      ,_UpperCAmelCase   :			Dict						):
									def _iter_archive_members(_UpperCAmelCase   :			List[Any]						):
											# this preserves the order of the members inside the ZIP archive
											_a	: int								=				Path(self.dummy_file						).parent
											_a	: List[str]								=				path.relative_to(_UpperCAmelCase						)
											with ZipFile(self.local_path_to_dummy_data						) as zip_file:
													_a	: List[Any]								=				zip_file.namelist()
											for member in members:
													if member.startswith(relative_path.as_posix()						):
															yield dummy_parent_path.joinpath(_UpperCAmelCase						)

									_a	: Dict								=				Path(_UpperCAmelCase						)
									_a	: Dict								=				_iter_archive_members(_UpperCAmelCase						) if self.use_local_dummy_data else path.rglob('*'						)
									for file_path in file_paths:
											if file_path.is_file() and not file_path.name.startswith(('.', '__')						):
													yield file_path.relative_to(_UpperCAmelCase						).as_posix(), file_path.open('rb'						)







							def 				__lowercase     (				self   :			Any      ,_UpperCAmelCase   :			Union[str, Any]						):
									if not isinstance(_UpperCAmelCase      ,_UpperCAmelCase						):
											_a	: str								=				[paths]
									for path in paths:
											if os.path.isfile(_UpperCAmelCase						):
													if os.path.basename(_UpperCAmelCase						).startswith(('.', '__')						):
															return
													yield path
											else:
													for dirpath, dirnames, filenames in os.walk(_UpperCAmelCase						):
															if os.path.basename(_UpperCAmelCase						).startswith(('.', '__')						):
																	continue
															dirnames.sort()
															for filename in sorted(_UpperCAmelCase						):
																	if filename.startswith(('.', '__')						):
																			continue
																	yield os.path.join(_UpperCAmelCase      ,_UpperCAmelCase						)




