



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


lowerCAmelCase						:		List[Any]						      = collections.namedtuple(
    '''CheckpointState''', ['''model_dict''', '''optimizer_dict''', '''scheduler_dict''', '''offset''', '''epoch''', '''encoder_params''']
)



def 	__lowerCAmelCase       (				lowerCamelCase :    str       ):






	'''simple docstring'''




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






class        UpperCAmelCase__						:



					def __init__(     self      ,					UpperCamelCase  )  ->      Dict:
						__lowerCAmelCase       				=	src_file



					def 							UpperCAmelCase_			(     self  )  ->      Union[str, Any]:
						raise NotImplementedError



					@staticmethod
					def 							UpperCAmelCase_			(     UpperCamelCase      ,					*UpperCamelCase      ,					**UpperCamelCase  )  ->      "DPRState":
						if comp_type.startswith("c"  ):
							return DPRContextEncoderState(*UpperCamelCase      ,					**UpperCamelCase  )
						if comp_type.startswith("q"  ):
							return DPRQuestionEncoderState(*UpperCamelCase      ,					**UpperCamelCase  )
						if comp_type.startswith("r"  ):
							return DPRReaderState(*UpperCamelCase      ,					**UpperCamelCase  )
						else:
							raise ValueError("Component type must be either 'ctx_encoder', 'question_encoder' or 'reader'."  )






class        UpperCAmelCase__						(		UpperCamelCase__			):



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






class        UpperCAmelCase__						(		UpperCamelCase__			):



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






class        UpperCAmelCase__						(		UpperCamelCase__			):



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



def 	__lowerCAmelCase       (				lowerCamelCase :    str					,		lowerCamelCase :    Path					,		lowerCamelCase :    Path       ):






	'''simple docstring'''




	__lowerCAmelCase       				=	Path(lowerCamelCase       )
	dest_dir.mkdir(exist_ok=lowerCamelCase       )

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


if __name__ == "__main__":
	lowerCAmelCase						:		List[Any]						      = 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.''')
	lowerCAmelCase						:		str						      = parser.parse_args()

	lowerCAmelCase						:		List[str]						      = Path(args.src)
	lowerCAmelCase						:		List[str]						      = f'converted-{src_file.name}' if args.dest is None else args.dest
	lowerCAmelCase						:		Tuple						      = 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 math



def 	__lowerCAmelCase       (				lowerCamelCase :    int       ):






	'''simple docstring'''




	if not isinstance(lowerCamelCase					,		lowerCamelCase       ):
		__lowerCAmelCase       				=	f'''Input value of [number={number}] must be an integer'''
		raise TypeError(lowerCamelCase       )
	if number < 1:
		__lowerCAmelCase       				=	f'''Input value of [number={number}] must be a positive integer'''
		raise ValueError(lowerCamelCase       )
	__lowerCAmelCase       				=	[number]
	while number != 1:
		if number % 2 == 0:
			__lowerCAmelCase       				=	math.floor(math.sqrt(lowerCamelCase       )       )
		else:
			__lowerCAmelCase       				=	math.floor(
			    math.sqrt(lowerCamelCase       ) * math.sqrt(lowerCamelCase       ) * math.sqrt(lowerCamelCase       )       )
		sequence.append(lowerCamelCase       )
	return sequence


if __name__ == "__main__":
	import doctest

	doctest.testmod()