





'''simple docstring'''






import bza
import gzip
import lzma
import os
import shutil
import struct
import tarfile
import warnings
import zipfile
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Dict, List, Optional, Type, Union

from .. import config
from .filelock import FileLock
from .logging import get_logger


lowerCAmelCase									=   get_logger(__name__)
class 		lowerCamelCase	:



	def __init__( self		,   a_ = None						):
			lowerCAmelCase  :     Any       						= (
			    os.path.join(a_		,   config.EXTRACTED_DATASETS_DIR						) if cache_dir else config.EXTRACTED_DATASETS_PATH
			)
			lowerCAmelCase  :     int       						= Extractor



	def 	_lowerCamelCase     ( self		,   a_						):
			from .file_utils import hash_url_to_filename

			# Path where we extract compressed archives
			# We extract in the cache dir, and get the extracted path name by hashing the original path"
			lowerCAmelCase  :     List[Any]       						= os.path.abspath(a_						)
			return os.path.join(self.extract_dir		,   hash_url_to_filename(a_						)						)



	def 	_lowerCamelCase     ( self		,   a_		,   a_						):
			return force_extract or (
			    not os.path.isfile(a_						) and not (os.path.isdir(a_						) and os.listdir(a_						))
			)





	def 	_lowerCamelCase     ( self		,   a_		,   a_ = False						):
			lowerCAmelCase  :     str       						= self.extractor.infer_extractor_format(a_						)
			if not extractor_format:
					return input_path
			lowerCAmelCase  :     Dict       						= self._get_output_path(a_						)
			if self._do_extract(a_		,   a_						):
					self.extractor.extract(a_		,   a_		,   a_						)
			return output_path






class 		lowerCamelCase	(   _A							):



	@classmethod
	@abstractmethod
	def 	_lowerCamelCase     ( cls		,   a_		,   **a_						):
			...





	@staticmethod
	@abstractmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			...






class 		lowerCamelCase	(   _A   , _A							):
	snake_case_   = []



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			with open(a_		,   "rb"						) as f:
					return f.read(a_						)





	@classmethod
	def 	_lowerCamelCase     ( cls		,   a_		,   a_ = b""						):
			if not magic_number:
					lowerCAmelCase  :     Optional[Any]       						= max(len(a_						) for cls_magic_number in cls.magic_numbers						)
					try:
							lowerCAmelCase  :     Tuple       						= cls.read_magic_number(a_		,   a_						)
					except OSError:
							return False
			return any(magic_number.startswith(a_						) for cls_magic_number in cls.magic_numbers						)






class 		lowerCamelCase	(   _A							):



	@classmethod
	def 	_lowerCamelCase     ( cls		,   a_		,   **a_						):
			return tarfile.is_tarfile(a_						)



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):

			def resolved(a_						) -> str:
					return os.path.realpath(os.path.abspath(a_						)						)

			def badpath(a_		,   a_						) -> bool:
					# joinpath will ignore base if path is absolute
					return not resolved(os.path.join(a_		,   a_						)						).startswith(a_						)

			def badlink(a_		,   a_						) -> bool:
					# Links are interpreted relative to the directory containing the link
					lowerCAmelCase  :     Optional[int]       						= resolved(os.path.join(a_		,   os.path.dirname(info.name						)						)						)
					return badpath(info.linkname		,   base=a_						)

			lowerCAmelCase  :     Optional[int]       						= resolved(a_						)

			for finfo in members:
					if badpath(finfo.name		,   a_						):
							logger.error(F'''Extraction of {finfo.name} is blocked (illegal path)'''						)
					elif finfo.issym() and badlink(a_		,   a_						):
							logger.error(F'''Extraction of {finfo.name} is blocked: Symlink to {finfo.linkname}'''						)
					elif finfo.islnk() and badlink(a_		,   a_						):
							logger.error(F'''Extraction of {finfo.name} is blocked: Hard link to {finfo.linkname}'''						)
					else:
							yield finfo





	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			os.makedirs(a_		,   exist_ok=a_						)
			lowerCAmelCase  :     Tuple       						= tarfile.open(a_						)
			tar_file.extractall(a_		,   members=TarExtractor.safemembers(a_		,   a_						)						)
			tar_file.close()






class 		lowerCamelCase	(   _A							):
	snake_case_   = [b"\x1F\x8B"]



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			with gzip.open(a_		,   "rb"						) as gzip_file:
					with open(a_		,   "wb"						) as extracted_file:
							shutil.copyfileobj(a_		,   a_						)






class 		lowerCamelCase	(   _A							):
	snake_case_   = [
	    b"PK\x03\x04",
	    b"PK\x05\x06",  # empty archive
	    b"PK\x07\x08",  # spanned archive
	]



	@classmethod
	def 	_lowerCamelCase     ( cls		,   a_		,   a_ = b""						):
			if super().is_extractable(a_		,   magic_number=a_						):
					return True
			try:
					# Alternative version of zipfile.is_zipfile that has less false positives, but misses executable zip archives.
					# From: https://github.com/python/cpython/pull/5053
					from zipfile import (
					    _CD_SIGNATURE,
					    _ECD_DISK_NUMBER,
					    _ECD_DISK_START,
					    _ECD_ENTRIES_TOTAL,
					    _ECD_OFFSET,
					    _ECD_SIZE,
					    _EndRecData,
					    sizeCentralDir,
					    stringCentralDir,
					    structCentralDir,
					)

					with open(a_		,   "rb"						) as fp:
							lowerCAmelCase  :     Optional[Any]       						= _EndRecData(a_						)
							if endrec:
									if endrec[_ECD_ENTRIES_TOTAL] == 0 and endrec[_ECD_SIZE] == 0 and endrec[_ECD_OFFSET] == 0:
											return True  # Empty zipfiles are still zipfiles
									elif endrec[_ECD_DISK_NUMBER] == endrec[_ECD_DISK_START]:
											fp.seek(endrec[_ECD_OFFSET]						)  # Central directory is on the same disk
											if fp.tell() == endrec[_ECD_OFFSET] and endrec[_ECD_SIZE] >= sizeCentralDir:
													lowerCAmelCase  :     Any       						= fp.read(a_						)  # CD is where we expect it to be
													if len(a_						) == sizeCentralDir:
															lowerCAmelCase  :     int       						= struct.unpack(a_		,   a_						)  # CD is the right size
															if centdir[_CD_SIGNATURE] == stringCentralDir:
																	return True  # First central directory entry  has correct magic number
					return False
			except Exception:  # catch all errors in case future python versions change the zipfile internals
					return False





	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			os.makedirs(a_		,   exist_ok=a_						)
			with zipfile.ZipFile(a_		,   "r"						) as zip_file:
					zip_file.extractall(a_						)
					zip_file.close()






class 		lowerCamelCase	(   _A							):
	snake_case_   = [b"\xFD\x37\x7A\x58\x5A\x00"]



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			with lzma.open(a_						) as compressed_file:
					with open(a_		,   "wb"						) as extracted_file:
							shutil.copyfileobj(a_		,   a_						)






class 		lowerCamelCase	(   _A							):
	snake_case_   = [b"Rar!\x1a\x07\x00", b"Rar!\x1a\x07\x01\x00"]  # RAR_ID  # RAR5_ID



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			if not config.RARFILE_AVAILABLE:
					raise ImportError("Please pip install rarfile"						)
			import rarfile

			os.makedirs(a_		,   exist_ok=a_						)
			lowerCAmelCase  :     Optional[Any]       						= rarfile.RarFile(a_						)
			rf.extractall(a_						)
			rf.close()






class 		lowerCamelCase	(   _A							):
	snake_case_   = [b"\x28\xb5\x2F\xFD"]



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			if not config.ZSTANDARD_AVAILABLE:
					raise ImportError("Please pip install zstandard"						)
			import zstandard as zstd

			lowerCAmelCase  :     List[Any]       						= zstd.ZstdDecompressor()
			with open(a_		,   "rb"						) as ifh, open(a_		,   "wb"						) as ofh:
					dctx.copy_stream(a_		,   a_						)






class 		lowerCamelCase	(   _A							):
	snake_case_   = [b"\x42\x5A\x68"]



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			with bza.open(a_		,   "rb"						) as compressed_file:
					with open(a_		,   "wb"						) as extracted_file:
							shutil.copyfileobj(a_		,   a_						)






class 		lowerCamelCase	(   _A							):
	snake_case_   = [b"\x37\x7A\xBC\xAF\x27\x1C"]



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			if not config.PY7ZR_AVAILABLE:
					raise ImportError("Please pip install py7zr"						)
			import pyazr

			os.makedirs(a_		,   exist_ok=a_						)
			with pyazr.SevenZipFile(a_		,   "r"						) as archive:
					archive.extractall(a_						)






class 		lowerCamelCase	(   _A							):
	snake_case_   = [b"\x04\x22\x4D\x18"]



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			if not config.LZ4_AVAILABLE:
					raise ImportError("Please pip install lz4"						)
			import lza.frame

			with lza.frame.open(a_		,   "rb"						) as compressed_file:
					with open(a_		,   "wb"						) as extracted_file:
							shutil.copyfileobj(a_		,   a_						)






class 		lowerCamelCase	:
	#  Put zip file to the last, b/c it is possible wrongly detected as zip (I guess it means: as tar or gzip)
	snake_case_   = {
	    "tar": TarExtractor,
	    "gzip": GzipExtractor,
	    "zip": ZipExtractor,
	    "xz": XzExtractor,
	    "rar": RarExtractor,
	    "zstd": ZstdExtractor,
	    "bz2": BzipaExtractor,
	    "7z": SevenZipExtractor,  # <Added version="2.4.0"/>
	    "lz4": LzaExtractor,  # <Added version="2.4.0"/>
	}



	@classmethod
	def 	_lowerCamelCase     ( cls						):
			return max(
			    len(a_						)
			    for extractor in cls.extractors.values()
			    if issubclass(a_		,   a_						)
			    for extractor_magic_number in extractor.magic_numbers						)



	@staticmethod
	def 	_lowerCamelCase     ( a_		,   a_						):
			try:
					return MagicNumberBaseExtractor.read_magic_number(a_		,   magic_number_length=a_						)
			except OSError:
					return b""



	@classmethod
	def 	_lowerCamelCase     ( cls		,   a_		,   a_ = False						):
			warnings.warn(
			    "Method 'is_extractable' was deprecated in version 2.4.0 and will be removed in 3.0.0. "
			    "Use 'infer_extractor_format' instead."		,   category=a_		,   )
			lowerCAmelCase  :     Any       						= cls.infer_extractor_format(a_						)
			if extractor_format:
					return True if not return_extractor else (True, cls.extractors[extractor_format])
			return False if not return_extractor else (False, None)



	@classmethod
	def 	_lowerCamelCase     ( cls		,   a_						):  # <Added version="2.4.0"/>
			lowerCAmelCase  :     List[str]       						= cls._get_magic_number_max_length()
			lowerCAmelCase  :     int       						= cls._read_magic_number(a_		,   a_						)
			for extractor_format, extractor in cls.extractors.items():
					if extractor.is_extractable(a_		,   magic_number=a_						):
							return extractor_format





	@classmethod
	def 	_lowerCamelCase     ( cls		,   a_		,   a_		,   a_ = None		,   a_ = "deprecated"		,   ):
			os.makedirs(os.path.dirname(a_						)		,   exist_ok=a_						)
			# Prevent parallel extractions
			lowerCAmelCase  :     Union[str, Any]       						= str(Path(a_						).with_suffix(".lock"						)						)
			with FileLock(a_						):
					shutil.rmtree(a_		,   ignore_errors=a_						)
					if extractor_format or extractor != "deprecated":
							if extractor != "deprecated" or not isinstance(a_		,   a_						):  # passed as positional arg
									warnings.warn(
									    "Parameter 'extractor' was deprecated in version 2.4.0 and will be removed in 3.0.0. "
									    "Use 'extractor_format' instead."		,   category=a_		,   )
									lowerCAmelCase  :     int       						= extractor if extractor != "deprecated" else extractor_format
							else:
									lowerCAmelCase  :     int       						= cls.extractors[extractor_format]
							return extractor.extract(a_		,   a_						)
					else:
							warnings.warn(
							    "Parameter 'extractor_format' was made required in version 2.4.0 and not passing it will raise an "
							    "exception in 3.0.0."		,   category=a_		,   )
							for extractor in cls.extractors.values():
									if extractor.is_extractable(a_						):
											return extractor.extract(a_		,   a_						)










'''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 		lowerCamelCase	(   _A							):



	def __init__( self		,   *a_		,   a_=None		,   a_=None		,   a_=None		,   **a_						):
			super().__init__(*a_		,   **a_						)
			lowerCAmelCase  :     Tuple       						= eval_examples
			lowerCAmelCase  :     List[Any]       						= post_process_function
			lowerCAmelCase  :     List[Any]       						= quant_trainer_args
			lowerCAmelCase  :     Dict       						= 128  # default number of calibration samples



	def 	_lowerCamelCase     ( self		,   a_=None						):
			if calib_dataset is None and self.calib_dataset is None:
					raise ValueError("Trainer: calibration requires an calib_dataset."						)
			lowerCAmelCase  :     List[str]       						= calib_dataset if calib_dataset is not None else self.calib_dataset

			lowerCAmelCase  :     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 	_lowerCamelCase     ( self		,   a_=None						):
			lowerCAmelCase  :     str       						= self.train_dataset if calib_dataset is None else calib_dataset
			lowerCAmelCase  :     str       						= self.get_calib_dataloader(a_						)

			lowerCAmelCase  :     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
					lowerCAmelCase						,		lowerCAmelCase						,		lowerCAmelCase  :     Tuple       						= 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						)
			lowerCAmelCase  :     Union[str, Any]       						= model



	def 	_lowerCamelCase     ( self		,   a_=None		,   a_=None		,   a_=None		,   a_ = "eval"						):
			lowerCAmelCase  :     Optional[Any]       						= self.eval_dataset if eval_dataset is None else eval_dataset
			lowerCAmelCase  :     Optional[Any]       						= self.get_eval_dataloader(a_						)
			lowerCAmelCase  :     Union[str, Any]       						= self.eval_examples if eval_examples is None else eval_examples

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

			if self.post_process_function is not None and self.compute_metrics is not None:
					lowerCAmelCase  :     Optional[Any]       						= self.post_process_function(a_		,   a_		,   output.predictions						)
					lowerCAmelCase  :     Optional[Any]       						= 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}_'''						):
									lowerCAmelCase  :     Any       						= metrics.pop(a_						)

					self.log(a_						)
			else:
					lowerCAmelCase  :     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()						)

			lowerCAmelCase  :     Dict       						= self.callback_handler.on_evaluate(self.args		,   self.state		,   self.control		,   a_						)
			return metrics



	def 	_lowerCamelCase     ( self		,   a_		,   a_		,   a_=None		,   a_ = "test"						):
			lowerCAmelCase  :     int       						= self.get_test_dataloader(a_						)

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

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

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

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





	def 	_lowerCamelCase     ( self		,   a_="./"						):
			lowerCAmelCase  :     Optional[Any]       						= self.eval_dataset
			lowerCAmelCase  :     Optional[int]       						= self.get_eval_dataloader(a_						)

			lowerCAmelCase  :     Union[str, Any]       						= next(iter(a_						)						)

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

			# convert to tuple
			lowerCAmelCase  :     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

			lowerCAmelCase  :     Any       						= True

			lowerCAmelCase  :     Dict       						= self.model.to(a_						)

			model.eval()
			model.float()

			lowerCAmelCase  :     Union[str, Any]       						= model.module if hasattr(a_		,   "module"						) else model
			quant_trainer.configure_model(a_		,   self.quant_trainer_args						)

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

			lowerCAmelCase  :     List[Any]       						= {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"						)



