




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

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

UpperCAmelCase_									=	{
    """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_									=	{
    """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_									=	{
    """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_									=	{
    """facebook/dpr-ctx_encoder-single-nq-base""": 512,
    """facebook/dpr-ctx_encoder-multiset-base""": 512,
}
UpperCAmelCase_									=	{
    """facebook/dpr-question_encoder-single-nq-base""": 512,
    """facebook/dpr-question_encoder-multiset-base""": 512,
}
UpperCAmelCase_									=	{
    """facebook/dpr-reader-single-nq-base""": 512,
    """facebook/dpr-reader-multiset-base""": 512,
}


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





class     lowerCamelCase__(					UpperCAmelCase__):
      UpperCAmelCase__      :	Any    		=   VOCAB_FILES_NAMES
      UpperCAmelCase__      :	Any    		=   CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
      UpperCAmelCase__      :	int    		=   CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      UpperCAmelCase__      :	Union[str, Any]    		=   CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION
      UpperCAmelCase__      :	int    		=   DPRContextEncoderTokenizer





class     lowerCamelCase__(					UpperCAmelCase__):
      UpperCAmelCase__      :	Any    		=   VOCAB_FILES_NAMES
      UpperCAmelCase__      :	Any    		=   QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
      UpperCAmelCase__      :	Union[str, Any]    		=   QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      UpperCAmelCase__      :	List[str]    		=   QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION
      UpperCAmelCase__      :	Optional[Any]    		=   DPRQuestionEncoderTokenizer



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

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


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

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

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

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

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

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

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

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

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

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





@add_start_docstrings(UpperCAmelCase__)
class     lowerCamelCase__:
      def __call__(				self: Optional[Any]    ,       UpperCamelCase_: Optional[int]    ,       UpperCamelCase_: Optional[int] = None    ,       UpperCamelCase_: List[str] = None    ,       UpperCamelCase_: Optional[Any] = False    ,       UpperCamelCase_: str = False    ,       UpperCamelCase_: Optional[Any] = None    ,       UpperCamelCase_: List[str] = None    ,       UpperCamelCase_: Union[str, Any] = None    ,       **UpperCamelCase_: Dict    ,       ):
             if titles is None and texts is None:
                    return super().__call__(
                        _a    ,       padding=_a    ,       truncation=_a    ,       max_length=_a    ,       return_tensors=_a    ,       return_attention_mask=_a    ,       **_a    ,       )
             elif titles is None or texts is None:
                    __lowerCamelCase   =						titles if texts is None else texts
                    return super().__call__(
                        _a    ,       _a    ,       padding=_a    ,       truncation=_a    ,       max_length=_a    ,       return_tensors=_a    ,       return_attention_mask=_a    ,       **_a    ,       )
             __lowerCamelCase   =						titles if not isinstance(_a    ,       _a   ) else [titles]
             __lowerCamelCase   =						texts if not isinstance(_a    ,       _a   ) else [texts]
             __lowerCamelCase   =						len(_a   )
             __lowerCamelCase   =						questions if not isinstance(_a    ,       _a   ) else [questions] * n_passages
             assert len(_a   ) == len(
                 _a   ), F'There should be as many titles than texts but got {len(_a   )} titles and {len(_a   )} texts.'
             __lowerCamelCase   =						super().__call__(_a    ,       _a    ,       padding=_a    ,       truncation=_a   )["""input_ids"""]
             __lowerCamelCase   =						super().__call__(_a    ,       add_special_tokens=_a    ,       padding=_a    ,       truncation=_a   )["""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(_a    ,       _a   )
                 ]
             }
             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(_a    ,       padding=_a    ,       max_length=_a    ,       return_tensors=_a   )
      def 			lowerCAmelCase__						(				self: Any    ,       UpperCamelCase_: int    ,       UpperCamelCase_: List[Any]    ,       UpperCamelCase_: List[str] = 16    ,       UpperCamelCase_: int = 64    ,       UpperCamelCase_: Optional[int] = 4    ,       ):
             __lowerCamelCase   =						reader_input["""input_ids"""]
             __lowerCamelCase,   __lowerCamelCase,   __lowerCamelCase   =						reader_output[:3]
             __lowerCamelCase   =						len(_a   )
             __lowerCamelCase   =						sorted(range(_a   )    ,       reverse=_a    ,       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(_a   )

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




      def 			lowerCAmelCase__						(				self: Dict    ,       UpperCamelCase_: Optional[Any]    ,       UpperCamelCase_: Optional[int]    ,       UpperCamelCase_: Optional[int]    ,       UpperCamelCase_: List[str]    ,       ):
             __lowerCamelCase   =						[]
             for start_index, start_score in enumerate(_a   ):
                    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(_a    ,       key=lambda UpperCamelCase_   : x[1]    ,       reverse=_a   )
             __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(_a   ) == top_spans:
                           break
             return chosen_span_intervals





@add_end_docstrings(UpperCAmelCase__)
class     lowerCamelCase__(					UpperCAmelCase__ ,     UpperCAmelCase__):
      UpperCAmelCase__      :	Optional[Any]    		=   VOCAB_FILES_NAMES
      UpperCAmelCase__      :	str    		=   READER_PRETRAINED_VOCAB_FILES_MAP
      UpperCAmelCase__      :	List[str]    		=   READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      UpperCAmelCase__      :	List[str]    		=   READER_PRETRAINED_INIT_CONFIGURATION
      UpperCAmelCase__      :	Optional[int]    		=   ['input_ids', 'attention_mask']
      UpperCAmelCase__      :	Union[str, Any]    		=   DPRReaderTokenizer







"""simple docstring"""




import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					dataset
							lowerCamelCase			      =					process
							lowerCamelCase			      =					params





				def __len__(       self ):





							"""simple docstring"""
							return len(self.dataset )





				def __getitem__(       self		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					self.dataset[i]
							lowerCamelCase			      =					self.process(_a		, **self.params )
							return processed






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a		, _a		, _a=None ):





							"""simple docstring"""
							lowerCamelCase			      =					loader
							lowerCamelCase			      =					infer
							lowerCamelCase			      =					params
							if loader_batch_size == 1:
										# Let's spare some time by deactivating altogether
										lowerCamelCase			      =					None
							lowerCamelCase			      =					loader_batch_size

							# Internal bookkeeping
							lowerCamelCase			      =					None
							lowerCamelCase			      =					None





				def __len__(       self ):





							"""simple docstring"""
							return len(self.loader )





				def __iter__(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					iter(self.loader )
							return self





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							if isinstance(self._loader_batch_data		, torch.Tensor ):
										# Batch data is simple tensor, just fetch the slice
										lowerCamelCase			      =					self._loader_batch_data[self._loader_batch_index]
							else:
										# Batch data is assumed to be BaseModelOutput (or dict)
										lowerCamelCase			      =					{}
										for k, element in self._loader_batch_data.items():
													if isinstance(_a		, _a ):
																# Convert ModelOutput to tuple first
																lowerCamelCase			      =					element.to_tuple()
																if isinstance(element[0]		, torch.Tensor ):
																			lowerCamelCase			      =					tuple(el[self._loader_batch_index].unsqueeze(0 ) for el in element )
																elif isinstance(element[0]		, np.ndarray ):
																			lowerCamelCase			      =					tuple(np.expand_dims(el[self._loader_batch_index]		, 0 ) for el in element )
																continue
													if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(_a		, _a ):
																# Those are stored as lists of tensors so need specific unbatching.
																if isinstance(element[0]		, torch.Tensor ):
																			lowerCamelCase			      =					tuple(el[self._loader_batch_index].unsqueeze(0 ) for el in element )
																elif isinstance(element[0]		, np.ndarray ):
																			lowerCamelCase			      =					tuple(np.expand_dims(el[self._loader_batch_index]		, 0 ) for el in element )
																continue
													if element is None:
																# This can happen for optional data that get passed around
																lowerCamelCase			      =					None
													elif isinstance(element[self._loader_batch_index]		, torch.Tensor ):
																# Take correct batch data, but make it looked like batch_size=1
																# For compatibility with other methods within transformers

																lowerCamelCase			      =					element[self._loader_batch_index].unsqueeze(0 )
													elif isinstance(element[self._loader_batch_index]		, np.ndarray ):
																# Take correct batch data, but make it looked like batch_size=1
																# For compatibility with other methods within transformers
																lowerCamelCase			      =					np.expand_dims(element[self._loader_batch_index]		, 0 )
													else:
																# This is typically a list, so no need to `unsqueeze`.
																lowerCamelCase			      =					element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
										lowerCamelCase			      =					self._loader_batch_data.__class__(_a )
							self._loader_batch_index += 1
							return result





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
										# We are currently unrolling a batch so we just need to return
										# the current item within a batch
										return self.loader_batch_item()

							# We're out of items within a batch
							lowerCamelCase			      =					next(self.iterator )
							lowerCamelCase			      =					self.infer(_a		, **self.params )
							# We now have a batch of "inferred things".
							if self.loader_batch_size is not None:
										# Try to infer the size of the batch
										if isinstance(_a		, torch.Tensor ):
													lowerCamelCase			      =					processed
										else:
													lowerCamelCase			      =					list(processed.keys() )[0]
													lowerCamelCase			      =					processed[key]
										if isinstance(_a		, _a ):
													lowerCamelCase			      =					len(_a )
										else:
													lowerCamelCase			      =					first_tensor.shape[0]
										if 0 < observed_batch_size < self.loader_batch_size:
													# could be last batch so we can't unroll as many
													# elements.
													lowerCamelCase			      =					observed_batch_size
										# Setting internal index to unwrap the batch
										lowerCamelCase			      =					processed
										lowerCamelCase			      =					0
										return self.loader_batch_item()
							else:
										# We're not unrolling batches
										return processed






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a		, _a		, _a=None ):





							"""simple docstring"""
							super().__init__(_a		, _a		, _a )





				def __iter__(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					iter(self.loader )
							lowerCamelCase			      =					None
							return self





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							if self.subiterator is None:
										lowerCamelCase			      =					self.infer(next(self.iterator )		, **self.params )
							try:
										# Try to return next item
										lowerCamelCase			      =					next(self.subiterator )
							except StopIteration:
										# When a preprocess iterator ends, we can start lookig at the next item
										# ChunkIterator will keep feeding until ALL elements of iterator
										# all have created their subiterator and have been iterating against.
										#
										# Another way to look at it, is we're basically flattening lists of lists
										# into a single list, but with generators
										lowerCamelCase			      =					self.infer(next(self.iterator )		, **self.params )
										lowerCamelCase			      =					next(self.subiterator )
							return processed






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __iter__(       self ):





							"""simple docstring"""
							lowerCamelCase			      =					iter(self.loader )
							return self





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							# Extremely similar to PipelineIterator in its unpacking mechanism
							# BUT, we have an extra required item which is the presence of `is_last`
							# That is because everything is flattened by `PipelineChunkIterator` we
							# need to keep track of how to regroup here in the original `process`
							# boundaries so that `process` and `postprocess` see the same data.

							# This iterator accumulates items (possibly while unbatching) until it
							# its a `is_last` and then just passes it on to the caller.
							lowerCamelCase			      =					False
							lowerCamelCase			      =					[]
							if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
										while self._loader_batch_index < self.loader_batch_size:
													lowerCamelCase			      =					self.loader_batch_item()
													lowerCamelCase			      =					item.pop("""is_last""" )
													accumulator.append(_a )
													if is_last:
																return accumulator

							while not is_last:
										lowerCamelCase			      =					self.infer(next(self.iterator )		, **self.params )
										if self.loader_batch_size is not None:
													if isinstance(_a		, torch.Tensor ):
																lowerCamelCase			      =					processed
													else:
																lowerCamelCase			      =					list(processed.keys() )[0]
																lowerCamelCase			      =					processed[key]
													if isinstance(_a		, _a ):
																lowerCamelCase			      =					len(_a )
													else:
																lowerCamelCase			      =					first_tensor.shape[0]
													if 0 < observed_batch_size < self.loader_batch_size:
																# could be last batch so we can't unroll as many
																# elements.
																lowerCamelCase			      =					observed_batch_size
													lowerCamelCase			      =					processed
													lowerCamelCase			      =					0
													while self._loader_batch_index < self.loader_batch_size:
																lowerCamelCase			      =					self.loader_batch_item()
																lowerCamelCase			      =					item.pop("""is_last""" )
																accumulator.append(_a )
																if is_last:
																			return accumulator
										else:
													lowerCamelCase			      =					processed
													lowerCamelCase			      =					item.pop("""is_last""" )
													accumulator.append(_a )
							return accumulator






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					dataset
							lowerCamelCase			      =					key





				def __len__(       self ):





							"""simple docstring"""
							return len(self.dataset )





				def __getitem__(       self		, _a ):





							"""simple docstring"""
							return self.dataset[i][self.key]






class     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''





				def __init__(       self		, _a		, _a		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					dataset
							lowerCamelCase			      =					keya
							lowerCamelCase			      =					keya





				def __len__(       self ):





							"""simple docstring"""
							return len(self.dataset )





				def __getitem__(       self		, _a ):





							"""simple docstring"""
							return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}



