






import inspect
import unittest

from transformers import ViTMSNConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch
 from torch import nn

 from transformers import ViTMSNForImageClassification, ViTMSNModel
 from transformers.models.vit_msn.modeling_vit_msn import VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import ViTImageProcessor



class 							_UpperCamelCase			:


      def __init__(	self					,						__UpperCamelCase					,						__UpperCamelCase=1_3					,						__UpperCamelCase=3_0					,						__UpperCamelCase=2					,						__UpperCamelCase=3					,						__UpperCamelCase=True					,						__UpperCamelCase=True					,						__UpperCamelCase=3_2					,						__UpperCamelCase=5					,						__UpperCamelCase=4					,						__UpperCamelCase=3_7					,						__UpperCamelCase="gelu"					,						__UpperCamelCase=0.1					,						__UpperCamelCase=0.1					,						__UpperCamelCase=1_0					,						__UpperCamelCase=0.0_2					,						__UpperCamelCase=None					,						)->	Optional[int]:
           __lowerCAmelCase								=					parent
           __lowerCAmelCase								=					batch_size
           __lowerCAmelCase								=					image_size
           __lowerCAmelCase								=					patch_size
           __lowerCAmelCase								=					num_channels
           __lowerCAmelCase								=					is_training
           __lowerCAmelCase								=					use_labels
           __lowerCAmelCase								=					hidden_size
           __lowerCAmelCase								=					num_hidden_layers
           __lowerCAmelCase								=					num_attention_heads
           __lowerCAmelCase								=					intermediate_size
           __lowerCAmelCase								=					hidden_act
           __lowerCAmelCase								=					hidden_dropout_prob
           __lowerCAmelCase								=					attention_probs_dropout_prob
           __lowerCAmelCase								=					type_sequence_label_size
           __lowerCAmelCase								=					initializer_range
           __lowerCAmelCase								=					scope

           # in ViT MSN, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
           __lowerCAmelCase								=					(image_size // patch_size) ** 2
           __lowerCAmelCase								=					num_patches + 1


      def  __UpperCAmelCase		(	self	)->	Optional[int]:
           __lowerCAmelCase								=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

           __lowerCAmelCase								=					None
           if self.use_labels:
                __lowerCAmelCase								=					ids_tensor([self.batch_size]					,						self.type_sequence_label_size	)

           __lowerCAmelCase								=					self.get_config()

           return config, pixel_values, labels


      def  __UpperCAmelCase		(	self	)->	List[str]:
           return ViTMSNConfig(
               image_size=self.image_size					,						patch_size=self.patch_size					,						num_channels=self.num_channels					,						hidden_size=self.hidden_size					,						num_hidden_layers=self.num_hidden_layers					,						num_attention_heads=self.num_attention_heads					,						intermediate_size=self.intermediate_size					,						hidden_act=self.hidden_act					,						hidden_dropout_prob=self.hidden_dropout_prob					,						attention_probs_dropout_prob=self.attention_probs_dropout_prob					,						initializer_range=self.initializer_range					,						)


      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase					,						__UpperCamelCase	)->	int:
           __lowerCAmelCase								=					ViTMSNModel(config=__UpperCamelCase	)
           model.to(__UpperCamelCase	)
           model.eval()
           __lowerCAmelCase								=					model(__UpperCamelCase	)
           self.parent.assertEqual(result.last_hidden_state.shape					,						(self.batch_size, self.seq_length, self.hidden_size)	)


      def  __UpperCAmelCase		(	self					,						__UpperCamelCase					,						__UpperCamelCase					,						__UpperCamelCase	)->	int:
           __lowerCAmelCase								=					self.type_sequence_label_size
           __lowerCAmelCase								=					ViTMSNForImageClassification(__UpperCamelCase	)
           model.to(__UpperCamelCase	)
           model.eval()
           __lowerCAmelCase								=					model(__UpperCamelCase					,						labels=__UpperCamelCase	)
           print("Pixel and labels shape: {pixel_values.shape}, {labels.shape}"	)
           print("Labels: {labels}"	)
           self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.type_sequence_label_size)	)

           # test greyscale images
           __lowerCAmelCase								=					1
           __lowerCAmelCase								=					ViTMSNForImageClassification(__UpperCamelCase	)
           model.to(__UpperCamelCase	)
           model.eval()

           __lowerCAmelCase								=					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
           __lowerCAmelCase								=					model(__UpperCamelCase	)
           self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.type_sequence_label_size)	)


      def  __UpperCAmelCase		(	self	)->	List[str]:
           __lowerCAmelCase								=					self.prepare_config_and_inputs()
           __lowerCAmelCase			, __lowerCAmelCase			, __lowerCAmelCase								=					config_and_inputs
           __lowerCAmelCase								=					{"pixel_values": pixel_values}
           return config, inputs_dict



@require_torch
class 							_UpperCamelCase			(a_							,		a_							,		unittest.TestCase			):
      snake_case_  =							(ViTMSNModel, ViTMSNForImageClassification) if is_torch_available() else ()
      snake_case_  =							(
          {"""feature-extraction""": ViTMSNModel, """image-classification""": ViTMSNForImageClassification}
          if is_torch_available()
          else {}
      )

      snake_case_  =							False
      snake_case_  =							False
      snake_case_  =							False
      snake_case_  =							False


      def  __UpperCAmelCase		(	self	)->	Dict:
           __lowerCAmelCase								=					ViTMSNModelTester(self	)
           __lowerCAmelCase								=					ConfigTester(self					,						config_class=__UpperCamelCase					,						has_text_modality=__UpperCamelCase					,						hidden_size=3_7	)


      def  __UpperCAmelCase		(	self	)->	str:
           self.config_tester.run_common_tests()


      @unittest.skip(reason="ViTMSN does not use inputs_embeds"	)
      def  __UpperCAmelCase		(	self	)->	Tuple:
           pass


      def  __UpperCAmelCase		(	self	)->	Optional[int]:
           __lowerCAmelCase			, __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                __lowerCAmelCase								=					model_class(__UpperCamelCase	)
                self.assertIsInstance(model.get_input_embeddings()					,						(nn.Module)	)
                __lowerCAmelCase								=					model.get_output_embeddings()
                self.assertTrue(x is None or isinstance(__UpperCamelCase					,						nn.Linear	)	)


      def  __UpperCAmelCase		(	self	)->	List[str]:
           __lowerCAmelCase			, __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                __lowerCAmelCase								=					model_class(__UpperCamelCase	)
                __lowerCAmelCase								=					inspect.signature(model.forward	)
                # signature.parameters is an OrderedDict => so arg_names order is deterministic
                __lowerCAmelCase								=					[*signature.parameters.keys()]

                __lowerCAmelCase								=					["pixel_values"]
                self.assertListEqual(arg_names[:1]					,						__UpperCamelCase	)


      def  __UpperCAmelCase		(	self	)->	Any:
           __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*__UpperCamelCase	)


      def  __UpperCAmelCase		(	self	)->	Optional[int]:
           __lowerCAmelCase								=					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*__UpperCamelCase	)


      @slow
      def  __UpperCAmelCase		(	self	)->	List[str]:
           for model_name in VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                __lowerCAmelCase								=					ViTMSNModel.from_pretrained(__UpperCamelCase	)
                self.assertIsNotNone(__UpperCamelCase	)





def       __lowerCAmelCase      (						):
     __lowerCAmelCase								=					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"       )
     return image



@require_torch
@require_vision
class 							_UpperCamelCase			(unittest.TestCase			):


      @cached_property
      def  __UpperCAmelCase		(	self	)->	Dict:
           return ViTImageProcessor.from_pretrained("facebook/vit-msn-small"	) if is_vision_available() else None


      @slow
      def  __UpperCAmelCase		(	self	)->	List[str]:
           torch.manual_seed(2	)
           __lowerCAmelCase								=					ViTMSNForImageClassification.from_pretrained("facebook/vit-msn-small"	).to(__UpperCamelCase	)

           __lowerCAmelCase								=					self.default_image_processor
           __lowerCAmelCase								=					prepare_img()
           __lowerCAmelCase								=					image_processor(images=__UpperCamelCase					,						return_tensors="pt"	).to(__UpperCamelCase	)

           # forward pass
           with torch.no_grad():
                __lowerCAmelCase								=					model(**__UpperCamelCase	)

           # verify the logits
           __lowerCAmelCase								=					torch.Size((1, 1_0_0_0)	)
           self.assertEqual(outputs.logits.shape					,						__UpperCamelCase	)

           __lowerCAmelCase								=					torch.tensor([-0.0_8_0_3, -0.4_4_5_4, -0.2_3_7_5]	).to(__UpperCamelCase	)

           self.assertTrue(torch.allclose(outputs.logits[0, :3]					,						__UpperCamelCase					,						atol=1e-4	)	)














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






