






import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor, ResNetConfig
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase			:					str         =	logging.get_logger(__name__)





def       __lowerCAmelCase      (						__snake_case       ):
     # initialize config
     if "resnet-50" in model_name:
          __lowerCAmelCase								=					ResNetConfig.from_pretrained("microsoft/resnet-50"       )
     elif "resnet-101" in model_name:
          __lowerCAmelCase								=					ResNetConfig.from_pretrained("microsoft/resnet-101"       )
     else:
          raise ValueError("Model name should include either resnet50 or resnet101"       )

     __lowerCAmelCase								=					DetrConfig(use_timm_backbone=__snake_case					, backbone_config=__snake_case       )

     # set label attributes
     __lowerCAmelCase								=					"panoptic" in model_name
     if is_panoptic:
          __lowerCAmelCase								=					250
     else:
          __lowerCAmelCase								=					91
          __lowerCAmelCase								=					"huggingface/label-files"
          __lowerCAmelCase								=					"coco-detection-id2label.json"
          __lowerCAmelCase								=					json.load(open(hf_hub_download(__snake_case					, __snake_case					, repo_type="dataset"       )					, "r"       )       )
          __lowerCAmelCase								=					{int(__snake_case       ): v for k, v in idalabel.items()}
          __lowerCAmelCase								=					idalabel
          __lowerCAmelCase								=					{v: k for k, v in idalabel.items()}

     return config, is_panoptic





def       __lowerCAmelCase      (						__snake_case       ):
     # here we list all keys to be renamed (original name on the left, our name on the right)
     __lowerCAmelCase								=					[]

     # stem
     # fmt: off
     rename_keys.append(("backbone.0.body.conv1.weight", "backbone.conv_encoder.model.embedder.embedder.convolution.weight")       )
     rename_keys.append(("backbone.0.body.bn1.weight", "backbone.conv_encoder.model.embedder.embedder.normalization.weight")       )
     rename_keys.append(("backbone.0.body.bn1.bias", "backbone.conv_encoder.model.embedder.embedder.normalization.bias")       )
     rename_keys.append(("backbone.0.body.bn1.running_mean", "backbone.conv_encoder.model.embedder.embedder.normalization.running_mean")       )
     rename_keys.append(("backbone.0.body.bn1.running_var", "backbone.conv_encoder.model.embedder.embedder.normalization.running_var")       )
     # stages
     for stage_idx in range(len(config.backbone_config.depths       )       ):
          for layer_idx in range(config.backbone_config.depths[stage_idx]       ):
               # shortcut
               if layer_idx == 0:
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.0.weight""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.convolution.weight""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.weight""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.weight""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.bias""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.bias""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_mean""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_mean""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.downsample.1.running_var""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.shortcut.normalization.running_var""",
                        )       )
               # 3 convs
               for i in range(3       ):
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.conv{i+1}.weight""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.convolution.weight""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.weight""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.weight""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.bias""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.bias""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_mean""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_mean""",
                        )       )
                    rename_keys.append(
                        (
                            F"""backbone.0.body.layer{stage_idx + 1}.{layer_idx}.bn{i+1}.running_var""",
                            F"""backbone.conv_encoder.model.encoder.stages.{stage_idx}.layers.{layer_idx}.layer.{i}.normalization.running_var""",
                        )       )
    # fmt: on

     for i in range(config.encoder_layers       ):
          # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
          rename_keys.append(
              (
                  F"""transformer.encoder.layers.{i}.self_attn.out_proj.weight""",
                  F"""encoder.layers.{i}.self_attn.out_proj.weight""",
              )       )
          rename_keys.append(
              (F"""transformer.encoder.layers.{i}.self_attn.out_proj.bias""", F"""encoder.layers.{i}.self_attn.out_proj.bias""")       )
          rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.weight""", F"""encoder.layers.{i}.fc1.weight""")       )
          rename_keys.append((F"""transformer.encoder.layers.{i}.linear1.bias""", F"""encoder.layers.{i}.fc1.bias""")       )
          rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.weight""", F"""encoder.layers.{i}.fc2.weight""")       )
          rename_keys.append((F"""transformer.encoder.layers.{i}.linear2.bias""", F"""encoder.layers.{i}.fc2.bias""")       )
          rename_keys.append(
              (F"""transformer.encoder.layers.{i}.norm1.weight""", F"""encoder.layers.{i}.self_attn_layer_norm.weight""")       )
          rename_keys.append(
              (F"""transformer.encoder.layers.{i}.norm1.bias""", F"""encoder.layers.{i}.self_attn_layer_norm.bias""")       )
          rename_keys.append(
              (F"""transformer.encoder.layers.{i}.norm2.weight""", F"""encoder.layers.{i}.final_layer_norm.weight""")       )
          rename_keys.append((F"""transformer.encoder.layers.{i}.norm2.bias""", F"""encoder.layers.{i}.final_layer_norm.bias""")       )
          # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
          rename_keys.append(
              (
                  F"""transformer.decoder.layers.{i}.self_attn.out_proj.weight""",
                  F"""decoder.layers.{i}.self_attn.out_proj.weight""",
              )       )
          rename_keys.append(
              (F"""transformer.decoder.layers.{i}.self_attn.out_proj.bias""", F"""decoder.layers.{i}.self_attn.out_proj.bias""")       )
          rename_keys.append(
              (
                  F"""transformer.decoder.layers.{i}.multihead_attn.out_proj.weight""",
                  F"""decoder.layers.{i}.encoder_attn.out_proj.weight""",
              )       )
          rename_keys.append(
              (
                  F"""transformer.decoder.layers.{i}.multihead_attn.out_proj.bias""",
                  F"""decoder.layers.{i}.encoder_attn.out_proj.bias""",
              )       )
          rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.weight""", F"""decoder.layers.{i}.fc1.weight""")       )
          rename_keys.append((F"""transformer.decoder.layers.{i}.linear1.bias""", F"""decoder.layers.{i}.fc1.bias""")       )
          rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.weight""", F"""decoder.layers.{i}.fc2.weight""")       )
          rename_keys.append((F"""transformer.decoder.layers.{i}.linear2.bias""", F"""decoder.layers.{i}.fc2.bias""")       )
          rename_keys.append(
              (F"""transformer.decoder.layers.{i}.norm1.weight""", F"""decoder.layers.{i}.self_attn_layer_norm.weight""")       )
          rename_keys.append(
              (F"""transformer.decoder.layers.{i}.norm1.bias""", F"""decoder.layers.{i}.self_attn_layer_norm.bias""")       )
          rename_keys.append(
              (F"""transformer.decoder.layers.{i}.norm2.weight""", F"""decoder.layers.{i}.encoder_attn_layer_norm.weight""")       )
          rename_keys.append(
              (F"""transformer.decoder.layers.{i}.norm2.bias""", F"""decoder.layers.{i}.encoder_attn_layer_norm.bias""")       )
          rename_keys.append(
              (F"""transformer.decoder.layers.{i}.norm3.weight""", F"""decoder.layers.{i}.final_layer_norm.weight""")       )
          rename_keys.append((F"""transformer.decoder.layers.{i}.norm3.bias""", F"""decoder.layers.{i}.final_layer_norm.bias""")       )

     # convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
     rename_keys.extend(
         [
             ("input_proj.weight", "input_projection.weight"),
             ("input_proj.bias", "input_projection.bias"),
             ("query_embed.weight", "query_position_embeddings.weight"),
             ("transformer.decoder.norm.weight", "decoder.layernorm.weight"),
             ("transformer.decoder.norm.bias", "decoder.layernorm.bias"),
             ("class_embed.weight", "class_labels_classifier.weight"),
             ("class_embed.bias", "class_labels_classifier.bias"),
             ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"),
             ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"),
             ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"),
             ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"),
             ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"),
             ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"),
         ]       )

     return rename_keys





def       __lowerCAmelCase      (						__snake_case					, __snake_case					, __snake_case       ):
     __lowerCAmelCase								=					state_dict.pop(__snake_case       )
     __lowerCAmelCase								=					val





def       __lowerCAmelCase      (						__snake_case					, __snake_case=False       ):
     __lowerCAmelCase								=					""
     if is_panoptic:
          __lowerCAmelCase								=					"detr."

     # first: transformer encoder
     for i in range(6       ):
          # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
          __lowerCAmelCase								=					state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""       )
          __lowerCAmelCase								=					state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""       )
          # next, add query, keys and values (in that order) to the state dict
          __lowerCAmelCase								=					in_proj_weight[:256, :]
          __lowerCAmelCase								=					in_proj_bias[:256]
          __lowerCAmelCase								=					in_proj_weight[256:512, :]
          __lowerCAmelCase								=					in_proj_bias[256:512]
          __lowerCAmelCase								=					in_proj_weight[-256:, :]
          __lowerCAmelCase								=					in_proj_bias[-256:]
     # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
     for i in range(6       ):
          # read in weights + bias of input projection layer of self-attention
          __lowerCAmelCase								=					state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""       )
          __lowerCAmelCase								=					state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"""       )
          # next, add query, keys and values (in that order) to the state dict
          __lowerCAmelCase								=					in_proj_weight[:256, :]
          __lowerCAmelCase								=					in_proj_bias[:256]
          __lowerCAmelCase								=					in_proj_weight[256:512, :]
          __lowerCAmelCase								=					in_proj_bias[256:512]
          __lowerCAmelCase								=					in_proj_weight[-256:, :]
          __lowerCAmelCase								=					in_proj_bias[-256:]
          # read in weights + bias of input projection layer of cross-attention
          __lowerCAmelCase								=					state_dict.pop(
              F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""       )
          __lowerCAmelCase								=					state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"""       )
          # next, add query, keys and values (in that order) of cross-attention to the state dict
          __lowerCAmelCase								=					in_proj_weight_cross_attn[:256, :]
          __lowerCAmelCase								=					in_proj_bias_cross_attn[:256]
          __lowerCAmelCase								=					in_proj_weight_cross_attn[256:512, :]
          __lowerCAmelCase								=					in_proj_bias_cross_attn[256:512]
          __lowerCAmelCase								=					in_proj_weight_cross_attn[-256:, :]
          __lowerCAmelCase								=					in_proj_bias_cross_attn[-256:]





def       __lowerCAmelCase      (						):
     __lowerCAmelCase								=					"http://images.cocodataset.org/val2017/000000039769.jpg"
     __lowerCAmelCase								=					Image.open(requests.get(__snake_case					, stream=__snake_case       ).raw       )

     return im





@torch.no_grad()
def       __lowerCAmelCase      (						__snake_case					, __snake_case=None					, __snake_case=False       ):
     __lowerCAmelCase			, __lowerCAmelCase								=					get_detr_config(__snake_case       )

     # load original model from torch hub
     __lowerCAmelCase								=					{
         "detr-resnet-50": "detr_resnet50",
         "detr-resnet-101": "detr_resnet101",
     }
     logger.info(F"""Converting model {model_name}..."""       )
     __lowerCAmelCase								=					torch.hub.load("facebookresearch/detr"					, model_name_to_original_name[model_name]					, pretrained=__snake_case       ).eval()
     __lowerCAmelCase								=					detr.state_dict()
     # rename keys
     for src, dest in create_rename_keys(__snake_case       ):
          if is_panoptic:
               __lowerCAmelCase								=					"detr." + src
          rename_key(__snake_case					, __snake_case					, __snake_case       )
     # query, key and value matrices need special treatment
     read_in_q_k_v(__snake_case					, is_panoptic=__snake_case       )
     # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
     __lowerCAmelCase								=					"detr.model." if is_panoptic else "model."
     for key in state_dict.copy().keys():
          if is_panoptic:
               if (
                   key.startswith("detr"       )
                   and not key.startswith("class_labels_classifier"       )
                   and not key.startswith("bbox_predictor"       )
               ):
                    __lowerCAmelCase								=					state_dict.pop(__snake_case       )
                    __lowerCAmelCase								=					val
               elif "class_labels_classifier" in key or "bbox_predictor" in key:
                    __lowerCAmelCase								=					state_dict.pop(__snake_case       )
                    __lowerCAmelCase								=					val
               elif key.startswith("bbox_attention"       ) or key.startswith("mask_head"       ):
                    continue
               else:
                    __lowerCAmelCase								=					state_dict.pop(__snake_case       )
                    __lowerCAmelCase								=					val
          else:
               if not key.startswith("class_labels_classifier"       ) and not key.startswith("bbox_predictor"       ):
                    __lowerCAmelCase								=					state_dict.pop(__snake_case       )
                    __lowerCAmelCase								=					val

    # finally, create HuggingFace model and load state dict
     __lowerCAmelCase								=					DetrForSegmentation(__snake_case       ) if is_panoptic else DetrForObjectDetection(__snake_case       )
     model.load_state_dict(__snake_case       )
     model.eval()

     # verify our conversion on an image
     __lowerCAmelCase								=					"coco_panoptic" if is_panoptic else "coco_detection"
     __lowerCAmelCase								=					DetrImageProcessor(format=__snake_case       )

     __lowerCAmelCase								=					processor(images=prepare_img()					, return_tensors="pt"       )
     __lowerCAmelCase								=					encoding["pixel_values"]

     __lowerCAmelCase								=					detr(__snake_case       )
     __lowerCAmelCase								=					model(__snake_case       )

     assert torch.allclose(outputs.logits					, original_outputs["pred_logits"]					, atol=1E-3       )
     assert torch.allclose(outputs.pred_boxes					, original_outputs["pred_boxes"]					, atol=1E-3       )
     if is_panoptic:
          assert torch.allclose(outputs.pred_masks					, original_outputs["pred_masks"]					, atol=1E-4       )
     print("Looks ok!"       )

     if pytorch_dump_folder_path is not None:
          # Save model and image processor
          logger.info(F"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""       )
          Path(__snake_case       ).mkdir(exist_ok=__snake_case       )
          model.save_pretrained(__snake_case       )
          processor.save_pretrained(__snake_case       )

     if push_to_hub:
          # Upload model and image processor to the hub
          logger.info("Uploading PyTorch model and image processor to the hub..."       )
          model.push_to_hub(F"""nielsr/{model_name}"""       )
          processor.push_to_hub(F"""nielsr/{model_name}"""       )


if __name__ == "__main__":
 lowerCamelCase			:					Any         =	argparse.ArgumentParser()

 parser.add_argument(
     '''--model_name''',
     default='''detr-resnet-50''',
     type=str,
     choices=['''detr-resnet-50''', '''detr-resnet-101'''],
     help='''Name of the DETR model you\'d like to convert.''',
 )
 parser.add_argument(
     '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
 )
 parser.add_argument('''--push_to_hub''', action='''store_true''', help='''Whether to push the model to the hub or not.''')
 lowerCamelCase			:					int         =	parser.parse_args()
 convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)














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			:					Any         =	logging.get_logger(__name__)

lowerCamelCase			:					Tuple         =	{'''vocab_file''': '''vocab.txt''', '''tokenizer_file''': '''tokenizer.json'''}

lowerCamelCase			:					Optional[int]         =	{
    '''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			:					List[str]         =	{
    '''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			:					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			:					str         =	{
    '''facebook/dpr-ctx_encoder-single-nq-base''': 512,
    '''facebook/dpr-ctx_encoder-multiset-base''': 512,
}
lowerCamelCase			:					Any         =	{
    '''facebook/dpr-question_encoder-single-nq-base''': 512,
    '''facebook/dpr-question_encoder-multiset-base''': 512,
}
lowerCamelCase			:					Tuple         =	{
    '''facebook/dpr-reader-single-nq-base''': 512,
    '''facebook/dpr-reader-multiset-base''': 512,
}


lowerCamelCase			:					Tuple         =	{
    '''facebook/dpr-ctx_encoder-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-ctx_encoder-multiset-base''': {'''do_lower_case''': True},
}
lowerCamelCase			:					int         =	{
    '''facebook/dpr-question_encoder-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-question_encoder-multiset-base''': {'''do_lower_case''': True},
}
lowerCamelCase			:					List[Any]         =	{
    '''facebook/dpr-reader-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-reader-multiset-base''': {'''do_lower_case''': True},
}



class 							_UpperCamelCase			(a_			):
      snake_case_  =							VOCAB_FILES_NAMES
      snake_case_  =							CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
      snake_case_  =							CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      snake_case_  =							CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION
      snake_case_  =							DPRContextEncoderTokenizer



class 							_UpperCamelCase			(a_			):
      snake_case_  =							VOCAB_FILES_NAMES
      snake_case_  =							QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
      snake_case_  =							QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      snake_case_  =							QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION
      snake_case_  =							DPRQuestionEncoderTokenizer

lowerCamelCase			:					Union[str, Any]         =	collections.namedtuple(
    '''DPRSpanPrediction''', ['''span_score''', '''relevance_score''', '''doc_id''', '''start_index''', '''end_index''', '''text''']
)

lowerCamelCase			:					List[str]         =	collections.namedtuple('''DPRReaderOutput''', ['''start_logits''', '''end_logits''', '''relevance_logits'''])


lowerCamelCase			:					List[Any]         =	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(a_			)
class 							_UpperCamelCase			:


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


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

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


      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase					,						__UpperCamelCase					,						__UpperCamelCase					,						)->	List[DPRSpanPrediction]:
           __lowerCAmelCase								=					[]
           for start_index, start_score in enumerate(__UpperCamelCase	):
                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)	)
           __lowerCAmelCase								=					sorted(__UpperCamelCase					,						key=lambda __UpperCamelCase	: x[1]					,						reverse=__UpperCamelCase	)
           __lowerCAmelCase								=					[]
           for (start_index, end_index), score in scores:
                assert start_index <= end_index, F"""Wrong span indices: [{start_index}:{end_index}]"""
                __lowerCAmelCase								=					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(__UpperCamelCase	) == top_spans:
                     break
           return chosen_span_intervals



@add_end_docstrings(a_			)
class 							_UpperCamelCase			(a_							,		a_			):
      snake_case_  =							VOCAB_FILES_NAMES
      snake_case_  =							READER_PRETRAINED_VOCAB_FILES_MAP
      snake_case_  =							READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      snake_case_  =							READER_PRETRAINED_INIT_CONFIGURATION
      snake_case_  =							["""input_ids""", """attention_mask"""]
      snake_case_  =							DPRReaderTokenizer






