





import collections
from typing import List, Optional, Union

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


UpperCamelCase__   : str					  = logging.get_logger(__name__)

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

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

UpperCamelCase__   : Optional[Any]					  = {
    """facebook/dpr-ctx_encoder-single-nq-base""": 5_12,
    """facebook/dpr-ctx_encoder-multiset-base""": 5_12,
}
UpperCamelCase__   : List[str]					  = {
    """facebook/dpr-question_encoder-single-nq-base""": 5_12,
    """facebook/dpr-question_encoder-multiset-base""": 5_12,
}
UpperCamelCase__   : Tuple					  = {
    """facebook/dpr-reader-single-nq-base""": 5_12,
    """facebook/dpr-reader-multiset-base""": 5_12,
}


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







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					lowerCamelCase     :     Optional[Any]			    =       VOCAB_FILES_NAMES
					lowerCamelCase     :     int			    =       CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
					lowerCamelCase     :     List[str]			    =       CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					lowerCamelCase     :     List[Any]			    =       CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION
					lowerCamelCase     :     int			    =       DPRContextEncoderTokenizer







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					lowerCamelCase     :     List[Any]			    =       VOCAB_FILES_NAMES
					lowerCamelCase     :     int			    =       QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
					lowerCamelCase     :     List[Any]			    =       QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					lowerCamelCase     :     Dict			    =       QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION
					lowerCamelCase     :     Tuple			    =       DPRQuestionEncoderTokenizer


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

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


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

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

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

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

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

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

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

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

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

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







@add_start_docstrings(A_  )
class   _UpperCamelCase  :







					'''simple docstring'''

					def __call__(						self    :  Union[str, Any] ,					__lowercase    :  Dict ,					__lowercase    :  Optional[Any] = None ,					__lowercase    :  Optional[Any] = None ,					__lowercase    :  Optional[Any] = False ,					__lowercase    :  int = False ,					__lowercase    :  Tuple = None ,					__lowercase    :  List[str] = None ,					__lowercase    :  Dict = None ,					**__lowercase    :  Union[str, Any] ,					):






						'''simple docstring'''







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





					def        SCREAMING_SNAKE_CASE							(						self    :  Dict ,					__lowercase    :  int ,					__lowercase    :  str ,					__lowercase    :  Tuple = 16 ,					__lowercase    :  Any = 64 ,					__lowercase    :  Any = 4 ,					):






						'''simple docstring'''







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

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





					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any] ,					__lowercase    :  Tuple ,					__lowercase    :  Optional[Any] ,					__lowercase    :  str ,					__lowercase    :  List[str] ,					):






						'''simple docstring'''







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







@add_end_docstrings(A_  )
class   _UpperCamelCase  (       A_						,       A_  ):







					'''simple docstring'''

					lowerCamelCase     :     Union[str, Any]			    =       VOCAB_FILES_NAMES
					lowerCamelCase     :     List[str]			    =       READER_PRETRAINED_VOCAB_FILES_MAP
					lowerCamelCase     :     str			    =       READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					lowerCamelCase     :     str			    =       READER_PRETRAINED_INIT_CONFIGURATION
					lowerCamelCase     :     Tuple			    =       ['input_ids', 'attention_mask']
					lowerCamelCase     :     List[str]			    =       DPRReaderTokenizer









import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases




def    A_(			A				,      A   ):
	assert isinstance(A				,      A   )
	assert dataset.num_rows == 4
	assert dataset.num_columns == 1
	assert dataset.column_names == ["text"]
	for feature, expected_dtype in expected_features.items():
		assert dataset.features[feature].dtype == expected_dtype




@pytest.mark.parametrize("""keep_in_memory"""				,      [False, True]   )
def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     tmp_path / """cache"""
	UpperCAmelCase_								=     {"""text""": """string"""}
	with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
		UpperCAmelCase_								=     TextDatasetReader(A				,      cache_dir=A				,      keep_in_memory=A   ).read()
	_check_text_dataset(A				,      A   )




@pytest.mark.parametrize(
    """features"""				,      [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,      )
def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     tmp_path / """cache"""
	UpperCAmelCase_								=     {"""text""": """string"""}
	UpperCAmelCase_								=     features.copy() if features else default_expected_features
	UpperCAmelCase_								=     (
	    Features({feature: Value(A   ) for feature, dtype in features.items()}   ) if features is not None else None
	)
	UpperCAmelCase_								=     TextDatasetReader(A				,      features=A				,      cache_dir=A   ).read()
	_check_text_dataset(A				,      A   )




@pytest.mark.parametrize("""split"""				,      [None, NamedSplit("""train"""   ), """train""", """test"""]   )
def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     tmp_path / """cache"""
	UpperCAmelCase_								=     {"""text""": """string"""}
	UpperCAmelCase_								=     TextDatasetReader(A				,      cache_dir=A				,      split=A   ).read()
	_check_text_dataset(A				,      A   )
	assert dataset.split == split if split else "train"




@pytest.mark.parametrize("""path_type"""				,      [str, list]   )
def    A_(			A				,      A				,      A   ):
	if issubclass(A				,      A   ):
		UpperCAmelCase_								=     text_path
	elif issubclass(A				,      A   ):
		UpperCAmelCase_								=     [text_path]
	UpperCAmelCase_								=     tmp_path / """cache"""
	UpperCAmelCase_								=     {"""text""": """string"""}
	UpperCAmelCase_								=     TextDatasetReader(A				,      cache_dir=A   ).read()
	_check_text_dataset(A				,      A   )




def    A_(			A				,      A				,      A=("train",)   ):
	assert isinstance(A				,      A   )
	for split in splits:
		UpperCAmelCase_								=     dataset_dict[split]
		assert dataset.num_rows == 4
		assert dataset.num_columns == 1
		assert dataset.column_names == ["text"]
		for feature, expected_dtype in expected_features.items():
			assert dataset.features[feature].dtype == expected_dtype




@pytest.mark.parametrize("""keep_in_memory"""				,      [False, True]   )
def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     tmp_path / """cache"""
	UpperCAmelCase_								=     {"""text""": """string"""}
	with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
		UpperCAmelCase_								=     TextDatasetReader({"""train""": text_path}				,      cache_dir=A				,      keep_in_memory=A   ).read()
	_check_text_datasetdict(A				,      A   )




@pytest.mark.parametrize(
    """features"""				,      [
        None,
        {"""text""": """string"""},
        {"""text""": """int32"""},
        {"""text""": """float32"""},
    ]				,      )
def    A_(			A				,      A				,      A   ):
	UpperCAmelCase_								=     tmp_path / """cache"""
	# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
	UpperCAmelCase_								=     {"""text""": """string"""}
	UpperCAmelCase_								=     features.copy() if features else default_expected_features
	UpperCAmelCase_								=     (
	    Features({feature: Value(A   ) for feature, dtype in features.items()}   ) if features is not None else None
	)
	UpperCAmelCase_								=     TextDatasetReader({"""train""": text_path}				,      features=A				,      cache_dir=A   ).read()
	_check_text_datasetdict(A				,      A   )




@pytest.mark.parametrize("""split"""				,      [None, NamedSplit("""train"""   ), """train""", """test"""]   )
def    A_(			A				,      A				,      A   ):
	if split:
		UpperCAmelCase_								=     {split: text_path}
	else:
		UpperCAmelCase_								=     """train"""
		UpperCAmelCase_								=     {"""train""": text_path, """test""": text_path}
	UpperCAmelCase_								=     tmp_path / """cache"""
	UpperCAmelCase_								=     {"""text""": """string"""}
	UpperCAmelCase_								=     TextDatasetReader(A				,      cache_dir=A   ).read()
	_check_text_datasetdict(A				,      A				,      splits=list(path.keys()   )   )
	assert all(dataset[split].split == split for split in path.keys()   )


