






import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartTokenizer, MBartTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
)

from ...test_tokenization_common import TokenizerTesterMixin


SCREAMING_SNAKE_CASE__				     =			get_tests_dir('''fixtures/test_sentencepiece.model''')


if is_torch_available():
							from transformers.models.mbart.modeling_mbart import shift_tokens_right

SCREAMING_SNAKE_CASE__				     =			25_0004
SCREAMING_SNAKE_CASE__				     =			25_0020






@require_sentencepiece
@require_tokenizers
class 			_UpperCamelCase(       __lowerCamelCase		,				unittest.TestCase		):
		__SCREAMING_SNAKE_CASE  :       List[str]							=							MBartTokenizer
		__SCREAMING_SNAKE_CASE  :       Tuple							=							MBartTokenizerFast
		__SCREAMING_SNAKE_CASE  :       Optional[int]							=							True
		__SCREAMING_SNAKE_CASE  :       Tuple							=							True

		def    __lowerCAmelCase  (					self		:    Optional[Any]    ):



					'''simple docstring'''



					super().setUp()

					# We have a SentencePiece fixture for testing
					__a						:  Any									=     MBartTokenizer(SCREAMING_SNAKE_CASE__    ,       keep_accents=SCREAMING_SNAKE_CASE__    )
					tokenizer.save_pretrained(self.tmpdirname    )

		def    __lowerCAmelCase  (					self		:    Dict    ):



					'''simple docstring'''



					__a						:  List[Any]									=     MBartTokenizer(SCREAMING_SNAKE_CASE__    ,       keep_accents=SCREAMING_SNAKE_CASE__    )

					__a						:  List[Any]									=     tokenizer.tokenize('This is a test'    )
					self.assertListEqual(SCREAMING_SNAKE_CASE__    ,       ['▁This', '▁is', '▁a', '▁t', 'est']    )

					self.assertListEqual(
					    tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__    )    ,       [value + tokenizer.fairseq_offset for value in [2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]]    ,       )

					__a						:  Optional[int]									=     tokenizer.tokenize('I was born in 92000, and this is falsé.'    )
					self.assertListEqual(
					    SCREAMING_SNAKE_CASE__    ,       [
					        SPIECE_UNDERLINE + 'I',
					        SPIECE_UNDERLINE + 'was',
					        SPIECE_UNDERLINE + 'b',
					        'or',
					        'n',
					        SPIECE_UNDERLINE + 'in',
					        SPIECE_UNDERLINE + '',
					        '9',
					        '2',
					        '0',
					        '0',
					        '0',
					        ',',
					        SPIECE_UNDERLINE + 'and',
					        SPIECE_UNDERLINE + 'this',
					        SPIECE_UNDERLINE + 'is',
					        SPIECE_UNDERLINE + 'f',
					        'al',
					        's',
					        'é',
					        '.',
					    ]    ,       )
					__a						:  Optional[Any]									=     tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__    )
					self.assertListEqual(
					    SCREAMING_SNAKE_CASE__    ,       [
					        value + tokenizer.fairseq_offset
					        for value in [8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 2, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 2, 4]
					        #                                       ^ unk: 2 + 1 = 3                  unk: 2 + 1 = 3 ^
					    ]    ,       )

					__a						:  List[str]									=     tokenizer.convert_ids_to_tokens(SCREAMING_SNAKE_CASE__    )
					self.assertListEqual(
					    SCREAMING_SNAKE_CASE__    ,       [
					        SPIECE_UNDERLINE + 'I',
					        SPIECE_UNDERLINE + 'was',
					        SPIECE_UNDERLINE + 'b',
					        'or',
					        'n',
					        SPIECE_UNDERLINE + 'in',
					        SPIECE_UNDERLINE + '',
					        '<unk>',
					        '2',
					        '0',
					        '0',
					        '0',
					        ',',
					        SPIECE_UNDERLINE + 'and',
					        SPIECE_UNDERLINE + 'this',
					        SPIECE_UNDERLINE + 'is',
					        SPIECE_UNDERLINE + 'f',
					        'al',
					        's',
					        '<unk>',
					        '.',
					    ]    ,       )



		def    __lowerCAmelCase  (					self		:    str    ):



					'''simple docstring'''



					if not self.test_slow_tokenizer:
								# as we don't have a slow version, we can't compare the outputs between slow and fast versions
								return

					__a						:  Tuple									=     (self.rust_tokenizer_class, 'hf-internal-testing/tiny-random-mbart', {})
					for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
								with self.subTest(f'''{tokenizer.__class__.__name__} ({pretrained_name})'''    ):
											__a						:  Tuple									=     self.rust_tokenizer_class.from_pretrained(SCREAMING_SNAKE_CASE__    ,       **SCREAMING_SNAKE_CASE__    )
											__a						:  Dict									=     self.tokenizer_class.from_pretrained(SCREAMING_SNAKE_CASE__    ,       **SCREAMING_SNAKE_CASE__    )

											__a						:  str									=     tempfile.mkdtemp()

											__a						:  Any									=     tokenizer_r.save_pretrained(SCREAMING_SNAKE_CASE__    )
											__a						:  Dict									=     tokenizer_p.save_pretrained(SCREAMING_SNAKE_CASE__    )

											# Checks it save with the same files + the tokenizer.json file for the fast one
											self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files    )    )
											__a						:  List[str]									=     tuple(f for f in tokenizer_r_files if 'tokenizer.json' not in f    )
											self.assertSequenceEqual(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )

											# Checks everything loads correctly in the same way
											__a						:  Dict									=     tokenizer_r.from_pretrained(SCREAMING_SNAKE_CASE__    )
											__a						:  List[Any]									=     tokenizer_p.from_pretrained(SCREAMING_SNAKE_CASE__    )

											# Check special tokens are set accordingly on Rust and Python
											for key in tokenizer_pp.special_tokens_map:
														self.assertTrue(hasattr(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )    )
														# self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
														# self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

											shutil.rmtree(SCREAMING_SNAKE_CASE__    )

											# Save tokenizer rust, legacy_format=True
											__a						:  Optional[Any]									=     tempfile.mkdtemp()

											__a						:  Optional[Any]									=     tokenizer_r.save_pretrained(SCREAMING_SNAKE_CASE__    ,       legacy_format=SCREAMING_SNAKE_CASE__    )
											__a						:  int									=     tokenizer_p.save_pretrained(SCREAMING_SNAKE_CASE__    )

											# Checks it save with the same files
											self.assertSequenceEqual(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )

											# Checks everything loads correctly in the same way
											__a						:  Optional[Any]									=     tokenizer_r.from_pretrained(SCREAMING_SNAKE_CASE__    )
											__a						:  Optional[int]									=     tokenizer_p.from_pretrained(SCREAMING_SNAKE_CASE__    )

											# Check special tokens are set accordingly on Rust and Python
											for key in tokenizer_pp.special_tokens_map:
														self.assertTrue(hasattr(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )    )

											shutil.rmtree(SCREAMING_SNAKE_CASE__    )

											# Save tokenizer rust, legacy_format=False
											__a						:  Tuple									=     tempfile.mkdtemp()

											__a						:  List[Any]									=     tokenizer_r.save_pretrained(SCREAMING_SNAKE_CASE__    ,       legacy_format=SCREAMING_SNAKE_CASE__    )
											__a						:  Tuple									=     tokenizer_p.save_pretrained(SCREAMING_SNAKE_CASE__    )

											# Checks it saved the tokenizer.json file
											self.assertTrue(any('tokenizer.json' in f for f in tokenizer_r_files    )    )

											# Checks everything loads correctly in the same way
											__a						:  Union[str, Any]									=     tokenizer_r.from_pretrained(SCREAMING_SNAKE_CASE__    )
											__a						:  Union[str, Any]									=     tokenizer_p.from_pretrained(SCREAMING_SNAKE_CASE__    )

											# Check special tokens are set accordingly on Rust and Python
											for key in tokenizer_pp.special_tokens_map:
														self.assertTrue(hasattr(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )    )

											shutil.rmtree(SCREAMING_SNAKE_CASE__    )






@require_torch
@require_sentencepiece
@require_tokenizers
class 			_UpperCamelCase(       unittest.TestCase		):
		__SCREAMING_SNAKE_CASE  :       List[Any]							=							'''facebook/mbart-large-en-ro'''
		__SCREAMING_SNAKE_CASE  :       Optional[Any]							=							[
		    ''' UN Chief Says There Is No Military Solution in Syria''',
		    ''' Secretary-General Ban Ki-moon says his response to Russia\'s stepped up military support for Syria is that "there is no military solution" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.''',
		]
		__SCREAMING_SNAKE_CASE  :       Dict							=							[
		    '''Şeful ONU declară că nu există o soluţie militară în Siria''',
		    '''Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei'''
		    ''' pentru Siria este că "nu există o soluţie militară" la conflictul de aproape cinci ani şi că noi arme nu vor'''
		    ''' face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.''',
		]
		__SCREAMING_SNAKE_CASE  :       Optional[int]							=							[8274, 12_7873, 2_5916, 7, 8622, 2071, 438, 6_7485, 53, 18_7895, 23, 5_1712, 2, EN_CODE]

		@classmethod
		def    __lowerCAmelCase  (					cls		:    List[str]    ):



					'''simple docstring'''



					__a						:  MBartTokenizer									=     MBartTokenizer.from_pretrained(
					    cls.checkpoint_name    ,       src_lang='en_XX'    ,       tgt_lang='ro_RO'    )
					__a						:  Dict									=     1
					return cls

		def    __lowerCAmelCase  (					self		:    Union[str, Any]    ):



					'''simple docstring'''



					self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ar_AR']    ,       2_5_0_0_0_1    )
					self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['en_EN']    ,       2_5_0_0_0_4    )
					self.assertEqual(self.tokenizer.fairseq_tokens_to_ids['ro_RO']    ,       2_5_0_0_2_0    )

		def    __lowerCAmelCase  (					self		:    Optional[Any]    ):



					'''simple docstring'''



					__a						:  Dict									=     self.tokenizer.batch_encode_plus(self.src_text    ).input_ids[0]
					self.assertListEqual(self.expected_src_tokens    ,       SCREAMING_SNAKE_CASE__    )

		def    __lowerCAmelCase  (					self		:    Any    ):



					'''simple docstring'''



					self.assertIn(SCREAMING_SNAKE_CASE__    ,       self.tokenizer.all_special_ids    )
					__a						:  List[str]									=     [RO_CODE, 8_8_4, 9_0_1_9, 9_6, 9, 9_1_6, 8_6_7_9_2, 3_6, 1_8_7_4_3, 1_5_5_9_6, 5, 2]
					__a						:  Union[str, Any]									=     self.tokenizer.decode(SCREAMING_SNAKE_CASE__    ,       skip_special_tokens=SCREAMING_SNAKE_CASE__    )
					__a						:  Tuple									=     self.tokenizer.decode(generated_ids[1:]    ,       skip_special_tokens=SCREAMING_SNAKE_CASE__    )
					self.assertEqual(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )
					self.assertNotIn(self.tokenizer.eos_token    ,       SCREAMING_SNAKE_CASE__    )

		def    __lowerCAmelCase  (					self		:    Optional[int]    ):



					'''simple docstring'''



					__a						:  Any									=     ['this is gunna be a long sentence ' * 2_0]
					assert isinstance(src_text[0]    ,       SCREAMING_SNAKE_CASE__    )
					__a						:  int									=     1_0
					__a						:  Dict									=     self.tokenizer(SCREAMING_SNAKE_CASE__    ,       max_length=SCREAMING_SNAKE_CASE__    ,       truncation=SCREAMING_SNAKE_CASE__    ).input_ids[0]
					self.assertEqual(ids[-2]    ,       2    )
					self.assertEqual(ids[-1]    ,       SCREAMING_SNAKE_CASE__    )
					self.assertEqual(len(SCREAMING_SNAKE_CASE__    )    ,       SCREAMING_SNAKE_CASE__    )

		def    __lowerCAmelCase  (					self		:    Optional[Any]    ):



					'''simple docstring'''



					self.assertListEqual(self.tokenizer.convert_tokens_to_ids(['<mask>', 'ar_AR']    )    ,       [2_5_0_0_2_6, 2_5_0_0_0_1]    )

		def    __lowerCAmelCase  (					self		:    str    ):



					'''simple docstring'''



					__a						:  Tuple									=     tempfile.mkdtemp()
					__a						:  Optional[Any]									=     self.tokenizer.fairseq_tokens_to_ids
					self.tokenizer.save_pretrained(SCREAMING_SNAKE_CASE__    )
					__a						:  List[Any]									=     MBartTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__    )
					self.assertDictEqual(new_tok.fairseq_tokens_to_ids    ,       SCREAMING_SNAKE_CASE__    )

		@require_torch
		def    __lowerCAmelCase  (					self		:    List[str]    ):



					'''simple docstring'''



					__a						:  str									=     self.tokenizer(self.src_text    ,       text_target=self.tgt_text    ,       padding=SCREAMING_SNAKE_CASE__    ,       return_tensors='pt'    )
					__a						:  Any									=     shift_tokens_right(batch['labels']    ,       self.tokenizer.pad_token_id    )

					# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
					assert batch.input_ids[1][-2:].tolist() == [2, EN_CODE]
					assert batch.decoder_input_ids[1][0].tolist() == RO_CODE
					assert batch.decoder_input_ids[1][-1] == 2
					assert batch.labels[1][-2:].tolist() == [2, RO_CODE]

		@require_torch
		def    __lowerCAmelCase  (					self		:    Optional[int]    ):



					'''simple docstring'''



					__a						:  str									=     self.tokenizer(
					    self.src_text    ,       text_target=self.tgt_text    ,       padding=SCREAMING_SNAKE_CASE__    ,       truncation=SCREAMING_SNAKE_CASE__    ,       max_length=len(self.expected_src_tokens    )    ,       return_tensors='pt'    ,       )

					__a						:  Optional[int]									=     shift_tokens_right(batch['labels']    ,       self.tokenizer.pad_token_id    )

					self.assertIsInstance(SCREAMING_SNAKE_CASE__    ,       SCREAMING_SNAKE_CASE__    )

					self.assertEqual((2, 1_4)    ,       batch.input_ids.shape    )
					self.assertEqual((2, 1_4)    ,       batch.attention_mask.shape    )
					__a						:  List[Any]									=     batch.input_ids.tolist()[0]
					self.assertListEqual(self.expected_src_tokens    ,       SCREAMING_SNAKE_CASE__    )
					self.assertEqual(2    ,       batch.decoder_input_ids[0, -1]    )  # EOS
					# Test that special tokens are reset
					self.assertEqual(self.tokenizer.prefix_tokens    ,       []    )
					self.assertEqual(self.tokenizer.suffix_tokens    ,       [self.tokenizer.eos_token_id, EN_CODE]    )

		def    __lowerCAmelCase  (					self		:    List[Any]    ):



					'''simple docstring'''



					__a						:  str									=     self.tokenizer(self.src_text    ,       padding=SCREAMING_SNAKE_CASE__    ,       truncation=SCREAMING_SNAKE_CASE__    ,       max_length=3    ,       return_tensors='pt'    )
					__a						:  Optional[int]									=     self.tokenizer(
					    text_target=self.tgt_text    ,       padding=SCREAMING_SNAKE_CASE__    ,       truncation=SCREAMING_SNAKE_CASE__    ,       max_length=1_0    ,       return_tensors='pt'    )
					__a						:  Optional[Any]									=     targets['input_ids']
					__a						:  Dict									=     shift_tokens_right(SCREAMING_SNAKE_CASE__    ,       self.tokenizer.pad_token_id    )

					self.assertEqual(batch.input_ids.shape[1]    ,       3    )
					self.assertEqual(batch.decoder_input_ids.shape[1]    ,       1_0    )



		@require_torch
		def    __lowerCAmelCase  (					self		:    int    ):



					'''simple docstring'''



					__a						:  Dict									=     self.tokenizer._build_translation_inputs(
					    'A test'    ,       return_tensors='pt'    ,       src_lang='en_XX'    ,       tgt_lang='ar_AR'    )

					self.assertEqual(
					    nested_simplify(SCREAMING_SNAKE_CASE__    )    ,       {
					        # A, test, EOS, en_XX
					        'input_ids': [[6_2, 3_0_3_4, 2, 2_5_0_0_0_4]],
					        'attention_mask': [[1, 1, 1, 1]],
					        # ar_AR
					        'forced_bos_token_id': 2_5_0_0_0_1,
					    }    ,       )













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   =  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 					SCREAMING_SNAKE_CASE_  (       _a ):
       """simple docstring"""


       def __init__(				self	:Dict,      *snake_case	:int,      snake_case	:Any=None,      snake_case	:str=None,      snake_case	:int=None,      **snake_case	:List[str]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Dict,      snake_case	:List[str]=None):



          """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(snake_case,      description='Calibration')

          return DataLoader(
              snake_case,      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=snake_case,      )





       def 				UpperCamelCase__						(				self	:Any,      snake_case	:Union[str, Any]=None):



          """simple docstring"""

          _lowercase              =self.train_dataset if calib_dataset is None else calib_dataset
          _lowercase              =self.get_calib_dataloader(snake_case)

          _lowercase              =self.model
          quant_trainer.configure_model(snake_case,      self.quant_trainer_args,      calib=snake_case)
          model.eval()
          quant_trainer.enable_calibration(snake_case)

          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(snake_case):
             # Prediction step
             _lowercase       ,       _lowercase       ,       _lowercase              =self.prediction_step(snake_case,      snake_case,      prediction_loss_only=snake_case)
             if (step + 1) * calib_dataloader.batch_size >= self.calib_num:
                break

          quant_trainer.finish_calibration(snake_case,      self.quant_trainer_args)
          _lowercase              =model





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:Union[str, Any]=None,      snake_case	:Optional[int]=None,      snake_case	:Union[str, Any]=None,      snake_case	:str = "eval"):



          """simple docstring"""

          _lowercase              =self.eval_dataset if eval_dataset is None else eval_dataset
          _lowercase              =self.get_eval_dataloader(snake_case)
          _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(
                 snake_case,      description='Evaluation',      prediction_loss_only=True if compute_metrics is None else None,      ignore_keys=snake_case,      )
          finally:
             _lowercase              =compute_metrics

          if self.post_process_function is not None and self.compute_metrics is not None:
             _lowercase              =self.post_process_function(snake_case,      snake_case,      output.predictions)
             _lowercase              =self.compute_metrics(snake_case)

             # Prefix all keys with metric_key_prefix + '_'
             for key in list(metrics.keys()):
                if not key.startswith(f'''{metric_key_prefix}_'''):
                   _lowercase              =metrics.pop(snake_case)

             self.log(snake_case)
          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,      snake_case)
          return metrics





       def 				UpperCamelCase__						(				self	:Optional[int],      snake_case	:str,      snake_case	:Optional[Any],      snake_case	:Any=None,      snake_case	:str = "test"):



          """simple docstring"""

          _lowercase              =self.get_test_dataloader(snake_case)

          # 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(
                 snake_case,      description='Prediction',      prediction_loss_only=True if compute_metrics is None else None,      ignore_keys=snake_case,      )
          finally:
             _lowercase              =compute_metrics

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

          _lowercase              =self.post_process_function(snake_case,      snake_case,      output.predictions,      'predict')
          _lowercase              =self.compute_metrics(snake_case)

          # Prefix all keys with metric_key_prefix + '_'
          for key in list(metrics.keys()):
             if not key.startswith(f'''{metric_key_prefix}_'''):
                _lowercase              =metrics.pop(snake_case)

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





       def 				UpperCamelCase__						(				self	:List[str],      snake_case	:Dict="./"):



          """simple docstring"""

          _lowercase              =self.eval_dataset
          _lowercase              =self.get_eval_dataloader(snake_case)

          _lowercase              =next(iter(snake_case))

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

          # convert to tuple
          _lowercase              =tuple(v.to(snake_case) 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(snake_case)

          model.eval()
          model.float()

          _lowercase              =model.module if hasattr(snake_case,      'module') else model
          quant_trainer.configure_model(snake_case,      self.quant_trainer_args)

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

          _lowercase              ={0: 'batch_size', 1: 'seq_len'}

          torch.onnx.export(
              snake_case,      snake_case,      snake_case,      export_params=snake_case,      opset_version=13,      do_constant_folding=snake_case,      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=snake_case,      )
          logger.info('onnx export finished')


