






"""simple docstring"""



import argparse
import collections
from pathlib import Path

import torch
from torch.serialization import default_restore_location

from transformers import BertConfig, DPRConfig, DPRContextEncoder, DPRQuestionEncoder, DPRReader


_UpperCamelCase						:    Optional[Any]		=						collections.namedtuple(
    'CheckpointState', ['model_dict', 'optimizer_dict', 'scheduler_dict', 'offset', 'epoch', 'encoder_params']
)
def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    str       ):





  '''simple docstring'''





  print(f'Reading saved model from {model_file}'       )
  lowercase				       =			torch.load(__snake_case						,   map_location=lambda __snake_case						,   __snake_case       : default_restore_location(__snake_case						,   'cpu'       )       )
  return CheckpointState(**__snake_case       )

class 			a	:







    def __init__(     self				,						_lowerCamelCase			):
      lowercase				       =			src_file







    def      UpperCamelCase_ (     self			):
      raise NotImplementedError







    @staticmethod
    def      UpperCamelCase_ (     _lowerCamelCase				,						*_lowerCamelCase				,						**_lowerCamelCase			):
      if comp_type.startswith('c'			):
        return DPRContextEncoderState(*_lowerCamelCase				,						**_lowerCamelCase			)
      if comp_type.startswith('q'			):
        return DPRQuestionEncoderState(*_lowerCamelCase				,						**_lowerCamelCase			)
      if comp_type.startswith('r'			):
        return DPRReaderState(*_lowerCamelCase				,						**_lowerCamelCase			)
      else:
        raise ValueError('Component type must be either \'ctx_encoder\', \'question_encoder\' or \'reader\'.'			)


class 			a	(  a_      ):







    def      UpperCamelCase_ (     self			):
      lowercase				       =			DPRContextEncoder(DPRConfig(**BertConfig.get_config_dict('bert-base-uncased'			)[0]			)			)
      print(F'Loading DPR biencoder from {self.src_file}'			)
      lowercase				       =			load_states_from_checkpoint(self.src_file			)
      lowercase       ,	lowercase				       =			model.ctx_encoder, 'ctx_model.'
      # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
      lowercase				       =			{'bert_model.embeddings.position_ids': model.ctx_encoder.bert_model.embeddings.position_ids}
      for key, value in saved_state.model_dict.items():
        if key.startswith(_lowerCamelCase			):
          lowercase				       =			key[len(_lowerCamelCase			) :]
          if not key.startswith('encode_proj.'			):
            lowercase				       =			'bert_model.' + key
          lowercase				       =			value
      encoder.load_state_dict(_lowerCamelCase			)
      return model


class 			a	(  a_      ):







    def      UpperCamelCase_ (     self			):
      lowercase				       =			DPRQuestionEncoder(DPRConfig(**BertConfig.get_config_dict('bert-base-uncased'			)[0]			)			)
      print(F'Loading DPR biencoder from {self.src_file}'			)
      lowercase				       =			load_states_from_checkpoint(self.src_file			)
      lowercase       ,	lowercase				       =			model.question_encoder, 'question_model.'
      # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
      lowercase				       =			{'bert_model.embeddings.position_ids': model.question_encoder.bert_model.embeddings.position_ids}
      for key, value in saved_state.model_dict.items():
        if key.startswith(_lowerCamelCase			):
          lowercase				       =			key[len(_lowerCamelCase			) :]
          if not key.startswith('encode_proj.'			):
            lowercase				       =			'bert_model.' + key
          lowercase				       =			value
      encoder.load_state_dict(_lowerCamelCase			)
      return model


class 			a	(  a_      ):







    def      UpperCamelCase_ (     self			):
      lowercase				       =			DPRReader(DPRConfig(**BertConfig.get_config_dict('bert-base-uncased'			)[0]			)			)
      print(F'Loading DPR reader from {self.src_file}'			)
      lowercase				       =			load_states_from_checkpoint(self.src_file			)
      # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
      lowercase				       =			{
          'encoder.bert_model.embeddings.position_ids': model.span_predictor.encoder.bert_model.embeddings.position_ids
      }
      for key, value in saved_state.model_dict.items():
        if key.startswith('encoder.'			) and not key.startswith('encoder.encode_proj'			):
          lowercase				       =			'encoder.bert_model.' + key[len('encoder.'			) :]
        lowercase				       =			value
      model.span_predictor.load_state_dict(_lowerCamelCase			)
      return model


def 		_SCREAMING_SNAKE_CASE					(       __snake_case     :    str						,   __snake_case     :    Path						,   __snake_case     :    Path       ):





  '''simple docstring'''





  lowercase				       =			Path(__snake_case       )
  dest_dir.mkdir(exist_ok=__snake_case       )

  lowercase				       =			DPRState.from_type(__snake_case						,   src_file=__snake_case       )
  lowercase				       =			dpr_state.load_dpr_model()
  model.save_pretrained(__snake_case       )
  model.from_pretrained(__snake_case       )  # sanity check


if __name__ == "__main__":
    _UpperCamelCase						:    Optional[int]		=						argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        '--type', type=str, help='Type of the component to convert: \'ctx_encoder\', \'question_encoder\' or \'reader\'.'
    )
    parser.add_argument(
        '--src',
        type=str,
        help=(
            'Path to the dpr checkpoint file. They can be downloaded from the official DPR repo'
            ' https://github.com/facebookresearch/DPR. Note that in the official repo, both encoders are stored in the'
            ' \'retriever\' checkpoints.'
        ),
    )
    parser.add_argument('--dest', type=str, default=None, help='Path to the output PyTorch model directory.')
    _UpperCamelCase						:    List[Any]		=						parser.parse_args()

    _UpperCamelCase						:    int		=						Path(args.src)
    _UpperCamelCase						:    List[Any]		=						F'''converted-{src_file.name}''' if args.dest is None else args.dest
    _UpperCamelCase						:    Optional[Any]		=						Path(dest_dir)
    assert src_file.exists()
    assert (
        args.type is not None
    ), "Please specify the component type of the DPR model to convert: 'ctx_encoder', 'question_encoder' or 'reader'."
    convert(args.type, src_file, dest_dir)






from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCamelCase			:					List[Any]         =							logging.get_logger(__name__)

UpperCamelCase			:					Tuple         =							{
    '''mnaylor/mega-base-wikitext''': '''https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json''',
}

class      __snake_case				(   _UpperCAmelCase):
					SCREAMING_SNAKE_CASE_		:       str								=			"mega"
					def __init__(						self       :						Optional[int]	,      _lowerCAmelCase       :						List[Any]=3_0_5_2_2	,      _lowerCAmelCase       :						List[str]=1_2_8	,      _lowerCAmelCase       :						Dict=4	,      _lowerCAmelCase       :						Optional[Any]=2_5_6	,      _lowerCAmelCase       :						Optional[Any]=1_6	,      _lowerCAmelCase       :						Dict=True	,      _lowerCAmelCase       :						Optional[Any]=6_4	,      _lowerCAmelCase       :						Optional[Any]=False	,      _lowerCAmelCase       :						Optional[int]=-1	,      _lowerCAmelCase       :						Any=None	,      _lowerCAmelCase       :						List[Any]=True	,      _lowerCAmelCase       :						List[Any]="scalenorm"	,      _lowerCAmelCase       :						str=True	,      _lowerCAmelCase       :						Union[str, Any]="silu"	,      _lowerCAmelCase       :						Optional[Any]="softmax"	,      _lowerCAmelCase       :						Optional[Any]=0.1	,      _lowerCAmelCase       :						Dict=0.1	,      _lowerCAmelCase       :						Any=0.1	,      _lowerCAmelCase       :						Optional[Any]=False	,      _lowerCAmelCase       :						Any=True	,      _lowerCAmelCase       :						int=2_5_6	,      _lowerCAmelCase       :						List[str]=True	,      _lowerCAmelCase       :						str=0.1	,      _lowerCAmelCase       :						List[str]=2_0_4_8	,      _lowerCAmelCase       :						Optional[Any]=False	,      _lowerCAmelCase       :						str=2	,      _lowerCAmelCase       :						Dict=0.02	,      _lowerCAmelCase       :						str=0.2	,      _lowerCAmelCase       :						List[str]=0.02	,      _lowerCAmelCase       :						int=1.0	,      _lowerCAmelCase       :						Optional[int]=1	,      _lowerCAmelCase       :						List[str]=0	,      _lowerCAmelCase       :						str=2	,      _lowerCAmelCase       :						Optional[Any]="rotary"	,      _lowerCAmelCase       :						int=None	,      _lowerCAmelCase       :						int=True	,      _lowerCAmelCase       :						Optional[Any]=True	,      **_lowerCAmelCase       :						str	,      )							->					Any:



											'''simple docstring'''





											super().__init__(pad_token_id=_lowerCAmelCase	,      bos_token_id=_lowerCAmelCase	,      eos_token_id=_lowerCAmelCase	,      **_lowerCAmelCase  )

											__magic_name__					=      vocab_size
											__magic_name__					=      hidden_size
											__magic_name__					=      num_hidden_layers
											__magic_name__					=      activation
											__magic_name__					=      attention_activation
											__magic_name__					=      intermediate_size
											__magic_name__					=      ema_projection_size
											__magic_name__					=      bidirectional
											__magic_name__					=      shared_representation_size
											__magic_name__					=      use_chunking
											__magic_name__					=      chunk_size
											__magic_name__					=      truncation
											__magic_name__					=      normalize_before_mega
											__magic_name__					=      normalization_type
											__magic_name__					=      norm_affine
											__magic_name__					=      dropout_prob
											__magic_name__					=      hidden_dropout_prob
											__magic_name__					=      attention_probs_dropout_prob
											__magic_name__					=      use_feature_dropout
											__magic_name__					=      use_normalized_ffn
											__magic_name__					=      nffn_hidden_size
											__magic_name__					=      normalize_before_ffn
											__magic_name__					=      nffn_activation_dropout_prob
											__magic_name__					=      max_positions
											__magic_name__					=      add_token_type_embeddings
											__magic_name__					=      type_vocab_size
											__magic_name__					=      initializer_range
											__magic_name__					=      ema_delta_alpha_range
											__magic_name__					=      ema_beta_range
											__magic_name__					=      ema_gamma_omega_range
											__magic_name__					=      relative_positional_bias
											__magic_name__					=      use_cache
											__magic_name__					=      classifier_dropout
											__magic_name__					=      add_lm_hidden_dense_layer
											__magic_name__					=      1  # not used but required by Hugging Face







class      __snake_case				(   _UpperCAmelCase):
					@property
					def 	a			(						self       :						int  )							->					Mapping[str, Mapping[int, str]]:



											'''simple docstring'''





											if self.task == "multiple-choice":
																	__magic_name__					=      {0: 'batch', 1: 'choice', 2: 'sequence'}
											else:
																	__magic_name__					=      {0: 'batch', 1: 'sequence'}
											return OrderedDict(
											    [
											        ('input_ids', dynamic_axis),
											        ('attention_mask', dynamic_axis),
											    ]  )

