






'''simple docstring'''





import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

from ...generation.test_utils import GenerationTesterMixin
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel


class      lowercase__     :

	def __init__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase=99		,    __UpperCAmelCase=13		,    __UpperCAmelCase=7		,    __UpperCAmelCase=9		,    __UpperCAmelCase=True		,    __UpperCAmelCase=True		,    __UpperCAmelCase=False		,    __UpperCAmelCase=32		,    __UpperCAmelCase=5		,    __UpperCAmelCase=4		,    __UpperCAmelCase=37		,    __UpperCAmelCase=8		,    __UpperCAmelCase=0.1		,    __UpperCAmelCase=0.002		,    __UpperCAmelCase=1		,    __UpperCAmelCase=0		,    __UpperCAmelCase=0		,    __UpperCAmelCase=None		,    __UpperCAmelCase=None		,    )->					Optional[int]:







							'''simple docstring'''





							lowerCAmelCase__       =						parent
							lowerCAmelCase__       =						batch_size
							lowerCAmelCase__       =						encoder_seq_length
							lowerCAmelCase__       =						decoder_seq_length
							# For common tests
							lowerCAmelCase__       =						self.decoder_seq_length
							lowerCAmelCase__       =						is_training
							lowerCAmelCase__       =						use_attention_mask
							lowerCAmelCase__       =						use_labels
							lowerCAmelCase__       =						vocab_size
							lowerCAmelCase__       =						hidden_size
							lowerCAmelCase__       =						num_hidden_layers
							lowerCAmelCase__       =						num_attention_heads
							lowerCAmelCase__       =						d_ff
							lowerCAmelCase__       =						relative_attention_num_buckets
							lowerCAmelCase__       =						dropout_rate
							lowerCAmelCase__       =						initializer_factor
							lowerCAmelCase__       =						eos_token_id
							lowerCAmelCase__       =						pad_token_id
							lowerCAmelCase__       =						decoder_start_token_id
							lowerCAmelCase__       =						None
							lowerCAmelCase__       =						decoder_layers

	def 		UpperCAmelCase  (							self       )->					Tuple:







							'''simple docstring'''





							return TaConfig.from_pretrained("google/umt5-base"       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase=None		,    __UpperCAmelCase=None		,    __UpperCAmelCase=None		,    __UpperCAmelCase=None		,    __UpperCAmelCase=None		,    )->					int:







							'''simple docstring'''





							if attention_mask is None:
													lowerCAmelCase__       =						input_ids.ne(config.pad_token_id       )
							if decoder_attention_mask is None:
													lowerCAmelCase__       =						decoder_input_ids.ne(config.pad_token_id       )
							if head_mask is None:
													lowerCAmelCase__       =						torch.ones(config.num_hidden_layers		,    config.num_attention_heads		,    device=__UpperCAmelCase       )
							if decoder_head_mask is None:
													lowerCAmelCase__       =						torch.ones(config.num_decoder_layers		,    config.num_attention_heads		,    device=__UpperCAmelCase       )
							if cross_attn_head_mask is None:
													lowerCAmelCase__       =						torch.ones(
													    config.num_decoder_layers		,    config.num_attention_heads		,    device=__UpperCAmelCase       )
							return {
							    "input_ids": input_ids,
							    "decoder_input_ids": decoder_input_ids,
							    "attention_mask": attention_mask,
							    "decoder_attention_mask": decoder_attention_mask,
							    "head_mask": head_mask,
							    "decoder_head_mask": decoder_head_mask,
							    "cross_attn_head_mask": cross_attn_head_mask,
							}

	def 		UpperCAmelCase  (							self       )->					List[str]:







							'''simple docstring'''





							lowerCAmelCase__       =						ids_tensor([self.batch_size, self.encoder_seq_length]		,    self.vocab_size       )
							lowerCAmelCase__       =						ids_tensor([self.batch_size, self.decoder_seq_length]		,    self.vocab_size       )

							# we need to clamp the input ids here to avoid having pad token in between
							# this is because for NllbMoe the position_ids are prepared such that
							# all pad tokens have pos id = 2 and rest are between 2..seq_length
							# and the seq_length here is seq_length - num_pad_tokens
							# but when using past, there is no way of knowing if the past input ids had
							# pad tokens in them, which results in incorrect seq_lenth and which in turn results in
							# position_ids being off by num_pad_tokens in past input
							lowerCAmelCase__       =						input_ids.clamp(self.pad_token_id + 1       )
							lowerCAmelCase__       =						decoder_input_ids.clamp(self.pad_token_id + 1       )

							lowerCAmelCase__       =						self.get_config()
							lowerCAmelCase__       =						config.num_attention_heads
							lowerCAmelCase__       =						self.prepare_inputs_dict(__UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase       )
							return config, input_dict

	def 		UpperCAmelCase  (							self       )->					int:







							'''simple docstring'''





							lowerCAmelCase__      ,							lowerCAmelCase__       =						self.prepare_config_and_inputs()
							return config, inputs_dict

	def 		UpperCAmelCase  (							self       )->					Tuple:







							'''simple docstring'''





							return TaConfig(
							    vocab_size=166		,    d_model=self.hidden_size		,    d_ff=self.d_ff		,    d_kv=self.hidden_size // self.num_attention_heads		,    num_layers=self.num_hidden_layers		,    num_decoder_layers=self.decoder_layers		,    num_heads=self.num_attention_heads		,    relative_attention_num_buckets=self.relative_attention_num_buckets		,    dropout_rate=self.dropout_rate		,    initializer_factor=self.initializer_factor		,    eos_token_id=self.eos_token_id		,    bos_token_id=self.pad_token_id		,    pad_token_id=self.pad_token_id		,    decoder_start_token_id=self.decoder_start_token_id		,    )

	def 		UpperCAmelCase  (							self       )->					Tuple:







							'''simple docstring'''





							return TaConfig(
							    vocab_size=self.vocab_size		,    d_model=self.hidden_size		,    d_ff=self.d_ff		,    d_kv=self.hidden_size // self.num_attention_heads		,    num_layers=self.num_hidden_layers		,    num_decoder_layers=self.decoder_layers		,    num_heads=self.num_attention_heads		,    relative_attention_num_buckets=self.relative_attention_num_buckets		,    dropout_rate=self.dropout_rate		,    initializer_factor=self.initializer_factor		,    eos_token_id=self.eos_token_id		,    bos_token_id=self.pad_token_id		,    pad_token_id=self.pad_token_id		,    decoder_start_token_id=self.decoder_start_token_id		,    )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    )->					Union[str, Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						UMTaModel(config=__UpperCAmelCase       )
							model.to(__UpperCAmelCase       )
							model.eval()
							lowerCAmelCase__       =						model(
							    input_ids=__UpperCAmelCase		,    decoder_input_ids=__UpperCAmelCase		,    attention_mask=__UpperCAmelCase		,    decoder_attention_mask=__UpperCAmelCase		,    )
							lowerCAmelCase__       =						model(input_ids=__UpperCAmelCase		,    decoder_input_ids=__UpperCAmelCase       )
							lowerCAmelCase__       =						result.last_hidden_state
							lowerCAmelCase__       =						result.past_key_values
							lowerCAmelCase__       =						result.encoder_last_hidden_state

							self.parent.assertEqual(encoder_output.size()		,    (self.batch_size, self.encoder_seq_length, self.hidden_size)       )
							self.parent.assertEqual(decoder_output.size()		,    (self.batch_size, self.decoder_seq_length, self.hidden_size)       )
							# There should be `num_layers` key value embeddings stored in decoder_past
							self.parent.assertEqual(len(__UpperCAmelCase       )		,    config.num_layers       )
							# There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
							self.parent.assertEqual(len(decoder_past[0]       )		,    4       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    __UpperCAmelCase		,    )->					Union[str, Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						UMTaModel(config=__UpperCAmelCase       ).get_decoder().to(__UpperCAmelCase       ).eval()
							# first forward pass
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    use_cache=__UpperCAmelCase       )
							lowerCAmelCase__       =						model(__UpperCAmelCase       )
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    use_cache=__UpperCAmelCase       )

							self.parent.assertTrue(len(__UpperCAmelCase       ) == len(__UpperCAmelCase       )       )
							self.parent.assertTrue(len(__UpperCAmelCase       ) == len(__UpperCAmelCase       ) + 1       )

							lowerCAmelCase__      ,							lowerCAmelCase__       =						outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							lowerCAmelCase__       =						ids_tensor((self.batch_size, 1)		,    config.vocab_size       )

							# append to next input_ids and
							lowerCAmelCase__       =						torch.cat([input_ids, next_tokens]		,    dim=-1       )

							lowerCAmelCase__       =						model(__UpperCAmelCase       )["last_hidden_state"]
							lowerCAmelCase__       =						model(__UpperCAmelCase		,    past_key_values=__UpperCAmelCase       )["last_hidden_state"]

							# select random slice
							lowerCAmelCase__       =						ids_tensor((1,)		,    output_from_past.shape[-1]       ).item()
							lowerCAmelCase__       =						output_from_no_past[:, -1, random_slice_idx].detach()
							lowerCAmelCase__       =						output_from_past[:, 0, random_slice_idx].detach()

							# test that outputs are equal for slice
							self.parent.assertTrue(torch.allclose(__UpperCAmelCase		,    __UpperCAmelCase		,    atol=1E-3       )       )


	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase		,    )->					int:







							'''simple docstring'''





							lowerCAmelCase__       =						UMTaModel(config=__UpperCAmelCase       ).to(__UpperCAmelCase       ).half().eval()
							lowerCAmelCase__       =						model(**__UpperCAmelCase       )["last_hidden_state"]
							self.parent.assertFalse(torch.isnan(__UpperCAmelCase       ).any().item()       )





@require_torch
class      lowercase__     (			_UpperCAmelCase,					_UpperCAmelCase,					_UpperCAmelCase,					unittest.TestCase	):
	a_         =(
	    (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
	)
	a_         =(UMTaForConditionalGeneration,) if is_torch_available() else ()
	a_         =(
	    {
	        """conversational""": UMTaForConditionalGeneration,
	        """feature-extraction""": UMTaModel,
	        """summarization""": UMTaForConditionalGeneration,
	        """text2text-generation""": UMTaForConditionalGeneration,
	        """translation""": UMTaForConditionalGeneration,
	        """question-answering""": UMTaForQuestionAnswering,
	    }
	    if is_torch_available()
	    else {}
	)
	a_         =True
	a_         =False
	a_         =False
	a_         =True
	a_         =True
	# The small UMT5 model needs higher percentages for CPU/MP tests
	a_         =[0.8, 0.9]

	def 		UpperCAmelCase  (							self       )->					Dict:







							'''simple docstring'''





							lowerCAmelCase__       =						UMTaModelTester(self       )

	@unittest.skip("Test has a segmentation fault on torch 1.8.0"       )
	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							lowerCAmelCase__       =						UMTaModel(config_and_inputs[0]       ).to(__UpperCAmelCase       )
							with tempfile.TemporaryDirectory() as tmpdirname:
													torch.onnx.export(
													    __UpperCAmelCase		,    (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])		,    F"{tmpdirname}/t5_test.onnx"		,    export_params=__UpperCAmelCase		,    opset_version=9		,    input_names=["input_ids", "decoder_input_ids"]		,    )

	@unittest.skipIf(torch_device == "cpu"		,    "Cant do half precision"       )
	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model_fpaa_forward(*__UpperCAmelCase       )

	def 		UpperCAmelCase  (							self       )->					Optional[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						["encoder_attentions", "decoder_attentions", "cross_attentions"]
							lowerCAmelCase__       =						self.model_tester.prepare_config_and_inputs()
							lowerCAmelCase__       =						config_and_inputs[0]
							lowerCAmelCase__       =						UMTaForConditionalGeneration(__UpperCAmelCase       ).eval()
							model.to(__UpperCAmelCase       )

							lowerCAmelCase__       =						{
							    "head_mask": torch.zeros(config.num_layers		,    config.num_heads		,    device=__UpperCAmelCase       ),
							    "decoder_head_mask": torch.zeros(config.num_decoder_layers		,    config.num_heads		,    device=__UpperCAmelCase       ),
							    "cross_attn_head_mask": torch.zeros(config.num_decoder_layers		,    config.num_heads		,    device=__UpperCAmelCase       ),
							}

							for attn_name, (name, mask) in zip(__UpperCAmelCase		,    head_masking.items()       ):
													lowerCAmelCase__       =						{name: mask}
													# Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
													if name == "head_mask":
																			lowerCAmelCase__       =						torch.ones(
																			    config.num_decoder_layers		,    config.num_heads		,    device=__UpperCAmelCase       )

													lowerCAmelCase__       =						model.generate(
													    config_and_inputs[1]["input_ids"]		,    num_beams=1		,    max_length=3		,    output_attentions=__UpperCAmelCase		,    return_dict_in_generate=__UpperCAmelCase		,    **__UpperCAmelCase		,    )
													# We check the state of decoder_attentions and cross_attentions just from the last step
													lowerCAmelCase__       =						out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
													self.assertEqual(sum([w.sum().item() for w in attn_weights]       )		,    0.0       )


	@unittest.skip("Does not work on the tiny model as we keep hitting edge cases."       )
	def 		UpperCAmelCase  (							self       )->					Any:







							'''simple docstring'''





							pass





@require_torch
@require_sentencepiece
@require_tokenizers
class      lowercase__     (			unittest.TestCase	):

	@slow
	@unittest.skip(
	    "Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged"       )
	def 		UpperCAmelCase  (							self       )->					List[Any]:







							'''simple docstring'''





							lowerCAmelCase__       =						UMTaForConditionalGeneration.from_pretrained("google/umt5-small"		,    return_dict=__UpperCAmelCase       ).to(__UpperCAmelCase       )
							lowerCAmelCase__       =						AutoTokenizer.from_pretrained("google/umt5-small"		,    use_fast=__UpperCAmelCase		,    legacy=__UpperCAmelCase       )
							lowerCAmelCase__       =						[
							    "Bonjour monsieur <extra_id_0> bien <extra_id_1>.",
							    "No se como puedo <extra_id_0>.",
							    "This is the reason why we <extra_id_0> them.",
							    "The <extra_id_0> walks in <extra_id_1>, seats",
							    "A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.",
							]
							lowerCAmelCase__       =						tokenizer(__UpperCAmelCase		,    return_tensors="pt"		,    padding=__UpperCAmelCase       ).input_ids
							# fmt: off
							lowerCAmelCase__       =						torch.tensor(
							    [
							        [ 38530, 210703, 256299, 1410, 256298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [   826, 321, 671, 25922, 256299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [  1460, 339, 312, 19014, 10620, 758, 256299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
							        [   517, 256299, 14869, 281, 301, 256298, 275, 119983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [   320, 256299, 14869, 281, 2234, 289, 2275, 333,61391, 289, 256298, 543, 256297, 168714, 329, 256296,274, 1],
							    ]       )
							# fmt: on
							torch.testing.assert_allclose(__UpperCAmelCase		,    __UpperCAmelCase       )

							lowerCAmelCase__       =						model.generate(input_ids.to(__UpperCAmelCase       )       )
							lowerCAmelCase__       =						[
							    "<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>",
							    "<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							    "<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							    "<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							    "<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							]
							lowerCAmelCase__       =						tokenizer.batch_decode(__UpperCAmelCase       )
							self.assertEqual(__UpperCAmelCase		,    __UpperCAmelCase       )














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 import BertTokenizer


a_							=    logging.get_logger(__name__)

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

a_							=    {
    '''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'''
        ),
    },
}
a_							=    {
    '''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'''
        ),
    },
}
a_							=    {
    '''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'''
        ),
    },
}

a_							=    {
    '''facebook/dpr-ctx_encoder-single-nq-base''': 512,
    '''facebook/dpr-ctx_encoder-multiset-base''': 512,
}
a_							=    {
    '''facebook/dpr-question_encoder-single-nq-base''': 512,
    '''facebook/dpr-question_encoder-multiset-base''': 512,
}
a_							=    {
    '''facebook/dpr-reader-single-nq-base''': 512,
    '''facebook/dpr-reader-multiset-base''': 512,
}


a_							=    {
    '''facebook/dpr-ctx_encoder-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-ctx_encoder-multiset-base''': {'''do_lower_case''': True},
}
a_							=    {
    '''facebook/dpr-question_encoder-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-question_encoder-multiset-base''': {'''do_lower_case''': True},
}
a_							=    {
    '''facebook/dpr-reader-single-nq-base''': {'''do_lower_case''': True},
    '''facebook/dpr-reader-multiset-base''': {'''do_lower_case''': True},
}


class      lowercase__     (			_UpperCAmelCase	):
	a_         =VOCAB_FILES_NAMES
	a_         =CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
	a_         =CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	a_         =CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION





class      lowercase__     (			_UpperCAmelCase	):
	a_         =VOCAB_FILES_NAMES
	a_         =QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
	a_         =QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	a_         =QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION





a_							=    collections.namedtuple(
    '''DPRSpanPrediction''', ['''span_score''', '''relevance_score''', '''doc_id''', '''start_index''', '''end_index''', '''text''']
)

a_							=    collections.namedtuple('''DPRReaderOutput''', ['''start_logits''', '''end_logits''', '''relevance_logits'''])


a_							=    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)

    Returns:
        `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      lowercase__     :

	def __call__(							self		,    __UpperCAmelCase		,    __UpperCAmelCase = None		,    __UpperCAmelCase = None		,    __UpperCAmelCase = False		,    __UpperCAmelCase = False		,    __UpperCAmelCase = None		,    __UpperCAmelCase = None		,    __UpperCAmelCase = None		,    **__UpperCAmelCase		,    )->					BatchEncoding:







							'''simple docstring'''





							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
							if len(__UpperCAmelCase       ) != len(__UpperCAmelCase       ):
													raise ValueError(
													    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 = 16		,    __UpperCAmelCase = 64		,    __UpperCAmelCase = 4		,    )->					List[DPRSpanPrediction]:







							'''simple docstring'''





							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]:







							'''simple docstring'''





							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:
													if start_index > end_index:
																			raise ValueError(F"Wrong span indices: [{start_index}:{end_index}]"       )
													lowerCAmelCase__       =						end_index - start_index + 1
													if length > max_answer_length:
																			raise ValueError(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(_UpperCAmelCase	)
class      lowercase__     (			_UpperCAmelCase,					_UpperCAmelCase	):
	a_         =VOCAB_FILES_NAMES
	a_         =READER_PRETRAINED_VOCAB_FILES_MAP
	a_         =READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	a_         =READER_PRETRAINED_INIT_CONFIGURATION
	a_         =["""input_ids""", """attention_mask"""]






