




"""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				:     Any					=				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__							):






				'''simple docstring'''





				def __init__(			self						,					*_a						,					_a=None						,					_a=None						,					_a=None						,					**_a					):





							"""simple docstring"""
							super().__init__(*_a						,					**_a					)
							lowerCamelCase      =   eval_examples
							lowerCamelCase      =   post_process_function
							lowerCamelCase      =   quant_trainer_args
							lowerCamelCase      =   128  # default number of calibration samples





				def   _lowerCAmelCase							(			self						,					_a=None					):





							"""simple docstring"""
							if calib_dataset is None and self.calib_dataset is None:
										raise ValueError("""Trainer: calibration requires an calib_dataset."""					)
							lowerCamelCase      =   calib_dataset if calib_dataset is not None else self.calib_dataset

							lowerCamelCase      =   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					):





							"""simple docstring"""
							lowerCamelCase      =   self.train_dataset if calib_dataset is None else calib_dataset
							lowerCamelCase      =   self.get_calib_dataloader(_a					)

							lowerCamelCase      =   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      =   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      =   model





				def   _lowerCAmelCase							(			self						,					_a=None						,					_a=None						,					_a=None						,					_a = "eval"					):





							"""simple docstring"""
							lowerCamelCase      =   self.eval_dataset if eval_dataset is None else eval_dataset
							lowerCamelCase      =   self.get_eval_dataloader(_a					)
							lowerCamelCase      =   self.eval_examples if eval_examples is None else eval_examples

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

							if self.post_process_function is not None and self.compute_metrics is not None:
										lowerCamelCase      =   self.post_process_function(_a						,					_a						,					output.predictions					)
										lowerCamelCase      =   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      =   metrics.pop(_a					)

										self.log(_a					)
							else:
										lowerCamelCase      =   {}

							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      =   self.callback_handler.on_evaluate(self.args						,					self.state						,					self.control						,					_a					)
							return metrics





				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a=None						,					_a = "test"					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_test_dataloader(_a					)

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

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

							lowerCamelCase      =   self.post_process_function(_a						,					_a						,					output.predictions						,					"""predict"""					)
							lowerCamelCase      =   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      =   metrics.pop(_a					)

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





				def   _lowerCAmelCase							(			self						,					_a="./"					):





							"""simple docstring"""
							lowerCamelCase      =   self.eval_dataset
							lowerCamelCase      =   self.get_eval_dataloader(_a					)

							lowerCamelCase      =   next(iter(_a					)					)

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

							# convert to tuple
							lowerCamelCase      =   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      =   True

							lowerCamelCase      =   self.model.to(_a					)

							model.eval()
							model.float()

							lowerCamelCase      =   model.module if hasattr(_a						,					"""module"""					) else model
							quant_trainer.configure_model(_a						,					self.quant_trainer_args					)

							lowerCamelCase      =   os.path.join(_a						,					"""model.onnx"""					)
							logger.info(f'exporting model to {output_model_file}'					)

							lowerCamelCase      =   {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"""




def        a__					(							snake_case__ = 50_00_00_00							)		->   int:
			lowerCamelCase      =   set()
			lowerCamelCase      =   int((limit - 24) ** (1 / 2)							)

			lowerCamelCase      =   set(range(3			,					prime_square_limit + 1			,					2							)							)
			primes.add(2							)
			for p in range(3			,					prime_square_limit + 1			,					2							):
						if p not in primes:
									continue
						primes.difference_update(set(range(p * p			,					prime_square_limit + 1			,					snake_case__							)							)							)

			for primea in primes:
						lowerCamelCase      =   primea * primea
						for primea in primes:
									lowerCamelCase      =   primea * primea * primea
									if square + cube >= limit - 16:
												break
									for primea in primes:
												lowerCamelCase      =   primea * primea * primea * primea
												lowerCamelCase      =   square + cube + tetr
												if total >= limit:
															break
												ret.add(snake_case__							)

			return len(snake_case__							)


if __name__ == "__main__":
					print(F"""{solution() = }""")



