





"""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_           =collections.namedtuple(
    """CheckpointState""", ["""model_dict""", """optimizer_dict""", """scheduler_dict""", """offset""", """epoch""", """encoder_params"""]
)





def        a_				(							_lowercase   ):
						print(F"""Reading saved model from {model_file}"""   )
						_UpperCamelCase			:					Optional[Any]     =    torch.load(_lowercase       ,						map_location=lambda _lowercase       ,						_lowercase   : default_restore_location(_lowercase       ,						'''cpu'''   )   )
						return CheckpointState(**_lowercase   )




class  _a						:






				def __init__( self		:    str, lowerCAmelCase__		:    Path							)	->      Dict:




										'''simple docstring'''



										_UpperCamelCase			:					Union[str, Any]     =    src_file






				def 				snake_case	( self		:    Dict							)	->      Optional[Any]:




										'''simple docstring'''



										raise NotImplementedError






				@staticmethod
				def 				snake_case	( lowerCAmelCase__		:    str, *lowerCAmelCase__		:    Tuple, **lowerCAmelCase__		:    Any							)	->      "DPRState":




										'''simple docstring'''



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






				def 				snake_case	( self		:    List[Any]							)	->      Any:




										'''simple docstring'''



										_UpperCamelCase			:					Optional[int]     =    DPRContextEncoder(DPRConfig(**BertConfig.get_config_dict('''bert-base-uncased'''							)[0]							)							)
										print(f"""Loading DPR biencoder from {self.src_file}"""							)
										_UpperCamelCase			:					Optional[int]     =    load_states_from_checkpoint(self.src_file							)
										_UpperCamelCase   ,							_UpperCamelCase			:					Optional[int]     =    model.ctx_encoder, '''ctx_model.'''
										# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
										_UpperCamelCase			:					Any     =    {'''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__							):
																						_UpperCamelCase			:					Optional[Any]     =    key[len(lowerCAmelCase__							) :]
																						if not key.startswith('''encode_proj.'''							):
																												_UpperCamelCase			:					Any     =    '''bert_model.''' + key
																						_UpperCamelCase			:					str     =    value
										encoder.load_state_dict(lowerCAmelCase__							)
										return model




class  _a						(      _lowerCAmelCase				):






				def 				snake_case	( self		:    List[Any]							)	->      Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase			:					Union[str, Any]     =    DPRQuestionEncoder(DPRConfig(**BertConfig.get_config_dict('''bert-base-uncased'''							)[0]							)							)
										print(f"""Loading DPR biencoder from {self.src_file}"""							)
										_UpperCamelCase			:					Optional[Any]     =    load_states_from_checkpoint(self.src_file							)
										_UpperCamelCase   ,							_UpperCamelCase			:					str     =    model.question_encoder, '''question_model.'''
										# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
										_UpperCamelCase			:					List[str]     =    {'''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__							):
																						_UpperCamelCase			:					Optional[Any]     =    key[len(lowerCAmelCase__							) :]
																						if not key.startswith('''encode_proj.'''							):
																												_UpperCamelCase			:					List[Any]     =    '''bert_model.''' + key
																						_UpperCamelCase			:					int     =    value
										encoder.load_state_dict(lowerCAmelCase__							)
										return model




class  _a						(      _lowerCAmelCase				):






				def 				snake_case	( self		:    Tuple							)	->      List[str]:




										'''simple docstring'''



										_UpperCamelCase			:					int     =    DPRReader(DPRConfig(**BertConfig.get_config_dict('''bert-base-uncased'''							)[0]							)							)
										print(f"""Loading DPR reader from {self.src_file}"""							)
										_UpperCamelCase			:					List[str]     =    load_states_from_checkpoint(self.src_file							)
										# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
										_UpperCamelCase			:					Union[str, Any]     =    {
										    '''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'''							):
																						_UpperCamelCase			:					Optional[int]     =    '''encoder.bert_model.''' + key[len('''encoder.'''							) :]
																_UpperCamelCase			:					int     =    value
										model.span_predictor.load_state_dict(lowerCAmelCase__							)
										return model





def        a_				(							_lowercase       ,						_lowercase       ,						_lowercase   ):
						_UpperCamelCase			:					Dict     =    Path(_lowercase   )
						dest_dir.mkdir(exist_ok=_lowercase   )

						_UpperCamelCase			:					int     =    DPRState.from_type(_lowercase       ,						src_file=_lowercase   )
						_UpperCamelCase			:					int     =    dpr_state.load_dpr_model()
						model.save_pretrained(_lowercase   )
						model.from_pretrained(_lowercase   )  # sanity check


if __name__ == "__main__":
						UpperCamelCase_           =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_           =parser.parse_args()

						UpperCamelCase_           =Path(args.src)
						UpperCamelCase_           =F"converted-{src_file.name}" if args.dest is None else args.dest
						UpperCamelCase_           =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)










"""simple docstring"""




import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
						from transformers.models.electra.modeling_flax_electra import (
						    FlaxElectraForCausalLM,
						    FlaxElectraForMaskedLM,
						    FlaxElectraForMultipleChoice,
						    FlaxElectraForPreTraining,
						    FlaxElectraForQuestionAnswering,
						    FlaxElectraForSequenceClassification,
						    FlaxElectraForTokenClassification,
						    FlaxElectraModel,
						)




class  _a						(      unittest.TestCase				):






				def __init__( self		:    str, lowerCAmelCase__		:    Tuple, lowerCAmelCase__		:    List[str]=1_3, lowerCAmelCase__		:    Optional[int]=7, lowerCAmelCase__		:    Optional[Any]=True, lowerCAmelCase__		:    Optional[int]=True, lowerCAmelCase__		:    Any=True, lowerCAmelCase__		:    str=True, lowerCAmelCase__		:    Any=9_9, lowerCAmelCase__		:    Union[str, Any]=2_4, lowerCAmelCase__		:    Optional[Any]=3_2, lowerCAmelCase__		:    Tuple=5, lowerCAmelCase__		:    int=4, lowerCAmelCase__		:    List[Any]=3_7, lowerCAmelCase__		:    Dict="gelu", lowerCAmelCase__		:    Any=0.1, lowerCAmelCase__		:    Any=0.1, lowerCAmelCase__		:    str=5_1_2, lowerCAmelCase__		:    Tuple=1_6, lowerCAmelCase__		:    int=2, lowerCAmelCase__		:    int=0.02, lowerCAmelCase__		:    Optional[int]=4, )	->      Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase			:					Dict     =    parent
										_UpperCamelCase			:					Tuple     =    batch_size
										_UpperCamelCase			:					str     =    seq_length
										_UpperCamelCase			:					Any     =    is_training
										_UpperCamelCase			:					Union[str, Any]     =    use_attention_mask
										_UpperCamelCase			:					List[Any]     =    use_token_type_ids
										_UpperCamelCase			:					str     =    use_labels
										_UpperCamelCase			:					Union[str, Any]     =    vocab_size
										_UpperCamelCase			:					int     =    hidden_size
										_UpperCamelCase			:					Optional[int]     =    embedding_size
										_UpperCamelCase			:					int     =    num_hidden_layers
										_UpperCamelCase			:					int     =    num_attention_heads
										_UpperCamelCase			:					Tuple     =    intermediate_size
										_UpperCamelCase			:					Optional[int]     =    hidden_act
										_UpperCamelCase			:					Any     =    hidden_dropout_prob
										_UpperCamelCase			:					Union[str, Any]     =    attention_probs_dropout_prob
										_UpperCamelCase			:					Any     =    max_position_embeddings
										_UpperCamelCase			:					List[str]     =    type_vocab_size
										_UpperCamelCase			:					Union[str, Any]     =    type_sequence_label_size
										_UpperCamelCase			:					str     =    initializer_range
										_UpperCamelCase			:					Union[str, Any]     =    num_choices






				def 				snake_case	( self		:    Optional[Any]							)	->      Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase			:					Optional[Any]     =    ids_tensor([self.batch_size, self.seq_length], self.vocab_size							)

										_UpperCamelCase			:					Dict     =    None
										if self.use_attention_mask:
																_UpperCamelCase			:					Union[str, Any]     =    random_attention_mask([self.batch_size, self.seq_length]							)

										_UpperCamelCase			:					Optional[Any]     =    None
										if self.use_token_type_ids:
																_UpperCamelCase			:					Union[str, Any]     =    ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size							)

										_UpperCamelCase			:					Optional[Any]     =    ElectraConfig(
										    vocab_size=self.vocab_size, hidden_size=self.hidden_size, embedding_size=self.embedding_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, max_position_embeddings=self.max_position_embeddings, type_vocab_size=self.type_vocab_size, initializer_range=self.initializer_range, )

										return config, input_ids, token_type_ids, attention_mask






				def 				snake_case	( self		:    Any							)	->      Optional[int]:




										'''simple docstring'''



										_UpperCamelCase			:					Dict     =    self.prepare_config_and_inputs()
										_UpperCamelCase   ,							_UpperCamelCase   ,							_UpperCamelCase   ,							_UpperCamelCase			:					Tuple     =    config_and_inputs
										_UpperCamelCase			:					List[Any]     =    {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
										return config, inputs_dict




@require_flax
class  _a						(      _lowerCAmelCase	,				unittest.TestCase				):
				UpperCamelCase								=    True

				UpperCamelCase								=    (
				    (
				        FlaxElectraModel,
				        FlaxElectraForCausalLM,
				        FlaxElectraForMaskedLM,
				        FlaxElectraForPreTraining,
				        FlaxElectraForTokenClassification,
				        FlaxElectraForQuestionAnswering,
				        FlaxElectraForMultipleChoice,
				        FlaxElectraForSequenceClassification,
				    )
				    if is_flax_available()
				    else ()
				)






				def 				snake_case	( self		:    Union[str, Any]							)	->      Optional[Any]:




										'''simple docstring'''



										_UpperCamelCase			:					List[Any]     =    FlaxElectraModelTester(self							)






				@slow
				def 				snake_case	( self		:    Optional[Any]							)	->      str:




										'''simple docstring'''



										for model_class_name in self.all_model_classes:
																if model_class_name == FlaxElectraForMaskedLM:
																						_UpperCamelCase			:					int     =    model_class_name.from_pretrained('''google/electra-small-generator'''							)
																else:
																						_UpperCamelCase			:					Any     =    model_class_name.from_pretrained('''google/electra-small-discriminator'''							)
																_UpperCamelCase			:					List[str]     =    model(np.ones((1, 1)							)							)
																self.assertIsNotNone(lowerCAmelCase__							)



