





"""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						:       Dict,					lowerCAmelCase__						:       Path				)      -> List[str]:




										'''simple docstring'''



										_UpperCamelCase     :   Tuple      =       src_file






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




										'''simple docstring'''



										raise NotImplementedError






				@staticmethod
				def 							snake_case      (					lowerCAmelCase__						:       str,					*lowerCAmelCase__						:       Optional[Any],					**lowerCAmelCase__						:       List[str]				)      -> "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						:       Union[str, Any]				)      -> int:




										'''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     :   Tuple      =       load_states_from_checkpoint(self.src_file				)
										_UpperCamelCase  ,				_UpperCamelCase     :   Optional[Any]      =       model.ctx_encoder, '''ctx_model.'''
										# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
										_UpperCamelCase     :   str      =       {'''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     :   Tuple      =       key[len(lowerCAmelCase__				) :]
																						if not key.startswith('''encode_proj.'''				):
																												_UpperCamelCase     :   Union[str, Any]      =       '''bert_model.''' + key
																						_UpperCamelCase     :   Union[str, Any]      =       value
										encoder.load_state_dict(lowerCAmelCase__				)
										return model




class 					_a		( _lowerCAmelCase							):






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




										'''simple docstring'''



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




class 					_a		( _lowerCAmelCase							):






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




										'''simple docstring'''



										_UpperCamelCase     :   str      =       DPRReader(DPRConfig(**BertConfig.get_config_dict('''bert-base-uncased'''				)[0]				)				)
										print(f"""Loading DPR reader from {self.src_file}"""				)
										_UpperCamelCase     :   Tuple      =       load_states_from_checkpoint(self.src_file				)
										# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
										_UpperCamelCase     :   str      =       {
										    '''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     :   int      =       '''encoder.bert_model.''' + key[len('''encoder.'''				) :]
																_UpperCamelCase     :   Any      =       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     :   List[Any]      =       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"""
#
# Sort large text files in a minimum amount of memory
#
import argparse
import os

class        _UpperCamelCase	:


  """simple docstring"""





  snake_case_        =    'block_{0}.dat'



  def __init__( self					: int       ,       snake_case					: int      )	->			List[str]:


         '''simple docstring'''





         __magic_name__			:		Optional[int]        =     filename
         __magic_name__			:		Any        =     []



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


         '''simple docstring'''





         __magic_name__			:		Optional[Any]        =     self.BLOCK_FILENAME_FORMAT.format(snake_case      )
         with open(snake_case       ,       '''w'''      ) as file:
                file.write(snake_case      )
         self.block_filenames.append(snake_case      )



  def 							_UpperCAmelCase      ( self					: str      )	->			List[str]:


         '''simple docstring'''





         return self.block_filenames



  def 							_UpperCAmelCase      ( self					: Union[str, Any]       ,       snake_case					: Dict       ,       snake_case					: int=None      )	->			Union[str, Any]:


         '''simple docstring'''





         __magic_name__			:		List[str]        =     0
         with open(self.filename      ) as file:
                while True:
                       __magic_name__			:		Tuple        =     file.readlines(snake_case      )

                       if lines == []:
                              break

                       if sort_key is None:
                              lines.sort()
                       else:
                              lines.sort(key=snake_case      )

                       self.write_block(''''''.join(snake_case      )       ,       snake_case      )
                       i += 1




  def 							_UpperCAmelCase      ( self					: Tuple      )	->			Optional[int]:


         '''simple docstring'''





         map(os.remove       ,       self.block_filenames      )

class        _UpperCamelCase	:


  """simple docstring"""





  def 							_UpperCAmelCase      ( self					: List[Any]       ,       snake_case					: Dict      )	->			int:


         '''simple docstring'''





         __magic_name__			:		Any        =     -1
         __magic_name__			:		Union[str, Any]        =     None

         for i in range(len(snake_case      )      ):
                if min_str is None or choices[i] < min_str:
                       __magic_name__			:		Dict        =     i

         return min_index

class        _UpperCamelCase	:


  """simple docstring"""





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


         '''simple docstring'''





         __magic_name__			:		Dict        =     files
         __magic_name__			:		List[str]        =     set()
         __magic_name__			:		Any        =     len(snake_case      )
         __magic_name__			:		int        =     {i: None for i in range(self.num_buffers      )}



  def 							_UpperCAmelCase      ( self					: List[str]      )	->			Optional[Any]:


         '''simple docstring'''





         return {
             i: self.buffers[i] for i in range(self.num_buffers      ) if i not in self.empty
         }



  def 							_UpperCAmelCase      ( self					: int      )	->			int:


         '''simple docstring'''





         for i in range(self.num_buffers      ):
                if self.buffers[i] is None and i not in self.empty:
                       __magic_name__			:		int        =     self.files[i].readline()

                       if self.buffers[i] == "":
                              self.empty.add(snake_case      )
                              self.files[i].close()

         if len(self.empty      ) == self.num_buffers:
                return False

         return True




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


         '''simple docstring'''





         __magic_name__			:		List[str]        =     self.buffers[index]
         __magic_name__			:		Dict        =     None

         return value

class        _UpperCamelCase	:


  """simple docstring"""





  def __init__( self					: Any       ,       snake_case					: str      )	->			Dict:


         '''simple docstring'''





         __magic_name__			:		Optional[Any]        =     merge_strategy



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


         '''simple docstring'''





         __magic_name__			:		Union[str, Any]        =     FilesArray(self.get_file_handles(snake_case       ,       snake_case      )      )
         with open(snake_case       ,       '''w'''       ,       snake_case      ) as outfile:
                while buffers.refresh():
                       __magic_name__			:		List[str]        =     self.merge_strategy.select(buffers.get_dict()      )
                       outfile.write(buffers.unshift(snake_case      )      )




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


         '''simple docstring'''





         __magic_name__			:		str        =     {}

         for i in range(len(snake_case      )      ):
                __magic_name__			:		Any        =     open(filenames[i]       ,       '''r'''       ,       snake_case      )  # noqa: UP015

         return files

class        _UpperCamelCase	:


  """simple docstring"""





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


         '''simple docstring'''





         __magic_name__			:		int        =     block_size



  def 							_UpperCAmelCase      ( self					: str       ,       snake_case					: Optional[int]       ,       snake_case					: Dict=None      )	->			Optional[Any]:


         '''simple docstring'''





         __magic_name__			:		Tuple        =     self.get_number_blocks(snake_case       ,       self.block_size      )
         __magic_name__			:		str        =     FileSplitter(snake_case      )
         splitter.split(self.block_size       ,       snake_case      )

         __magic_name__			:		int        =     FileMerger(NWayMerge()      )
         __magic_name__			:		Any        =     self.block_size / (num_blocks + 1)
         merger.merge(splitter.get_block_filenames()       ,       filename + '''.out'''       ,       snake_case      )

         splitter.cleanup()




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


         '''simple docstring'''





         return (os.stat(snake_case      ).st_size / block_size) + 1







def       UpperCamelCase_					(     lowerCamelCase							:					Union[str, Any]   )						->     Union[str, Any]:



       """simple docstring"""







       if string[-1].lower() == "k":
              return int(string[:-1]   ) * 1024
       elif string[-1].lower() == "m":
              return int(string[:-1]   ) * 1024 * 1024
       elif string[-1].lower() == "g":
              return int(string[:-1]   ) * 1024 * 1024 * 1024
       else:
              return int(lowerCamelCase   )







def       UpperCamelCase_					(     )						->     Optional[Any]:



       """simple docstring"""







       __magic_name__			:		Tuple        =     argparse.ArgumentParser()
       parser.add_argument(
           '''-m'''     , '''--mem'''     , help='''amount of memory to use for sorting'''     , default='''100M'''   )
       parser.add_argument(
           '''filename'''     , metavar='''<filename>'''     , nargs=1     , help='''name of file to sort'''   )
       __magic_name__			:		List[str]        =     parser.parse_args()

       __magic_name__			:		Any        =     ExternalSort(parse_memory(args.mem   )   )
       sorter.sort(args.filename[0]   )


if __name__ == "__main__":
 main()
