



import json
import os
import tempfile

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features import ArrayaD
from utils import generate_examples, get_duration


_lowerCamelCase         =	(30, 487)
_lowerCamelCase         =	(36, 1024)
_lowerCamelCase         =	(100, 100)
_lowerCamelCase         =	100

_lowerCamelCase         =	datasets.Features(
    {"""text""": ArrayaD(SHAPE_TEST_1, dtype="""float32"""), """image""": ArrayaD(SHAPE_TEST_2, dtype="""float32""")}
)

_lowerCamelCase     ,							_lowerCamelCase         =	os.path.split(__file__)
_lowerCamelCase         =	os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))






@get_duration
def 						_lowerCAmelCase    (					__lowerCamelCase :  List[str]		,						__lowerCamelCase :  str		,						__lowerCamelCase :  List[Any]  ):



 """simple docstring"""





 with ArrowWriter(features=__lowerCamelCase		,						path=os.path.join(__lowerCamelCase		,						"beta.arrow"  )  ) as writer:
  for key, record in dummy_data:
   __SCREAMING_SNAKE_CASE							:     List[str]							=						my_features.encode_example(__lowerCamelCase  )
   writer.write(__lowerCamelCase  )
  __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE							:     int							=						writer.finalize()






@get_duration
def 						_lowerCAmelCase    (					__lowerCamelCase :  int		,						__lowerCamelCase :  Optional[int]  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[Any]							=						datasets.Dataset.from_file(
     filename=os.path.join(__lowerCamelCase		,						"beta.arrow"  )		,						info=datasets.DatasetInfo(features=__lowerCamelCase  )  )
 for _ in dataset:
  pass






@get_duration
def 						_lowerCAmelCase    (					__lowerCamelCase :  List[Any]		,						__lowerCamelCase :  Any  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     Tuple							=						datasets.Dataset.from_file(
     filename=os.path.join(__lowerCamelCase		,						"beta.arrow"  )		,						info=datasets.DatasetInfo(features=__lowerCamelCase  )  )
 dataset.set_format("numpy"  )
 for _ in dataset:
  pass






@get_duration
def 						_lowerCAmelCase    (					__lowerCamelCase :  Optional[int]		,						__lowerCamelCase :  Optional[Any]  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[Any]							=						10
 __SCREAMING_SNAKE_CASE							:     Tuple							=						datasets.Dataset.from_file(
     filename=os.path.join(__lowerCamelCase		,						"beta.arrow"  )		,						info=datasets.DatasetInfo(features=__lowerCamelCase  )  )
 for i in range(0		,						len(__lowerCamelCase  )		,						__lowerCamelCase  ):
  __SCREAMING_SNAKE_CASE							:     Tuple							=						dataset[i : i + batch_size]






@get_duration
def 						_lowerCAmelCase    (					__lowerCamelCase :  List[Any]		,						__lowerCamelCase :  Any  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[Any]							=						10
 __SCREAMING_SNAKE_CASE							:     Any							=						datasets.Dataset.from_file(
     filename=os.path.join(__lowerCamelCase		,						"beta.arrow"  )		,						info=datasets.DatasetInfo(features=__lowerCamelCase  )  )
 dataset.set_format("numpy"  )
 for i in range(0		,						len(__lowerCamelCase  )		,						__lowerCamelCase  ):
  __SCREAMING_SNAKE_CASE							:     List[Any]							=						dataset[i : i + batch_size]






@get_duration
def 						_lowerCAmelCase    (					__lowerCamelCase :  Any		,						__lowerCamelCase :  str  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[str]							=						datasets.Dataset.from_file(
     filename=os.path.join(__lowerCamelCase		,						"beta.arrow"  )		,						info=datasets.DatasetInfo(features=__lowerCamelCase  )  )
 for col in feats:
  __SCREAMING_SNAKE_CASE							:     Dict							=						dataset[col]






@get_duration
def 						_lowerCAmelCase    (					__lowerCamelCase :  Dict		,						__lowerCamelCase :  Dict  ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[Any]							=						datasets.Dataset.from_file(
     filename=os.path.join(__lowerCamelCase		,						"beta.arrow"  )		,						info=datasets.DatasetInfo(features=__lowerCamelCase  )  )
 dataset.set_format("numpy"  )
 for col in feats:
  __SCREAMING_SNAKE_CASE							:     List[Any]							=						dataset[col]






def 						_lowerCAmelCase    (					):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE							:     List[str]							=						{}
 __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						(
     read_unformated,
     read_formatted_as_numpy,
     read_batch_unformated,
     read_batch_formatted_as_numpy,
     read_col_unformated,
     read_col_formatted_as_numpy,
 )
 with tempfile.TemporaryDirectory() as tmp_dir:
  __SCREAMING_SNAKE_CASE							:     Optional[Any]							=						datasets.Features({"image": ArrayaD(__lowerCamelCase		,						dtype="float32"  )}  )
  __SCREAMING_SNAKE_CASE							:     Union[str, Any]							=						generate_examples(features=__lowerCamelCase		,						num_examples=__lowerCamelCase  )
  __SCREAMING_SNAKE_CASE							:     Tuple							=						write(__lowerCamelCase		,						__lowerCamelCase		,						__lowerCamelCase  )
  for read_func in read_functions:
   __SCREAMING_SNAKE_CASE							:     Tuple							=						read_func(__lowerCamelCase		,						__lowerCamelCase  )

 with tempfile.TemporaryDirectory() as tmp_dir:
  # don't use fixed length for fair comparison
  # feats = datasets.Features(
  #     {"image": datasets.Sequence(datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[1]), SPEED_TEST_SHAPE[0])}
  # )
  __SCREAMING_SNAKE_CASE							:     Dict							=						datasets.Features({"image": datasets.Sequence(datasets.Sequence(datasets.Value("float32"  )  )  )}  )
  __SCREAMING_SNAKE_CASE							:     str							=						generate_examples(
      features=__lowerCamelCase		,						num_examples=__lowerCamelCase		,						seq_shapes={"image": SPEED_TEST_SHAPE}  )
  __SCREAMING_SNAKE_CASE							:     List[Any]							=						write(__lowerCamelCase		,						__lowerCamelCase		,						__lowerCamelCase  )
  for read_func in read_functions:
   __SCREAMING_SNAKE_CASE							:     int							=						read_func(__lowerCamelCase		,						__lowerCamelCase  )

 with tempfile.TemporaryDirectory() as tmp_dir:
  # don't use fixed length for fair comparison
  # feats = datasets.Features(
  #     {"image": datasets.Sequence(datasets.Value("float32"), SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1])}
  # )
  __SCREAMING_SNAKE_CASE							:     Tuple							=						datasets.Features({"image": datasets.Sequence(datasets.Value("float32"  )  )}  )
  __SCREAMING_SNAKE_CASE							:     int							=						generate_examples(
      features=__lowerCamelCase		,						num_examples=__lowerCamelCase		,						seq_shapes={"image": [SPEED_TEST_SHAPE[0] * SPEED_TEST_SHAPE[1]]}		,						)
  __SCREAMING_SNAKE_CASE							:     Tuple							=						write(__lowerCamelCase		,						__lowerCamelCase		,						__lowerCamelCase  )
  for read_func in read_functions:
   __SCREAMING_SNAKE_CASE							:     int							=						read_func(__lowerCamelCase		,						__lowerCamelCase  )

 with open(__lowerCamelCase		,						"wb"  ) as f:
  f.write(json.dumps(__lowerCamelCase  ).encode("utf-8"  )  )


if __name__ == "__main__":  # useful to run the profiler
      benchmark_array_xd()










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

def       _a	(					SCREAMING_SNAKE_CASE						:					str  ):







 """simple docstring"""



 print(F"Reading saved model from {model_file}"  )
 UpperCamelCase__       :     int 							= torch.load(SCREAMING_SNAKE_CASE     ,			map_location=lambda SCREAMING_SNAKE_CASE     ,			SCREAMING_SNAKE_CASE  : default_restore_location(SCREAMING_SNAKE_CASE     ,			'''cpu'''  )  )
 return CheckpointState(**SCREAMING_SNAKE_CASE  )

class    __magic_name__						:


     def __init__(   self				:  str				,						lowerCamelCase__				:  Path      )	->   Tuple:





      '''simple docstring'''






      UpperCamelCase__       :     List[str] 							= src_file


     def   UpperCAmelCase__						(   self				:  Dict      )	->   int:





      '''simple docstring'''






      raise NotImplementedError


     @staticmethod
     def   UpperCAmelCase__						(   lowerCamelCase__				:  str				,						*lowerCamelCase__				:  Union[str, Any]				,						**lowerCamelCase__				:  Optional[int]      )	->   "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    __magic_name__						(  __lowerCAmelCase):


     def   UpperCAmelCase__						(   self				:  List[str]      )	->   Dict:





      '''simple docstring'''






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



class    __magic_name__						(  __lowerCAmelCase):


     def   UpperCAmelCase__						(   self				:  Any      )	->   Dict:





      '''simple docstring'''






      UpperCamelCase__       :     List[Any] 							= DPRQuestionEncoder(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__       :     int 							= model.question_encoder, '''question_model.'''
      # Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
      UpperCamelCase__       :     Optional[int] 							= {'''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__       :     int 							= key[len(lowerCamelCase__      ) :]
        if not key.startswith('''encode_proj.'''      ):
         UpperCamelCase__       :     Union[str, Any] 							= '''bert_model.''' + key
        UpperCamelCase__       :     str 							= value
      encoder.load_state_dict(lowerCamelCase__      )
      return model



class    __magic_name__						(  __lowerCAmelCase):


     def   UpperCAmelCase__						(   self				:  int      )	->   Any:





      '''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__       :     Tuple 							= {
          '''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__       :     List[str] 							= '''encoder.bert_model.''' + key[len('''encoder.'''      ) :]
       UpperCamelCase__       :     List[Any] 							= value
      model.span_predictor.load_state_dict(lowerCamelCase__      )
      return model



def       _a	(					SCREAMING_SNAKE_CASE						:					str     ,			SCREAMING_SNAKE_CASE						:					Path     ,			SCREAMING_SNAKE_CASE						:					Path  ):







 """simple docstring"""



 UpperCamelCase__       :     Any 							= Path(SCREAMING_SNAKE_CASE  )
 dest_dir.mkdir(exist_ok=SCREAMING_SNAKE_CASE  )

 UpperCamelCase__       :     str 							= DPRState.from_type(SCREAMING_SNAKE_CASE     ,			src_file=SCREAMING_SNAKE_CASE  )
 UpperCamelCase__       :     str 							= dpr_state.load_dpr_model()
 model.save_pretrained(SCREAMING_SNAKE_CASE  )
 model.from_pretrained(SCREAMING_SNAKE_CASE  )  # sanity check


if __name__ == "__main__":
       __UpperCamelCase							:							Optional[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.")
       __UpperCamelCase							:							Union[str, Any]    							=					parser.parse_args()

       __UpperCamelCase							:							Optional[int]    							=					Path(args.src)
       __UpperCamelCase							:							Dict    							=					f"converted-{src_file.name}" if args.dest is None else args.dest
       __UpperCamelCase							:							str    							=					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)


