


import copy
import inspect
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers import VideoMAEConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch
	from torch import nn

	from transformers import (
	    MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING,
	    VideoMAEForPreTraining,
	    VideoMAEForVideoClassification,
	    VideoMAEModel,
	)
	from transformers.models.videomae.modeling_videomae import VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from transformers import VideoMAEImageProcessor





class  A       :


							def __init__(self ,							lowerCAmelCase ,							lowerCAmelCase=1_3 ,							lowerCAmelCase=1_0 ,							lowerCAmelCase=3 ,							lowerCAmelCase=2 ,							lowerCAmelCase=2 ,							lowerCAmelCase=2 ,							lowerCAmelCase=True ,							lowerCAmelCase=True ,							lowerCAmelCase=3_2 ,							lowerCAmelCase=5 ,							lowerCAmelCase=4 ,							lowerCAmelCase=3_7 ,							lowerCAmelCase="gelu" ,							lowerCAmelCase=0.1 ,							lowerCAmelCase=0.1 ,							lowerCAmelCase=1_0 ,							lowerCAmelCase=0.02 ,							lowerCAmelCase=0.9 ,							lowerCAmelCase=None ,							):
									__lowercase=     parent
									__lowercase=     batch_size
									__lowercase=     image_size
									__lowercase=     num_channels
									__lowercase=     patch_size
									__lowercase=     tubelet_size
									__lowercase=     num_frames
									__lowercase=     is_training
									__lowercase=     use_labels
									__lowercase=     hidden_size
									__lowercase=     num_hidden_layers
									__lowercase=     num_attention_heads
									__lowercase=     intermediate_size
									__lowercase=     hidden_act
									__lowercase=     hidden_dropout_prob
									__lowercase=     attention_probs_dropout_prob
									__lowercase=     type_sequence_label_size
									__lowercase=     initializer_range
									__lowercase=     mask_ratio
									__lowercase=     scope

									# in VideoMAE, the number of tokens equals num_frames/tubelet_size * num_patches per frame
									__lowercase=     (image_size // patch_size) ** 2
									__lowercase=     (num_frames // tubelet_size) * self.num_patches_per_frame

									# use this variable to define bool_masked_pos
									__lowercase=     int(mask_ratio * self.seq_length   )


							def      _A			(self   ):
									__lowercase=     floats_tensor(
									    [self.batch_size, self.num_frames, self.num_channels, self.image_size, self.image_size]   )

									__lowercase=     None
									if self.use_labels:
											__lowercase=     ids_tensor([self.batch_size] ,							self.type_sequence_label_size   )

									__lowercase=     self.get_config()

									return config, pixel_values, labels


							def      _A			(self   ):
									return VideoMAEConfig(
									    image_size=self.image_size ,							patch_size=self.patch_size ,							num_channels=self.num_channels ,							num_frames=self.num_frames ,							tubelet_size=self.tubelet_size ,							hidden_size=self.hidden_size ,							num_hidden_layers=self.num_hidden_layers ,							num_attention_heads=self.num_attention_heads ,							intermediate_size=self.intermediate_size ,							hidden_act=self.hidden_act ,							hidden_dropout_prob=self.hidden_dropout_prob ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob ,							is_decoder=__lowerCAmelCase ,							initializer_range=self.initializer_range ,							)


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase   ):
									__lowercase=     VideoMAEModel(config=__lowerCAmelCase   )
									model.to(__lowerCAmelCase   )
									model.eval()
									__lowercase=     model(__lowerCAmelCase   )
									self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.seq_length, self.hidden_size)   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase   ):
									__lowercase=     VideoMAEForPreTraining(__lowerCAmelCase   )
									model.to(__lowerCAmelCase   )
									model.eval()
									# important: each video needs to have the same number of masked patches
									# hence we define a single mask, which we then repeat for each example in the batch
									__lowercase=     torch.ones((self.num_masks,)   )
									__lowercase=     torch.cat([mask, torch.zeros(self.seq_length - mask.size(0   )   )]   )
									__lowercase=     mask.expand(self.batch_size ,							-1   ).bool()

									__lowercase=     model(__lowerCAmelCase ,							__lowerCAmelCase   )
									# model only returns predictions for masked patches
									__lowercase=     mask.sum().item()
									__lowercase=     3 * self.tubelet_size * self.patch_size**2
									self.parent.assertEqual(result.logits.shape ,							(self.batch_size, num_masked_patches, decoder_num_labels)   )







							def      _A			(self   ):
									__lowercase=     self.prepare_config_and_inputs()
									__lowercase=     config_and_inputs
									__lowercase=     {'''pixel_values''': pixel_values}
									return config, inputs_dict





@require_torch
class  A       (  A_   ,							A_   ,							unittest.TestCase					):
							UpperCamelCase_  :				str			=(
							    (VideoMAEModel, VideoMAEForPreTraining, VideoMAEForVideoClassification) if is_torch_available() else ()
							)
							UpperCamelCase_  :				List[str]			=(
							    {'''feature-extraction''': VideoMAEModel, '''video-classification''': VideoMAEForVideoClassification}
							    if is_torch_available()
							    else {}
							)

							UpperCamelCase_  :				Dict			=False
							UpperCamelCase_  :				Dict			=False
							UpperCamelCase_  :				Any			=False
							UpperCamelCase_  :				Optional[Any]			=False


							def      _A			(self   ):
									__lowercase=     VideoMAEModelTester(self   )
									__lowercase=     ConfigTester(self ,							config_class=__lowerCAmelCase ,							has_text_modality=__lowerCAmelCase ,							hidden_size=3_7   )


							def      _A			(self ,							lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase=False   ):
									__lowercase=     copy.deepcopy(__lowerCAmelCase   )

									if model_class == VideoMAEForPreTraining:
											# important: each video needs to have the same number of masked patches
											# hence we define a single mask, which we then repeat for each example in the batch
											__lowercase=     torch.ones((self.model_tester.num_masks,)   )
											__lowercase=     torch.cat([mask, torch.zeros(self.model_tester.seq_length - mask.size(0   )   )]   )
											__lowercase=     mask.expand(self.model_tester.batch_size ,							-1   ).bool()
											__lowercase=     bool_masked_pos.to(__lowerCAmelCase   )

									if return_labels:
											if model_class in [
											    *get_values(__lowerCAmelCase   ),
											]:
													__lowercase=     torch.zeros(
													    self.model_tester.batch_size ,							dtype=torch.long ,							device=__lowerCAmelCase   )

									return inputs_dict


							def      _A			(self   ):
									self.config_tester.run_common_tests()


							@unittest.skip(reason='VideoMAE does not use inputs_embeds'   )
							def      _A			(self   ):
									pass


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											__lowercase=     model_class(__lowerCAmelCase   )
											self.assertIsInstance(model.get_input_embeddings() ,							(nn.Module)   )
											__lowercase=     model.get_output_embeddings()
											self.assertTrue(x is None or isinstance(__lowerCAmelCase ,							nn.Linear   )   )


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											__lowercase=     model_class(__lowerCAmelCase   )
											__lowercase=     inspect.signature(model.forward   )
											# signature.parameters is an OrderedDict => so arg_names order is deterministic
											__lowercase=     [*signature.parameters.keys()]

											__lowercase=     ['''pixel_values''']
											self.assertListEqual(arg_names[:1] ,							__lowerCAmelCase   )


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__lowerCAmelCase   )


							def      _A			(self   ):
									__lowercase=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_pretraining(*__lowerCAmelCase   )


							@slow
							def      _A			(self   ):
									for model_name in VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__lowercase=     VideoMAEModel.from_pretrained(__lowerCAmelCase   )
											self.assertIsNotNone(__lowerCAmelCase   )


							def      _A			(self   ):
									if not self.has_attentions:
											pass

									else:
											__lowercase=     self.model_tester.prepare_config_and_inputs_for_common()
											__lowercase=     True

											for model_class in self.all_model_classes:
													__lowercase=     self.model_tester.seq_length - self.model_tester.num_masks
													__lowercase=     (
													    num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length
													)

													__lowercase=     True
													__lowercase=     False
													__lowercase=     True
													__lowercase=     model_class(__lowerCAmelCase   )
													model.to(__lowerCAmelCase   )
													model.eval()
													with torch.no_grad():
															__lowercase=     model(**self._prepare_for_class(__lowerCAmelCase ,							__lowerCAmelCase   )   )
													__lowercase=     outputs.attentions
													self.assertEqual(len(__lowerCAmelCase   ) ,							self.model_tester.num_hidden_layers   )

													# check that output_attentions also work using config
													del inputs_dict["output_attentions"]
													__lowercase=     True
													__lowercase=     model_class(__lowerCAmelCase   )
													model.to(__lowerCAmelCase   )
													model.eval()
													with torch.no_grad():
															__lowercase=     model(**self._prepare_for_class(__lowerCAmelCase ,							__lowerCAmelCase   )   )
													__lowercase=     outputs.attentions
													self.assertEqual(len(__lowerCAmelCase   ) ,							self.model_tester.num_hidden_layers   )

													self.assertListEqual(
													    list(attentions[0].shape[-3:]   ) ,							[self.model_tester.num_attention_heads, seq_len, seq_len] ,							)
													__lowercase=     len(__lowerCAmelCase   )

													# Check attention is always last and order is fine
													__lowercase=     True
													__lowercase=     True
													__lowercase=     model_class(__lowerCAmelCase   )
													model.to(__lowerCAmelCase   )
													model.eval()
													with torch.no_grad():
															__lowercase=     model(**self._prepare_for_class(__lowerCAmelCase ,							__lowerCAmelCase   )   )

													self.assertEqual(out_len + 1 ,							len(__lowerCAmelCase   )   )

													__lowercase=     outputs.attentions

													self.assertEqual(len(__lowerCAmelCase   ) ,							self.model_tester.num_hidden_layers   )
													self.assertListEqual(
													    list(self_attentions[0].shape[-3:]   ) ,							[self.model_tester.num_attention_heads, seq_len, seq_len] ,							)


							def      _A			(self   ):
									def check_hidden_states_output(lowerCAmelCase ,							lowerCAmelCase ,							lowerCAmelCase   ):
											__lowercase=     model_class(__lowerCAmelCase   )
											model.to(__lowerCAmelCase   )
											model.eval()

											with torch.no_grad():
													__lowercase=     model(**self._prepare_for_class(__lowerCAmelCase ,							__lowerCAmelCase   )   )

											__lowercase=     outputs.hidden_states
											__lowercase=     self.model_tester.num_hidden_layers + 1
											self.assertEqual(len(__lowerCAmelCase   ) ,							__lowerCAmelCase   )

											__lowercase=     self.model_tester.seq_length - self.model_tester.num_masks
											__lowercase=     num_visible_patches if model_class == VideoMAEForPreTraining else self.model_tester.seq_length

											self.assertListEqual(
											    list(hidden_states[0].shape[-2:]   ) ,							[seq_length, self.model_tester.hidden_size] ,							)

									__lowercase=     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											__lowercase=     True
											check_hidden_states_output(__lowerCAmelCase ,							__lowerCAmelCase ,							__lowerCAmelCase   )

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											__lowercase=     True

											check_hidden_states_output(__lowerCAmelCase ,							__lowerCAmelCase ,							__lowerCAmelCase   )







							@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'   )
							def      _A			(self   ):
									pass







def 			_lowerCamelCase(      )		->      List[str]:

		'''simple docstring'''







		__lowercase=     hf_hub_download(
		    repo_id='hf-internal-testing/spaghetti-video'  ,					filename='eating_spaghetti.npy'  ,					repo_type='dataset'      )
		__lowercase=     np.load(lowercase__      )
		return list(lowercase__      )





@require_torch
@require_vision
class  A       (  unittest.TestCase					):


							@cached_property
							def      _A			(self   ):
									return (
									    VideoMAEImageProcessor(image_mean=[0.5, 0.5, 0.5] ,							image_std=[0.5, 0.5, 0.5]   )
									    if is_vision_available()
									    else None
									)


							@slow
							def      _A			(self   ):
									__lowercase=     VideoMAEForVideoClassification.from_pretrained('MCG-NJU/videomae-base-finetuned-kinetics'   ).to(
									    __lowerCAmelCase   )

									__lowercase=     self.default_image_processor
									__lowercase=     prepare_video()
									__lowercase=     image_processor(__lowerCAmelCase ,							return_tensors='pt'   ).to(__lowerCAmelCase   )

									# forward pass
									with torch.no_grad():
											__lowercase=     model(**__lowerCAmelCase   )

									# verify the logits
									__lowercase=     torch.Size((1, 4_0_0)   )
									self.assertEqual(outputs.logits.shape ,							__lowerCAmelCase   )

									__lowercase=     torch.tensor([0.36_69, -0.06_88, -0.24_21]   ).to(__lowerCAmelCase   )

									self.assertTrue(torch.allclose(outputs.logits[0, :3] ,							__lowerCAmelCase ,							atol=1E-4   )   )







							@slow
							def      _A			(self   ):
									__lowercase=     VideoMAEForPreTraining.from_pretrained('MCG-NJU/videomae-base-short'   ).to(__lowerCAmelCase   )

									__lowercase=     self.default_image_processor
									__lowercase=     prepare_video()
									__lowercase=     image_processor(__lowerCAmelCase ,							return_tensors='pt'   ).to(__lowerCAmelCase   )

									# add boolean mask, indicating which patches to mask
									__lowercase=     hf_hub_download(repo_id='hf-internal-testing/bool-masked-pos' ,							filename='bool_masked_pos.pt'   )
									__lowercase=     torch.load(__lowerCAmelCase   )

									# forward pass
									with torch.no_grad():
											__lowercase=     model(**__lowerCAmelCase   )

									# verify the logits
									__lowercase=     torch.Size([1, 1_4_0_8, 1_5_3_6]   )
									__lowercase=     torch.tensor(
									    [[0.79_94, 0.96_12, 0.85_08], [0.74_01, 0.89_58, 0.83_02], [0.58_62, 0.74_68, 0.73_25]] ,							device=__lowerCAmelCase   )
									self.assertEqual(outputs.logits.shape ,							__lowerCAmelCase   )
									self.assertTrue(torch.allclose(outputs.logits[0, :3, :3] ,							__lowerCAmelCase ,							atol=1E-4   )   )

									# verify the loss (`config.norm_pix_loss` = `True`)
									__lowercase=     torch.tensor([0.51_42] ,							device=__lowerCAmelCase   )
									self.assertTrue(torch.allclose(outputs.loss ,							__lowerCAmelCase ,							atol=1E-4   )   )

									# verify the loss (`config.norm_pix_loss` = `False`)
									__lowercase=     VideoMAEForPreTraining.from_pretrained('MCG-NJU/videomae-base-short' ,							norm_pix_loss=__lowerCAmelCase   ).to(
									    __lowerCAmelCase   )

									with torch.no_grad():
											__lowercase=     model(**__lowerCAmelCase   )

									__lowercase=     torch.tensor(torch.tensor([0.64_69]   ) ,							device=__lowerCAmelCase   )
									self.assertTrue(torch.allclose(outputs.loss ,							__lowerCAmelCase ,							atol=1E-4   )   )














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


SCREAMING_SNAKE_CASE__   :      List[str]				=		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							,       *__lowerCAmelCase							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       **__lowerCAmelCase  ):
								"""simple docstring"""
								super().__init__(*__lowerCAmelCase							,       **__lowerCAmelCase  )
								__magic_name__		:Any   =   eval_examples
								__magic_name__		:str   =   post_process_function
								__magic_name__		:int   =   quant_trainer_args
								__magic_name__		:List[str]   =   1_2_8  # default number of calibration samples







				def       A						(    self							,       __lowerCAmelCase=None  ):
								"""simple docstring"""
								if calib_dataset is None and self.calib_dataset is None:
												raise ValueError('''Trainer: calibration requires an calib_dataset.'''  )
								__magic_name__		:Optional[Any]   =   calib_dataset if calib_dataset is not None else self.calib_dataset

								__magic_name__		: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       A						(    self							,       __lowerCAmelCase=None  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   self.train_dataset if calib_dataset is None else calib_dataset
								__magic_name__		:Any   =   self.get_calib_dataloader(__lowerCAmelCase  )

								__magic_name__		:List[str]   =   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
												__magic_name__  ,					__magic_name__  ,					__magic_name__		:str   =   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  )
								__magic_name__		:Any   =   model







				def       A						(    self							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       __lowerCAmelCase=None							,       __lowerCAmelCase = "eval"  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   self.eval_dataset if eval_dataset is None else eval_dataset
								__magic_name__		:Optional[Any]   =   self.get_eval_dataloader(__lowerCAmelCase  )
								__magic_name__		:str   =   self.eval_examples if eval_examples is None else eval_examples

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

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

												self.log(__lowerCAmelCase  )
								else:
												__magic_name__		:List[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()  )

								__magic_name__		:Optional[Any]   =   self.callback_handler.on_evaluate(self.args							,       self.state							,       self.control							,       __lowerCAmelCase  )
								return metrics







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase=None							,       __lowerCAmelCase = "test"  ):
								"""simple docstring"""
								__magic_name__		:int   =   self.get_test_dataloader(__lowerCAmelCase  )

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

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

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

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







				def       A						(    self							,       __lowerCAmelCase="./"  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   self.eval_dataset
								__magic_name__		:Any   =   self.get_eval_dataloader(__lowerCAmelCase  )

								__magic_name__		:int   =   next(iter(__lowerCAmelCase  )  )

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

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

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

								__magic_name__		:Any   =   True

								__magic_name__		:Optional[int]   =   self.model.to(__lowerCAmelCase  )

								model.eval()
								model.float()

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

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

								__magic_name__		:Dict   =   {0: '''batch_size''', 1: '''seq_len'''}

								torch.onnx.export(
								    __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       export_params=__lowerCAmelCase							,       opset_version=1_3							,       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'''  )






