


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 			_a				(     UpperCamelCase__     ):

							def __init__(				self:      Any       ,		*UpperCamelCase_:      Tuple       ,		UpperCamelCase_:      Any=None       ,		UpperCamelCase_:      Union[str, Any]=None       ,		UpperCamelCase_:      int=None       ,		**UpperCamelCase_:      Union[str, Any]				)     ->  Optional[Any]:





									"""simple docstring"""







									super().__init__(*UpperCamelCase_       ,		**UpperCamelCase_				)
									lowercase__				= eval_examples
									lowercase__				= post_process_function
									lowercase__				= quant_trainer_args
									lowercase__				= 128  # default number of calibration samples

							def  lowerCamelCase_					(				self:      Any       ,		UpperCamelCase_:      int=None				)     ->  Optional[Any]:





									"""simple docstring"""







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

									lowercase__				= 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  lowerCamelCase_					(				self:      Tuple       ,		UpperCamelCase_:      List[Any]=None				)     ->  Optional[Any]:





									"""simple docstring"""







									lowercase__				= self.train_dataset if calib_dataset is None else calib_dataset
									lowercase__				= self.get_calib_dataloader(UpperCamelCase_				)

									lowercase__				= 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
											lowercase__	,						lowercase__	,						lowercase__				= 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				)
									lowercase__				= model

							def  lowerCamelCase_					(				self:      Union[str, Any]       ,		UpperCamelCase_:      Optional[Any]=None       ,		UpperCamelCase_:      str=None       ,		UpperCamelCase_:      str=None       ,		UpperCamelCase_:      str = "eval"				)     ->  Union[str, Any]:





									"""simple docstring"""







									lowercase__				= self.eval_dataset if eval_dataset is None else eval_dataset
									lowercase__				= self.get_eval_dataloader(UpperCamelCase_				)
									lowercase__				= self.eval_examples if eval_examples is None else eval_examples

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

									if self.post_process_function is not None and self.compute_metrics is not None:
											lowercase__				= self.post_process_function(UpperCamelCase_       ,		UpperCamelCase_       ,		output.predictions				)
											lowercase__				= 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}_'				):
															lowercase__				= metrics.pop(UpperCamelCase_				)

											self.log(UpperCamelCase_				)
									else:
											lowercase__				= {}

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

									lowercase__				= self.callback_handler.on_evaluate(self.args       ,		self.state       ,		self.control       ,		UpperCamelCase_				)
									return metrics

							def  lowerCamelCase_					(				self:      Dict       ,		UpperCamelCase_:      Optional[int]       ,		UpperCamelCase_:      int       ,		UpperCamelCase_:      Union[str, Any]=None       ,		UpperCamelCase_:      str = "test"				)     ->  Tuple:





									"""simple docstring"""







									lowercase__				= self.get_test_dataloader(UpperCamelCase_				)

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

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

									lowercase__				= self.post_process_function(UpperCamelCase_       ,		UpperCamelCase_       ,		output.predictions       ,		'''predict'''				)
									lowercase__				= 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}_'				):
													lowercase__				= metrics.pop(UpperCamelCase_				)

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


							def  lowerCamelCase_					(				self:      Tuple       ,		UpperCamelCase_:      Union[str, Any]="./"				)     ->  Union[str, Any]:





									"""simple docstring"""







									lowercase__				= self.eval_dataset
									lowercase__				= self.get_eval_dataloader(UpperCamelCase_				)

									lowercase__				= next(iter(UpperCamelCase_				)				)

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

									# convert to tuple
									lowercase__				= tuple(v.to(UpperCamelCase_				) for k, v in batch.items()				)

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

									lowercase__				= True

									lowercase__				= self.model.to(UpperCamelCase_				)

									model.eval()
									model.float()

									lowercase__				= model.module if hasattr(UpperCamelCase_       ,		'''module'''				) else model
									quant_trainer.configure_model(UpperCamelCase_       ,		self.quant_trainer_args				)

									lowercase__				= os.path.join(UpperCamelCase_       ,		'''model.onnx'''				)
									logger.info(f'exporting model to {output_model_file}'				)

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




from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map



def        a__					(							snake_case__							)		->   Dict[str, torch.Tensor]:
			lowerCamelCase      =   []
			lowerCamelCase      =   []
			lowerCamelCase      =   []

			for rt in rc.restypes:
						lowerCamelCase      =   rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
						restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]							)
						lowerCamelCase      =   {name: i for i, name in enumerate(snake_case__							)}
						restype_atomaa_to_atomaa_list.append(
						    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]							)

						restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]							)

			# Add dummy mapping for restype 'UNK'
			restype_atomaa_to_atomaa_list.append([0] * 14							)
			restype_atomaa_to_atomaa_list.append([0] * 37							)
			restype_atomaa_mask_list.append([0.0] * 14							)

			lowerCamelCase      =   torch.tensor(
			    snake_case__			,					dtype=torch.intaa			,					device=protein["""aatype"""].device			,					)
			lowerCamelCase      =   torch.tensor(
			    snake_case__			,					dtype=torch.intaa			,					device=protein["""aatype"""].device			,					)
			lowerCamelCase      =   torch.tensor(
			    snake_case__			,					dtype=torch.floataa			,					device=protein["""aatype"""].device			,					)
			lowerCamelCase      =   protein["""aatype"""].to(torch.long							)

			# create the mapping for (residx, atom14) --> atom37, i.e. an array
			# with shape (num_res, 14) containing the atom37 indices for this protein
			lowerCamelCase      =   restype_atomaa_to_atomaa[protein_aatype]
			lowerCamelCase      =   restype_atomaa_mask[protein_aatype]

			lowerCamelCase      =   residx_atomaa_mask
			lowerCamelCase      =   residx_atomaa_to_atomaa.long()

			# create the gather indices for mapping back
			lowerCamelCase      =   restype_atomaa_to_atomaa[protein_aatype]
			lowerCamelCase      =   residx_atomaa_to_atomaa.long()

			# create the corresponding mask
			lowerCamelCase      =   torch.zeros([21, 37]			,					dtype=torch.floataa			,					device=protein["""aatype"""].device							)
			for restype, restype_letter in enumerate(rc.restypes							):
						lowerCamelCase      =   rc.restype_atoa[restype_letter]
						lowerCamelCase      =   rc.residue_atoms[restype_name]
						for atom_name in atom_names:
									lowerCamelCase      =   rc.atom_order[atom_name]
									lowerCamelCase      =   1

			lowerCamelCase      =   restype_atomaa_mask[protein_aatype]
			lowerCamelCase      =   residx_atomaa_mask

			return protein



def        a__					(							snake_case__							)		->   Dict[str, np.ndarray]:
			lowerCamelCase      =   tree_map(lambda snake_case__							: torch.tensor(snake_case__			,					device=batch["""aatype"""].device							)			,					snake_case__			,					np.ndarray							)
			lowerCamelCase      =   tensor_tree_map(lambda snake_case__							: np.array(snake_case__							)			,					make_atomaa_masks(snake_case__							)							)
			return out



