
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


lowerCamelCase__     :		Optional[int]        = logging.get_logger(__name__)

lowerCamelCase__     :		Dict        = {"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

lowerCamelCase__     :		Union[str, 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"""
        ),
    },
}
lowerCamelCase__     :		Tuple        = {
    """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"""
        ),
    },
}
lowerCamelCase__     :		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"""
        ),
    },
}

lowerCamelCase__     :		Tuple        = {
    """facebook/dpr-ctx_encoder-single-nq-base""": 5_1_2,
    """facebook/dpr-ctx_encoder-multiset-base""": 5_1_2,
}
lowerCamelCase__     :		Optional[Any]        = {
    """facebook/dpr-question_encoder-single-nq-base""": 5_1_2,
    """facebook/dpr-question_encoder-multiset-base""": 5_1_2,
}
lowerCamelCase__     :		Dict        = {
    """facebook/dpr-reader-single-nq-base""": 5_1_2,
    """facebook/dpr-reader-multiset-base""": 5_1_2,
}


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



class    _snake_case (						UpperCAmelCase_				):
 __lowerCAmelCase	:  Optional[int]	   =      VOCAB_FILES_NAMES
 __lowerCAmelCase	:  int	   =      CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase	:  List[Any]	   =      CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase	:  List[Any]	   =      CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION
 __lowerCAmelCase	:  Any	   =      DPRContextEncoderTokenizer



class    _snake_case (						UpperCAmelCase_				):
 __lowerCAmelCase	:  Optional[Any]	   =      VOCAB_FILES_NAMES
 __lowerCAmelCase	:  Any	   =      QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase	:  List[Any]	   =      QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase	:  Optional[int]	   =      QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION
 __lowerCAmelCase	:  List[str]	   =      DPRQuestionEncoderTokenizer

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

lowerCamelCase__     :		int        = collections.namedtuple("""DPRReaderOutput""", ["""start_logits""", """end_logits""", """relevance_logits"""])


lowerCamelCase__     :		Optional[int]        = R"""
    Return a dictionary with the token ids of the input strings and other information to give to `.decode_best_spans`.
    It converts the strings of a question and different passages (title and text) in a sequence of IDs (integers),
    using the tokenizer and vocabulary. The resulting `input_ids` is a matrix of size `(n_passages, sequence_length)`
    with the format:

    [CLS] <question token ids> [SEP] <titles ids> [SEP] <texts ids>

    Args:
        questions (`str` or `List[str]`):
            The questions to be encoded. You can specify one question for many passages. In this case, the question
            will be duplicated like `[questions] * n_passages`. Otherwise you have to specify as many questions as in
            `titles` or `texts`.
        titles (`str` or `List[str]`):
            The passages titles to be encoded. This can be a string or a list of strings if there are several passages.
        texts (`str` or `List[str]`):
            The passages texts to be encoded. This can be a string or a list of strings if there are several passages.
        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
            Activates and controls padding. Accepts the following values:

            - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single sequence
              if provided).
            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided.
            - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
              lengths).
        truncation (`bool`, `str` or [`~tokenization_utils_base.TruncationStrategy`], *optional*, defaults to `False`):
            Activates and controls truncation. Accepts the following values:

            - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or to
              the maximum acceptable input length for the model if that argument is not provided. This will truncate
              token by token, removing a token from the longest sequence in the pair if a pair of sequences (or a batch
              of pairs) is provided.
            - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the first
              sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the maximum
              acceptable input length for the model if that argument is not provided. This will only truncate the
              second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
            - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths
              greater than the model maximum admissible input size).
        max_length (`int`, *optional*):
                Controls the maximum length to use by one of the truncation/padding parameters.

                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
                is required by one of the truncation/padding parameters. If the model has no specific maximum input
                length (like XLNet) truncation/padding to a maximum length will be deactivated.
        return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors instead of list of python integers. Acceptable values are:

                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return Numpy `np.ndarray` objects.
        return_attention_mask (`bool`, *optional*):
            Whether or not to return the attention mask. If not set, will return the attention mask according to the
            specific tokenizer's default, defined by the `return_outputs` attribute.

            [What are attention masks?](../glossary#attention-mask)

    Return:
        `Dict[str, List[List[int]]]`: A dictionary with the following keys:

        - `input_ids`: List of token ids to be fed to a model.
        - `attention_mask`: List of indices specifying which tokens should be attended to by the model.
    """



@add_start_docstrings(UpperCAmelCase_				)
class    _snake_case :


 def __call__(       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_ = None       ,	SCREAMING_SNAKE_CASE_ = None       ,	SCREAMING_SNAKE_CASE_ = False       ,	SCREAMING_SNAKE_CASE_ = False       ,	SCREAMING_SNAKE_CASE_ = None       ,	SCREAMING_SNAKE_CASE_ = None       ,	SCREAMING_SNAKE_CASE_ = None       ,	**SCREAMING_SNAKE_CASE_       ,	):


   '''simple docstring'''




   if titles is None and texts is None:
     return super().__call__(
         SCREAMING_SNAKE_CASE_       ,	padding=SCREAMING_SNAKE_CASE_       ,	truncation=SCREAMING_SNAKE_CASE_       ,	max_length=SCREAMING_SNAKE_CASE_       ,	return_tensors=SCREAMING_SNAKE_CASE_       ,	return_attention_mask=SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_       ,	)
   elif titles is None or texts is None:
     lowercase__	: Union[str, Any]     		=   titles if texts is None else texts
     return super().__call__(
         SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	padding=SCREAMING_SNAKE_CASE_       ,	truncation=SCREAMING_SNAKE_CASE_       ,	max_length=SCREAMING_SNAKE_CASE_       ,	return_tensors=SCREAMING_SNAKE_CASE_       ,	return_attention_mask=SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_       ,	)
   lowercase__	: Optional[int]     		=   titles if not isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) else [titles]
   lowercase__	: List[Any]     		=   texts if not isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) else [texts]
   lowercase__	: int     		=   len(SCREAMING_SNAKE_CASE_)
   lowercase__	: Optional[int]     		=   questions if not isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) else [questions] * n_passages
   assert len(SCREAMING_SNAKE_CASE_) == len(
       SCREAMING_SNAKE_CASE_), f'There should be as many titles than texts but got {len(SCREAMING_SNAKE_CASE_)} titles and {len(SCREAMING_SNAKE_CASE_)} texts.'
   lowercase__	: Dict     		=   super().__call__(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	padding=SCREAMING_SNAKE_CASE_       ,	truncation=SCREAMING_SNAKE_CASE_)["""input_ids"""]
   lowercase__	: Union[str, Any]     		=   super().__call__(SCREAMING_SNAKE_CASE_       ,	add_special_tokens=SCREAMING_SNAKE_CASE_       ,	padding=SCREAMING_SNAKE_CASE_       ,	truncation=SCREAMING_SNAKE_CASE_)["""input_ids"""]
   lowercase__	: Union[str, Any]     		=   {
       """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(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)
       ]
   }
   if return_attention_mask is not False:
     lowercase__	: Optional[int]     		=   []
     for input_ids in encoded_inputs["input_ids"]:
       attention_mask.append([int(input_id != self.pad_token_id) for input_id in input_ids])
     lowercase__	: List[str]     		=   attention_mask
   return self.pad(SCREAMING_SNAKE_CASE_       ,	padding=SCREAMING_SNAKE_CASE_       ,	max_length=SCREAMING_SNAKE_CASE_       ,	return_tensors=SCREAMING_SNAKE_CASE_)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_ = 16       ,	SCREAMING_SNAKE_CASE_ = 64       ,	SCREAMING_SNAKE_CASE_ = 4       ,	):


   '''simple docstring'''




   lowercase__	: Dict     		=   reader_input["""input_ids"""]
   lowercase__		,     lowercase__		,     lowercase__	: str     		=   reader_output[:3]
   lowercase__	: List[Any]     		=   len(SCREAMING_SNAKE_CASE_)
   lowercase__	: Union[str, Any]     		=   sorted(range(SCREAMING_SNAKE_CASE_)       ,	reverse=SCREAMING_SNAKE_CASE_       ,	key=relevance_logits.__getitem__)
   lowercase__	: List[DPRReaderOutput]     		=   []
   for doc_id in sorted_docs:
     lowercase__	: Dict     		=   list(input_ids[doc_id])
     # assuming question & title information is at the beginning of the sequence
     lowercase__	: Optional[Any]     		=   sequence_ids.index(self.sep_token_id       ,	2) + 1  # second sep id
     if sequence_ids[-1] == self.pad_token_id:
       lowercase__	: int     		=   sequence_ids.index(self.pad_token_id)
     else:
       lowercase__	: Optional[int]     		=   len(SCREAMING_SNAKE_CASE_)

     lowercase__	: int     		=   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=SCREAMING_SNAKE_CASE_       ,	top_spans=SCREAMING_SNAKE_CASE_       ,	)
     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=SCREAMING_SNAKE_CASE_       ,	start_index=SCREAMING_SNAKE_CASE_       ,	end_index=SCREAMING_SNAKE_CASE_       ,	text=self.decode(sequence_ids[start_index : end_index + 1])       ,	))
     if len(SCREAMING_SNAKE_CASE_) >= num_spans:
       break
   return nbest_spans_predictions[:num_spans]



 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	):


   '''simple docstring'''




   lowercase__	: Tuple     		=   []
   for start_index, start_score in enumerate(SCREAMING_SNAKE_CASE_):
     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))
   lowercase__	: Optional[Any]     		=   sorted(SCREAMING_SNAKE_CASE_       ,	key=lambda SCREAMING_SNAKE_CASE_: x[1]       ,	reverse=SCREAMING_SNAKE_CASE_)
   lowercase__	: Tuple     		=   []
   for (start_index, end_index), score in scores:
     assert start_index <= end_index, f'Wrong span indices: [{start_index}:{end_index}]'
     lowercase__	: Dict     		=   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(SCREAMING_SNAKE_CASE_) == top_spans:
       break
   return chosen_span_intervals



@add_end_docstrings(UpperCAmelCase_				)
class    _snake_case (						UpperCAmelCase_	,    UpperCAmelCase_				):
 __lowerCAmelCase	:  int	   =      VOCAB_FILES_NAMES
 __lowerCAmelCase	:  Optional[Any]	   =      READER_PRETRAINED_VOCAB_FILES_MAP
 __lowerCAmelCase	:  Any	   =      READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 __lowerCAmelCase	:  Any	   =      READER_PRETRAINED_INIT_CONFIGURATION
 __lowerCAmelCase	:  Optional[Any]	   =      ['input_ids', 'attention_mask']
 __lowerCAmelCase	:  int	   =      DPRReaderTokenizer








import argparse
import os

import torch

from transformers import FlavaConfig, FlavaForPreTraining
from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint

def 		UpperCamelCase						(   lowercase_ )		->						Union[str, Any]:


  '''simple docstring'''






  return sum(param.float().sum() if """encoder.embeddings""" not in key else 0 for key, param in state_dict.items() )

def 		UpperCamelCase						(   lowercase_			,			lowercase_ )		->						List[Any]:


  '''simple docstring'''






  lowercase__	: int     		=   {}

  for key, value in state_dict.items():
    if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
      continue

    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.mim_head.cls.predictions"""			,			"""mmm_image_head""" )
    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.mlm_head.cls.predictions"""			,			"""mmm_text_head""" )
    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.itm_head.cls"""			,			"""itm_head""" )
    lowercase__	: Tuple     		=   key.replace("""heads.cmd.itm_head.pooler"""			,			"""itm_head.pooler""" )
    lowercase__	: Optional[Any]     		=   key.replace("""heads.cmd.clip_head.logit_scale"""			,			"""flava.logit_scale""" )
    lowercase__	: Optional[int]     		=   key.replace("""heads.fairseq_mlm.cls.predictions"""			,			"""mlm_head""" )
    lowercase__	: List[Any]     		=   key.replace("""heads.imagenet.mim_head.cls.predictions"""			,			"""mim_head""" )
    lowercase__	: int     		=   key.replace("""mm_text_projection"""			,			"""flava.text_to_mm_projection""" )
    lowercase__	: Optional[Any]     		=   key.replace("""mm_image_projection"""			,			"""flava.image_to_mm_projection""" )
    lowercase__	: Optional[Any]     		=   key.replace("""image_encoder.module"""			,			"""flava.image_model""" )
    lowercase__	: Any     		=   key.replace("""text_encoder.module"""			,			"""flava.text_model""" )
    lowercase__	: Optional[Any]     		=   key.replace("""mm_encoder.module.encoder.cls_token"""			,			"""flava.multimodal_model.cls_token""" )
    lowercase__	: Tuple     		=   key.replace("""mm_encoder.module"""			,			"""flava.multimodal_model""" )
    lowercase__	: Any     		=   key.replace("""text_projection"""			,			"""flava.text_projection""" )
    lowercase__	: List[Any]     		=   key.replace("""image_projection"""			,			"""flava.image_projection""" )

    lowercase__	: str     		=   value.float()

  for key, value in codebook_state_dict.items():
    lowercase__	: Any     		=   value

  return upgrade

@torch.no_grad()
def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_			,			lowercase_=None )		->						Union[str, Any]:


  '''simple docstring'''






  if config_path is not None:
    lowercase__	: int     		=   FlavaConfig.from_pretrained(lowercase_ )
  else:
    lowercase__	: Optional[int]     		=   FlavaConfig()

  lowercase__	: List[Any]     		=   FlavaForPreTraining(lowercase_ ).eval()

  lowercase__	: Dict     		=   convert_dalle_checkpoint(lowercase_			,			lowercase_			,			save_checkpoint=lowercase_ )

  if os.path.exists(lowercase_ ):
    lowercase__	: Dict     		=   torch.load(lowercase_			,			map_location="""cpu""" )
  else:
    lowercase__	: Dict     		=   torch.hub.load_state_dict_from_url(lowercase_			,			map_location="""cpu""" )

  lowercase__	: int     		=   upgrade_state_dict(lowercase_			,			lowercase_ )
  hf_model.load_state_dict(lowercase_ )
  lowercase__	: Optional[int]     		=   hf_model.state_dict()
  lowercase__	: Optional[int]     		=   count_parameters(lowercase_ )
  lowercase__	: Any     		=   count_parameters(lowercase_ ) + count_parameters(lowercase_ )

  assert torch.allclose(lowercase_			,			lowercase_			,			atol=1E-3 )

  hf_model.save_pretrained(lowercase_ )


if __name__ == "__main__":
     lowerCamelCase__     :		int        = argparse.ArgumentParser()
     parser.add_argument("""--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model.""")
     parser.add_argument("""--checkpoint_path""", default=None, type=str, help="""Path to flava checkpoint""")
     parser.add_argument("""--codebook_path""", default=None, type=str, help="""Path to flava codebook checkpoint""")
     parser.add_argument("""--config_path""", default=None, type=str, help="""Path to hf config.json of model to convert""")
     lowerCamelCase__     :		List[str]        = parser.parse_args()

     convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)






