


'''simple docstring'''




import collections
from typing import List, Optional, Union

from ...tokenization_utils_base import BatchEncoding
from ...utils import TensorType, add_end_docstrings, add_start_docstrings, logging
from ..bert.tokenization_bert_fast import BertTokenizerFast
from .tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer, DPRReaderTokenizer


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

UpperCAmelCase_      :		Any          =		{"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

UpperCAmelCase_      :		List[Any]          =		{
    "vocab_file": {
        "facebook/dpr-ctx_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-ctx_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-ctx_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-ctx_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}
UpperCAmelCase_      :		List[Any]          =		{
    "vocab_file": {
        "facebook/dpr-question_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-question_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-question_encoder-single-nq-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-question_encoder-multiset-base": (
            "https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}
UpperCAmelCase_      :		Optional[Any]          =		{
    "vocab_file": {
        "facebook/dpr-reader-single-nq-base": (
            "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/vocab.txt"
        ),
        "facebook/dpr-reader-multiset-base": (
            "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "facebook/dpr-reader-single-nq-base": (
            "https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/tokenizer.json"
        ),
        "facebook/dpr-reader-multiset-base": (
            "https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/tokenizer.json"
        ),
    },
}

UpperCAmelCase_      :		str          =		{
    "facebook/dpr-ctx_encoder-single-nq-base": 512,
    "facebook/dpr-ctx_encoder-multiset-base": 512,
}
UpperCAmelCase_      :		Tuple          =		{
    "facebook/dpr-question_encoder-single-nq-base": 512,
    "facebook/dpr-question_encoder-multiset-base": 512,
}
UpperCAmelCase_      :		Union[str, Any]          =		{
    "facebook/dpr-reader-single-nq-base": 512,
    "facebook/dpr-reader-multiset-base": 512,
}


UpperCAmelCase_      :		List[str]          =		{
    "facebook/dpr-ctx_encoder-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-ctx_encoder-multiset-base": {"do_lower_case": True},
}
UpperCAmelCase_      :		Optional[Any]          =		{
    "facebook/dpr-question_encoder-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-question_encoder-multiset-base": {"do_lower_case": True},
}
UpperCAmelCase_      :		int          =		{
    "facebook/dpr-reader-single-nq-base": {"do_lower_case": True},
    "facebook/dpr-reader-multiset-base": {"do_lower_case": True},
}







class 					a			(				snake_case__  ):
 '''simple docstring'''




 __lowerCAmelCase						: Union[str, Any]							  =		VOCAB_FILES_NAMES
 __lowerCAmelCase						: str							  =		CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase						: str							  =		CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase						: List[Any]							  =		CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION
 __lowerCAmelCase						: Optional[Any]							  =		DPRContextEncoderTokenizer







class 					a			(				snake_case__  ):
 '''simple docstring'''




 __lowerCAmelCase						: Optional[int]							  =		VOCAB_FILES_NAMES
 __lowerCAmelCase						: str							  =		QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase						: Optional[Any]							  =		QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase						: List[str]							  =		QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION
 __lowerCAmelCase						: Union[str, Any]							  =		DPRQuestionEncoderTokenizer





UpperCAmelCase_      :		List[Any]          =		collections.namedtuple(
    "DPRSpanPrediction", ["span_score", "relevance_score", "doc_id", "start_index", "end_index", "text"]
)

UpperCAmelCase_      :		Optional[Any]          =		collections.namedtuple("DPRReaderOutput", ["start_logits", "end_logits", "relevance_logits"])


UpperCAmelCase_      :		List[str]          =		R"\n    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.\n    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),\n    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`\n    with the format:\n\n    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>\n\n    Args:\n        questions (`str` or `List[str]`):\n            The questions to be encoded. You can specify one question for many passages. In this case, the question\n            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in\n            `titles` or `texts`.\n        titles (`str` or `List[str]`):\n            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.\n        texts (`str` or `List[str]`):\n            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.\n        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):\n            Activates and controls padding. Accepts the following values:\n\n            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence\n              if provided).\n            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided.\n            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different\n              lengths).\n        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):\n            Activates and controls truncation. Accepts the following values:\n\n            - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to\n              the maximum acceptable input length for the model if that argument is not provided. This will truncate\n              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch\n              of pairs) is provided.\n            - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the first\n              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum\n              acceptable input length for the model if that argument is not provided. This will only truncate the\n              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.\n            - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths\n              greater than the model maximum admissible input size).\n        max_length (`int`, *optional*):\n                Controls the maximum length to use by one of the truncation/padding parameters.\n\n                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length\n                is required by one of the truncation/padding parameters. If the model has no specific maximum input\n                length (like XLNet) truncation/padding to a maximum length will be deactivated.\n        return_tensors (`str` or [`~utils.TensorType`], *optional*):\n                If set, will return tensors instead of list of python integers. Acceptable values are:\n\n                - `'tf'`: Return TensorFlow `tf.constant` objects.\n                - `'pt'`: Return PyTorch `torch.Tensor` objects.\n                - `'np'`: Return Numpy `np.ndarray` objects.\n        return_attention_mask (`bool`, *optional*):\n            Whether or not to return the attention mask. If not set, will return the attention mask according to the\n            specific tokenizer's default, defined by the `return_outputs` attribute.\n\n            [What are attention masks?](../glossary#attention-mask)\n\n    Return:\n        `Dict[str, List[List[int]]]`: A dictionary with the following keys:\n\n        - `input_ids`: List of token ids to be fed to a model.\n        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.\n    "







@add_start_docstrings(snake_case__  )
class 					a			:
 '''simple docstring'''




 def __call__(    self		,      lowerCamelCase_		,      lowerCamelCase_ = None		,      lowerCamelCase_ = None		,      lowerCamelCase_ = False		,      lowerCamelCase_ = False		,      lowerCamelCase_ = None		,      lowerCamelCase_ = None		,      lowerCamelCase_ = None		,      **lowerCamelCase_		,      )      ->    BatchEncoding:
    if titles is None and texts is None:
       return super().__call__(
           lowerCamelCase_		,      padding=lowerCamelCase_		,      truncation=lowerCamelCase_		,      max_length=lowerCamelCase_		,      return_tensors=lowerCamelCase_		,      return_attention_mask=lowerCamelCase_		,      **lowerCamelCase_		,      )
    elif titles is None or texts is None:
       _a					:  str						=  titles if texts is None else texts
       return super().__call__(
           lowerCamelCase_		,      lowerCamelCase_		,      padding=lowerCamelCase_		,      truncation=lowerCamelCase_		,      max_length=lowerCamelCase_		,      return_tensors=lowerCamelCase_		,      return_attention_mask=lowerCamelCase_		,      **lowerCamelCase_		,      )
    _a					:  Optional[int]						=  titles if not isinstance(lowerCamelCase_		,      lowerCamelCase_      ) else [titles]
    _a					:  Dict						=  texts if not isinstance(lowerCamelCase_		,      lowerCamelCase_      ) else [texts]
    _a					:  Optional[Any]						=  len(lowerCamelCase_      )
    _a					:  Optional[Any]						=  questions if not isinstance(lowerCamelCase_		,      lowerCamelCase_      ) else [questions] * n_passages
    assert len(lowerCamelCase_      ) == len(
        lowerCamelCase_      ), F'''There should be as many titles than texts but got {len(lowerCamelCase_      )} titles and {len(lowerCamelCase_      )} texts.'''
    _a					:  Optional[int]						=  super().__call__(lowerCamelCase_		,      lowerCamelCase_		,      padding=lowerCamelCase_		,      truncation=lowerCamelCase_      )['input_ids']
    _a					:  List[Any]						=  super().__call__(lowerCamelCase_		,      add_special_tokens=lowerCamelCase_		,      padding=lowerCamelCase_		,      truncation=lowerCamelCase_      )['input_ids']
    _a					:  List[str]						=  {
        'input_ids': [
            (encoded_question_and_title + encoded_text)[:max_length]
            if max_length is not None and truncation
            else encoded_question_and_title + encoded_text
            for encoded_question_and_title, encoded_text in zip(lowerCamelCase_		,      lowerCamelCase_      )
        ]
    }
    if return_attention_mask is not False:
       _a					:  List[Any]						=  []
       for input_ids in encoded_inputs["input_ids"]:
          attention_mask.append([int(input_id != self.pad_token_id      ) for input_id in input_ids]      )
       _a					:  Union[str, Any]						=  attention_mask
    return self.pad(lowerCamelCase_		,      padding=lowerCamelCase_		,      max_length=lowerCamelCase_		,      return_tensors=lowerCamelCase_      )




 def   __UpperCamelCase   (    self		,      lowerCamelCase_		,      lowerCamelCase_		,      lowerCamelCase_ = 1_6		,      lowerCamelCase_ = 6_4		,      lowerCamelCase_ = 4		,      )      ->    List[DPRSpanPrediction]:
    _a					:  int						=  reader_input['input_ids']
    _a		,						_a		,						_a					:  Optional[Any]						=  reader_output[:3]
    _a					:  Tuple						=  len(lowerCamelCase_      )
    _a					:  Dict						=  sorted(range(lowerCamelCase_      )		,      reverse=lowerCamelCase_		,      key=relevance_logits.__getitem__      )
    _a					:  List[DPRReaderOutput]						=  []
    for doc_id in sorted_docs:
       _a					:  Tuple						=  list(input_ids[doc_id]      )
       # assuming question & title information is at the beginning of the sequence
       _a					:  List[Any]						=  sequence_ids.index(self.sep_token_id		,      2      ) + 1  # second sep id
       if sequence_ids[-1] == self.pad_token_id:
          _a					:  List[str]						=  sequence_ids.index(self.pad_token_id      )
       else:
          _a					:  Dict						=  len(lowerCamelCase_      )

       _a					:  Tuple						=  self._get_best_spans(
           start_logits=start_logits[doc_id][passage_offset:sequence_len]		,      end_logits=end_logits[doc_id][passage_offset:sequence_len]		,      max_answer_length=lowerCamelCase_		,      top_spans=lowerCamelCase_		,      )
       for start_index, end_index in best_spans:
          start_index += passage_offset
          end_index += passage_offset
          nbest_spans_predictions.append(
              DPRSpanPrediction(
                  span_score=start_logits[doc_id][start_index] + end_logits[doc_id][end_index]		,      relevance_score=relevance_logits[doc_id]		,      doc_id=lowerCamelCase_		,      start_index=lowerCamelCase_		,      end_index=lowerCamelCase_		,      text=self.decode(sequence_ids[start_index : end_index + 1]      )		,      )      )
       if len(lowerCamelCase_      ) >= num_spans:
          break
    return nbest_spans_predictions[:num_spans]




 def   __UpperCamelCase   (    self		,      lowerCamelCase_		,      lowerCamelCase_		,      lowerCamelCase_		,      lowerCamelCase_		,      )      ->    List[DPRSpanPrediction]:
    _a					:  Tuple						=  []
    for start_index, start_score in enumerate(lowerCamelCase_      ):
       for answer_length, end_score in enumerate(end_logits[start_index : start_index + max_answer_length]      ):
          scores.append(((start_index, start_index + answer_length), start_score + end_score)      )
    _a					:  List[Any]						=  sorted(lowerCamelCase_		,      key=lambda lowerCamelCase_      : x[1]		,      reverse=lowerCamelCase_      )
    _a					:  str						=  []
    for (start_index, end_index), score in scores:
       assert start_index <= end_index, F'''Wrong span indices: [{start_index}:{end_index}]'''
       _a					:  int						=  end_index - start_index + 1
       assert length <= max_answer_length, F'''Span is too long: {length} > {max_answer_length}'''
       if any(
           start_index <= prev_start_index <= prev_end_index <= end_index
           or prev_start_index <= start_index <= end_index <= prev_end_index
           for (prev_start_index, prev_end_index) in chosen_span_intervals      ):
          continue
       chosen_span_intervals.append((start_index, end_index)      )

       if len(lowerCamelCase_      ) == top_spans:
          break
    return chosen_span_intervals







@add_end_docstrings(snake_case__  )
class 					a			(				snake_case__			,							snake_case__  ):
 '''simple docstring'''




 __lowerCAmelCase						: Any							  =		VOCAB_FILES_NAMES
 __lowerCAmelCase						: Dict							  =		READER_PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase						: Tuple							  =		READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase						: Tuple							  =		READER_PRETRAINED_INIT_CONFIGURATION
 __lowerCAmelCase						: int							  =		["""input_ids""", """attention_mask"""]
 __lowerCAmelCase						: int							  =		DPRReaderTokenizer






'''simple docstring'''




import argparse

from torch import nn

# transformers_old should correspond to branch `save_old_prophetnet_model_structure` here
# original prophetnet_checkpoints are saved under `patrickvonplaten/..._old` respectively
from transformers_old.modeling_prophetnet import (
    ProphetNetForConditionalGeneration as ProphetNetForConditionalGenerationOld,
)
from transformers_old.modeling_xlm_prophetnet import (
    XLMProphetNetForConditionalGeneration as XLMProphetNetForConditionalGenerationOld,
)

from transformers import ProphetNetForConditionalGeneration, XLMProphetNetForConditionalGeneration, logging


UpperCAmelCase_      :		List[str]          =		logging.get_logger(__name__)
logging.set_verbosity_info()


def    UpperCAmelCase_  (       A		,	A      ):
   '''simple docstring'''





   if "xprophetnet" in prophetnet_checkpoint_path:
      _a					:  str						=  XLMProphetNetForConditionalGenerationOld.from_pretrained(A      )
      _a		,						_a					:  Any						=  XLMProphetNetForConditionalGeneration.from_pretrained(
          A		,	output_loading_info=A      )
   else:
      _a					:  Optional[Any]						=  ProphetNetForConditionalGenerationOld.from_pretrained(A      )
      _a		,						_a					:  Dict						=  ProphetNetForConditionalGeneration.from_pretrained(
          A		,	output_loading_info=A      )

   _a					:  Tuple						=  ['key_proj', 'value_proj', 'query_proj']

   _a					:  str						=  {
       'self_attn': 'ngram_self_attn',
       'cross_attn': 'encoder_attn',
       'cross_attn_layer_norm': 'encoder_attn_layer_norm',
       'feed_forward_layer_norm': 'final_layer_norm',
       'feed_forward': '',
       'intermediate': 'fc1',
       'output': 'fc2',
       'key_proj': 'k_proj',
       'query_proj': 'q_proj',
       'value_proj': 'v_proj',
       'word_embeddings': 'embed_tokens',
       'embeddings_layer_norm': 'emb_layer_norm',
       'relative_pos_embeddings': 'relative_linear',
       'ngram_embeddings': 'ngram_input_embed',
       'position_embeddings': 'embed_positions',
   }

   for key in loading_info["missing_keys"]:
      _a					:  Tuple						=  key.split('.'      )

      if attributes[0] == "lm_head":
         _a					:  str						=  prophet
         _a					:  Dict						=  prophet_old
      else:
         _a					:  Union[str, Any]						=  prophet.prophetnet
         _a					:  List[Any]						=  prophet_old.model

      _a					:  Dict						=  False
      for attribute in attributes:
         if attribute in mapping:
            _a					:  List[str]						=  mapping[attribute]
            if not hasattr(A		,	A      ) and len(A      ) > 0:
               _a					:  Any						=  attribute
         elif hasattr(A		,	A      ):
            _a					:  Dict						=  attribute

         if attribute == "weight":
            assert old_model.weight.shape == model.weight.shape, "Shapes have to match!"
            _a					:  str						=  old_model.weight
            logger.info(f'''{attribute} is initialized.'''      )
            _a					:  List[str]						=  True
            break
         elif attribute == "bias":
            assert old_model.bias.shape == model.bias.shape, "Shapes have to match!"
            _a					:  int						=  old_model.bias
            logger.info(f'''{attribute} is initialized'''      )
            _a					:  Optional[int]						=  True
            break
         elif attribute in special_keys and hasattr(A		,	'in_proj_weight'      ):
            _a					:  Optional[int]						=  old_model.in_proj_weight.shape[0] // 3
            _a					:  List[str]						=  getattr(A		,	A      )
            param.weight.shape == old_model.in_proj_weight[:embed_dim, :].shape, "Shapes have to match"
            param.bias.shape == old_model.in_proj_bias[:embed_dim].shape, "Shapes have to match"
            if attribute == "query_proj":
               _a					:  Optional[int]						=  nn.Parameter(old_model.in_proj_weight[:embed_dim, :]      )
               _a					:  Union[str, Any]						=  nn.Parameter(old_model.in_proj_bias[:embed_dim]      )

            elif attribute == "key_proj":
               _a					:  str						=  nn.Parameter(old_model.in_proj_weight[embed_dim : 2 * embed_dim, :]      )
               _a					:  List[Any]						=  nn.Parameter(old_model.in_proj_bias[embed_dim : 2 * embed_dim]      )
            elif attribute == "value_proj":
               _a					:  Optional[int]						=  nn.Parameter(old_model.in_proj_weight[2 * embed_dim :, :]      )
               _a					:  str						=  nn.Parameter(old_model.in_proj_bias[2 * embed_dim :]      )
            _a					:  int						=  True
            break
         elif attribute == "position_embeddings":
            assert (
                model.position_embeddings.weight.shape[-1] == old_model.embed_positions.weight.shape[-1]
            ), "Hidden size has to match"
            assert model.position_embeddings.weight.shape[0] == 5_1_2, "We want 512 position_embeddings."
            _a					:  Tuple						=  nn.Parameter(old_model.embed_positions.weight[:5_1_2, :]      )
            _a					:  Any						=  True
            break

         if attribute.isdigit():
            _a					:  List[str]						=  model[int(A      )]
            _a					:  Union[str, Any]						=  old_model[int(A      )]
         else:
            _a					:  List[Any]						=  getattr(A		,	A      )

            if old_attribute == "":
               _a					:  List[str]						=  old_model
            else:
               if not hasattr(A		,	A      ):
                  raise ValueError(f'''{old_model} does not have {old_attribute}'''      )
               _a					:  Optional[int]						=  getattr(A		,	A      )

      if not is_key_init:
         raise ValueError(f'''{key} was not correctly initialized!'''      )

   print(f'''Saving model to {pytorch_dump_folder_path}'''      )
   prophet.save_pretrained(A      )


if __name__ == "__main__":
 UpperCAmelCase_      :		Any          =		argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     "--prophetnet_checkpoint_path", default=None, type=str, required=True, help="Path the official PyTorch dump."
 )
 parser.add_argument(
     "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
 )
 UpperCAmelCase_      :		List[str]          =		parser.parse_args()
 convert_prophetnet_checkpoint_to_pytorch(args.prophetnet_checkpoint_path, args.pytorch_dump_folder_path)


