






import collections
from typing import List, Optional, Union

from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
from ..bert.tokenization_bert import BertTokenizer


SCREAMING_SNAKE_CASE__   :      Union[str, Any]				=		logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__   :      Optional[int]				=		{"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

SCREAMING_SNAKE_CASE__   :      Any				=		{
    """vocab_file""": {
        """facebook/dpr-ctx_encoder-single-nq-base""": (
            """https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/vocab.txt"""
        ),
        """facebook/dpr-ctx_encoder-multiset-base""": (
            """https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/vocab.txt"""
        ),
    },
    """tokenizer_file""": {
        """facebook/dpr-ctx_encoder-single-nq-base""": (
            """https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/tokenizer.json"""
        ),
        """facebook/dpr-ctx_encoder-multiset-base""": (
            """https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/tokenizer.json"""
        ),
    },
}
SCREAMING_SNAKE_CASE__   :      Tuple				=		{
    """vocab_file""": {
        """facebook/dpr-question_encoder-single-nq-base""": (
            """https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/vocab.txt"""
        ),
        """facebook/dpr-question_encoder-multiset-base""": (
            """https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/vocab.txt"""
        ),
    },
    """tokenizer_file""": {
        """facebook/dpr-question_encoder-single-nq-base""": (
            """https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/tokenizer.json"""
        ),
        """facebook/dpr-question_encoder-multiset-base""": (
            """https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/tokenizer.json"""
        ),
    },
}
SCREAMING_SNAKE_CASE__   :      Optional[Any]				=		{
    """vocab_file""": {
        """facebook/dpr-reader-single-nq-base""": (
            """https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/vocab.txt"""
        ),
        """facebook/dpr-reader-multiset-base""": (
            """https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/vocab.txt"""
        ),
    },
    """tokenizer_file""": {
        """facebook/dpr-reader-single-nq-base""": (
            """https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/tokenizer.json"""
        ),
        """facebook/dpr-reader-multiset-base""": (
            """https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/tokenizer.json"""
        ),
    },
}

SCREAMING_SNAKE_CASE__   :      List[Any]				=		{
    """facebook/dpr-ctx_encoder-single-nq-base""": 5_12,
    """facebook/dpr-ctx_encoder-multiset-base""": 5_12,
}
SCREAMING_SNAKE_CASE__   :      Optional[Any]				=		{
    """facebook/dpr-question_encoder-single-nq-base""": 5_12,
    """facebook/dpr-question_encoder-multiset-base""": 5_12,
}
SCREAMING_SNAKE_CASE__   :      List[Any]				=		{
    """facebook/dpr-reader-single-nq-base""": 5_12,
    """facebook/dpr-reader-multiset-base""": 5_12,
}


SCREAMING_SNAKE_CASE__   :      Dict				=		{
    """facebook/dpr-ctx_encoder-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-ctx_encoder-multiset-base""": {"""do_lower_case""": True},
}
SCREAMING_SNAKE_CASE__   :      List[str]				=		{
    """facebook/dpr-question_encoder-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-question_encoder-multiset-base""": {"""do_lower_case""": True},
}
SCREAMING_SNAKE_CASE__   :      Dict				=		{
    """facebook/dpr-reader-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-reader-multiset-base""": {"""do_lower_case""": True},
}
class     lowerCamelCase_					(   lowerCamelCase			):
				a__       							=	VOCAB_FILES_NAMES
				a__       							=	CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
				a__       							=	CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				a__       							=	CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION







class     lowerCamelCase_					(   lowerCamelCase			):
				a__       							=	VOCAB_FILES_NAMES
				a__       							=	QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
				a__       							=	QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				a__       							=	QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION







SCREAMING_SNAKE_CASE__   :      Dict				=		collections.namedtuple(
    """DPRSpanPrediction""", ["""span_score""", """relevance_score""", """doc_id""", """start_index""", """end_index""", """text"""]
)

SCREAMING_SNAKE_CASE__   :      List[str]				=		collections.namedtuple("""DPRReaderOutput""", ["""start_logits""", """end_logits""", """relevance_logits"""])


SCREAMING_SNAKE_CASE__   :      Any				=		r"""
    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.
    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),
    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`
    with the format:

    ```
    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>
    ```

    Args:
        questions (`str` or `List[str]`):
            The questions to be encoded. You can specify one question for many passages. In this case, the question
            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in
            `titles` or `texts`.
        titles (`str` or `List[str]`):
            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.
        texts (`str` or `List[str]`):
            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
            Activates and controls padding. Accepts the following values:

            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
              if provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
              lengths).
        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):
            Activates and controls truncation. Accepts the following values:

            - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to
              the maximum acceptable input length for the model if that argument is not provided. This will truncate
              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch
              of pairs) is provided.
            - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the first
              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the
              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths
              greater than the model maximum admissible input size).
        max_length (`int`, *optional*):
                Controls the maximum length to use by one of the truncation/padding parameters.

                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
                is required by one of the truncation/padding parameters. If the model has no specific maximum input
                length (like XLNet) truncation/padding to a maximum length will be deactivated.
        return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors instead of list of python integers. Acceptable values are:

                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return Numpy `np.ndarray` objects.
        return_attention_mask (`bool`, *optional*):
            Whether or not to return the attention mask. If not set, will return the attention mask according to the
            specific tokenizer's default, defined by the `return_outputs` attribute.

            [What are attention masks?](../glossary#attention-mask)

    Returns:
        `Dict[str, List[List[int]]]`: A dictionary with the following keys:

        - `input_ids`: List of token ids to be fed to a model.
        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.
    """
@add_start_docstrings(lowerCamelCase			)
class     lowerCamelCase_					:







				def __call__(    self							,       __lowerCAmelCase							,       __lowerCAmelCase = None							,       __lowerCAmelCase = None							,       __lowerCAmelCase = False							,       __lowerCAmelCase = False							,       __lowerCAmelCase = None							,       __lowerCAmelCase = None							,       __lowerCAmelCase = None							,       **__lowerCAmelCase							,       ):
								"""simple docstring"""
								if titles is None and texts is None:
												return super().__call__(
												    __lowerCAmelCase							,       padding=__lowerCAmelCase							,       truncation=__lowerCAmelCase							,       max_length=__lowerCAmelCase							,       return_tensors=__lowerCAmelCase							,       return_attention_mask=__lowerCAmelCase							,       **__lowerCAmelCase							,       )
								elif titles is None or texts is None:
												__magic_name__		:int   =   titles if texts is None else texts
												return super().__call__(
												    __lowerCAmelCase							,       __lowerCAmelCase							,       padding=__lowerCAmelCase							,       truncation=__lowerCAmelCase							,       max_length=__lowerCAmelCase							,       return_tensors=__lowerCAmelCase							,       return_attention_mask=__lowerCAmelCase							,       **__lowerCAmelCase							,       )
								__magic_name__		:Any   =   titles if not isinstance(__lowerCAmelCase							,       __lowerCAmelCase  ) else [titles]
								__magic_name__		:str   =   texts if not isinstance(__lowerCAmelCase							,       __lowerCAmelCase  ) else [texts]
								__magic_name__		:Union[str, Any]   =   len(__lowerCAmelCase  )
								__magic_name__		:Optional[Any]   =   questions if not isinstance(__lowerCAmelCase							,       __lowerCAmelCase  ) else [questions] * n_passages
								if len(__lowerCAmelCase  ) != len(__lowerCAmelCase  ):
												raise ValueError(
												    F'''There should be as many titles than texts but got {len(__lowerCAmelCase  )} titles and {len(__lowerCAmelCase  )} texts.'''  )
								__magic_name__		:List[str]   =   super().__call__(__lowerCAmelCase							,       __lowerCAmelCase							,       padding=__lowerCAmelCase							,       truncation=__lowerCAmelCase  )['''input_ids''']
								__magic_name__		:Optional[int]   =   super().__call__(__lowerCAmelCase							,       add_special_tokens=__lowerCAmelCase							,       padding=__lowerCAmelCase							,       truncation=__lowerCAmelCase  )['''input_ids''']
								__magic_name__		:Optional[Any]   =   {
								    '''input_ids''': [
								        (encoded_question_and_title + encoded_text)[:max_length]
								        if max_length is not None and truncation
								        else encoded_question_and_title + encoded_text
								        for encoded_question_and_title, encoded_text in zip(__lowerCAmelCase							,       __lowerCAmelCase  )
								    ]
								}
								if return_attention_mask is not False:
												__magic_name__		:Optional[int]   =   []
												for input_ids in encoded_inputs["input_ids"]:
																attention_mask.append([int(input_id != self.pad_token_id  ) for input_id in input_ids]  )
												__magic_name__		:Optional[int]   =   attention_mask
								return self.pad(__lowerCAmelCase							,       padding=__lowerCAmelCase							,       max_length=__lowerCAmelCase							,       return_tensors=__lowerCAmelCase  )







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase = 1_6							,       __lowerCAmelCase = 6_4							,       __lowerCAmelCase = 4							,       ):
								"""simple docstring"""
								__magic_name__		:str   =   reader_input['''input_ids''']
								__magic_name__  ,					__magic_name__  ,					__magic_name__		:str   =   reader_output[:3]
								__magic_name__		:List[str]   =   len(__lowerCAmelCase  )
								__magic_name__		:Tuple   =   sorted(range(__lowerCAmelCase  )							,       reverse=__lowerCAmelCase							,       key=relevance_logits.__getitem__  )
								__magic_name__		:List[DPRReaderOutput]   =   []
								for doc_id in sorted_docs:
												__magic_name__		:List[Any]   =   list(input_ids[doc_id]  )
												# assuming question & title information is at the beginning of the sequence
												__magic_name__		:Dict   =   sequence_ids.index(self.sep_token_id							,       2  ) + 1  # second sep id
												if sequence_ids[-1] == self.pad_token_id:
																__magic_name__		:Optional[Any]   =   sequence_ids.index(self.pad_token_id  )
												else:
																__magic_name__		:Optional[Any]   =   len(__lowerCAmelCase  )

												__magic_name__		:int   =   self._get_best_spans(
												    start_logits=start_logits[doc_id][passage_offset:sequence_len]							,       end_logits=end_logits[doc_id][passage_offset:sequence_len]							,       max_answer_length=__lowerCAmelCase							,       top_spans=__lowerCAmelCase							,       )
												for start_index, end_index in best_spans:
																start_index += passage_offset
																end_index += passage_offset
																nbest_spans_predictions.append(
																    DPRSpanPrediction(
																        span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index]							,       relevance_score=relevance_logits[doc_id]							,       doc_id=__lowerCAmelCase							,       start_index=__lowerCAmelCase							,       end_index=__lowerCAmelCase							,       text=self.decode(sequence_ids[start_index : end_index + 1]  )							,       )  )
												if len(__lowerCAmelCase  ) >= num_spans:
																break
								return nbest_spans_predictions[:num_spans]







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       __lowerCAmelCase							,       ):
								"""simple docstring"""
								__magic_name__		:Optional[Any]   =   []
								for start_index, start_score in enumerate(__lowerCAmelCase  ):
												for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]  ):
																scores.append(((start_index, start_index + answer_length), start_score + end_score)  )
								__magic_name__		:Any   =   sorted(__lowerCAmelCase							,       key=lambda __lowerCAmelCase  : x[1]							,       reverse=__lowerCAmelCase  )
								__magic_name__		:Optional[int]   =   []
								for (start_index, end_index), score in scores:
												if start_index > end_index:
																raise ValueError(F'''Wrong span indices: [{start_index}:{end_index}]'''  )
												__magic_name__		:Any   =   end_index - start_index + 1
												if length > max_answer_length:
																raise ValueError(F'''Span is too long: {length} > {max_answer_length}'''  )
												if any(
												    start_index <= prev_start_index <= prev_end_index <= end_index
												    or prev_start_index <= start_index <= end_index <= prev_end_index
												    for (prev_start_index, prev_end_index) in chosen_span_intervals  ):
																continue
												chosen_span_intervals.append((start_index, end_index)  )

												if len(__lowerCAmelCase  ) == top_spans:
																break
								return chosen_span_intervals







@add_end_docstrings(lowerCamelCase			)
class     lowerCamelCase_					(   lowerCamelCase      ,					lowerCamelCase			):
				a__       							=	VOCAB_FILES_NAMES
				a__       							=	READER_PRETRAINED_VOCAB_FILES_MAP
				a__       							=	READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				a__       							=	READER_PRETRAINED_INIT_CONFIGURATION
				a__       							=	['''input_ids''', '''attention_mask''']














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






