



"""simple docstring"""





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


lowerCamelCase       	=						logging.get_logger(__name__)

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

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


lowerCamelCase       	=						{
    """facebook/dpr-ctx_encoder-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-ctx_encoder-multiset-base""": {"""do_lower_case""": True},
}
lowerCamelCase       	=						{
    """facebook/dpr-question_encoder-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-question_encoder-multiset-base""": {"""do_lower_case""": True},
}
lowerCamelCase       	=						{
    """facebook/dpr-reader-single-nq-base""": {"""do_lower_case""": True},
    """facebook/dpr-reader-multiset-base""": {"""do_lower_case""": True},
}
class 				lowercase__       (   SCREAMING_SNAKE_CASE				):







       '''simple docstring'''


       UpperCamelCase						=					VOCAB_FILES_NAMES
       UpperCamelCase						=					CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
       UpperCamelCase						=					CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       UpperCamelCase						=					CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION



class 				lowercase__       (   SCREAMING_SNAKE_CASE				):







       '''simple docstring'''


       UpperCamelCase						=					VOCAB_FILES_NAMES
       UpperCamelCase						=					QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
       UpperCamelCase						=					QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       UpperCamelCase						=					QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION



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

lowerCamelCase       	=						collections.namedtuple("""DPRReaderOutput""", ["""start_logits""", """end_logits""", """relevance_logits"""])


lowerCamelCase       	=						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(SCREAMING_SNAKE_CASE				)
class 				lowercase__       :







       '''simple docstring'''


       def __call__(							self      :  Union[str, Any]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  Optional[str] = None		,    _UpperCAmelCase      :  Optional[str] = None		,    _UpperCAmelCase      :  Union[bool, str] = False		,    _UpperCAmelCase      :  Union[bool, str] = False		,    _UpperCAmelCase      :  Optional[int] = None		,    _UpperCAmelCase      :  Optional[Union[str, TensorType]] = None		,    _UpperCAmelCase      :  Optional[bool] = None		,    **_UpperCAmelCase      :  Tuple		,    )			->						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:
                   UpperCAmelCase_      			=							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		,    )
             UpperCAmelCase_      			=							titles if not isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ) else [titles]
             UpperCAmelCase_      			=							texts if not isinstance(_UpperCAmelCase		,    _UpperCAmelCase   ) else [texts]
             UpperCAmelCase_      			=							len(_UpperCAmelCase   )
             UpperCAmelCase_      			=							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."""   )
             UpperCAmelCase_      			=							super().__call__(_UpperCAmelCase		,    _UpperCAmelCase		,    padding=_UpperCAmelCase		,    truncation=_UpperCAmelCase   )["input_ids"]
             UpperCAmelCase_      			=							super().__call__(_UpperCAmelCase		,    add_special_tokens=_UpperCAmelCase		,    padding=_UpperCAmelCase		,    truncation=_UpperCAmelCase   )["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(_UpperCAmelCase		,    _UpperCAmelCase   )
                 ]
             }
             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(_UpperCAmelCase		,    padding=_UpperCAmelCase		,    max_length=_UpperCAmelCase		,    return_tensors=_UpperCAmelCase   )



       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  BatchEncoding		,    _UpperCAmelCase      :  DPRReaderOutput		,    _UpperCAmelCase      :  int = 16		,    _UpperCAmelCase      :  int = 64		,    _UpperCAmelCase      :  int = 4		,    )			->						List[DPRSpanPrediction]:


             '''simple docstring'''




             UpperCAmelCase_      			=							reader_input["input_ids"]
             UpperCAmelCase_			,   UpperCAmelCase_			,   UpperCAmelCase_      			=							reader_output[:3]
             UpperCAmelCase_      			=							len(_UpperCAmelCase   )
             UpperCAmelCase_      			=							sorted(range(_UpperCAmelCase   )		,    reverse=_UpperCAmelCase		,    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(_UpperCAmelCase   )

                   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=_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 					lowercase__    (							self      :  List[str]		,    _UpperCAmelCase      :  List[int]		,    _UpperCAmelCase      :  List[int]		,    _UpperCAmelCase      :  int		,    _UpperCAmelCase      :  int		,    )			->						List[DPRSpanPrediction]:


             '''simple docstring'''




             UpperCAmelCase_      			=							[]
             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)   )
             UpperCAmelCase_      			=							sorted(_UpperCAmelCase		,    key=lambda _UpperCAmelCase   : x[1]		,    reverse=_UpperCAmelCase   )
             UpperCAmelCase_      			=							[]
             for (start_index, end_index), score in scores:
                   if start_index > end_index:
                         raise ValueError(F"""Wrong span indices: [{start_index}:{end_index}]"""   )
                   UpperCAmelCase_      			=							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(SCREAMING_SNAKE_CASE				)
class 				lowercase__       (   SCREAMING_SNAKE_CASE    , SCREAMING_SNAKE_CASE				):







       '''simple docstring'''


       UpperCamelCase						=					VOCAB_FILES_NAMES
       UpperCamelCase						=					READER_PRETRAINED_VOCAB_FILES_MAP
       UpperCamelCase						=					READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       UpperCamelCase						=					READER_PRETRAINED_INIT_CONFIGURATION
       UpperCamelCase						=					['''input_ids''', '''attention_mask''']









"""simple docstring"""





# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
# hack it in for now:
import sys
from pathlib import Path


lowerCamelCase       	=						Path(__file__).resolve().parents[3] / """src"""
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(42)

lowerCamelCase       	=						{"""base""": """patrickvonplaten/wav2vec2_tiny_random""", """robust""": """patrickvonplaten/wav2vec2_tiny_random_robust"""}

lowerCamelCase       	=						"""zero2"""
lowerCamelCase       	=						"""zero3"""
lowerCamelCase       	=						[ZEROa, ZEROa]
def        a__     (				lowerCAmelCase__     ,      lowerCAmelCase__     ,      lowerCAmelCase__				):
      # customize the test name generator function as we want both params to appear in the sub-test
      # name, as by default it shows only the first param
      UpperCAmelCase_      			=							parameterized.to_safe_name("_".join(str(lowerCAmelCase__				) for x in param.args				)				)
      return f"""{func.__name__}_{param_based_name}"""


# Cartesian-product of zero stages with models to test
lowerCamelCase       	=						list(itertools.product(stages, models.keys()))


@slow
@require_deepspeed
@require_torch_gpu
class 				lowercase__       (   SCREAMING_SNAKE_CASE				):







       '''simple docstring'''


       @parameterized.expand(_UpperCAmelCase		,    name_func=_UpperCAmelCase   )
       def 					lowercase__    (							self      :  Dict		,    _UpperCAmelCase      :  int		,    _UpperCAmelCase      :  Any   )			->						Optional[int]:


             '''simple docstring'''




             self.run_and_check(
                 stage=_UpperCAmelCase		,    model=_UpperCAmelCase		,    distributed=_UpperCAmelCase		,    fpaa=_UpperCAmelCase		,    )



       @require_torch_multi_gpu
       @parameterized.expand(_UpperCAmelCase		,    name_func=_UpperCAmelCase   )
       def 					lowercase__    (							self      :  Any		,    _UpperCAmelCase      :  Optional[Any]		,    _UpperCAmelCase      :  str   )			->						Any:


             '''simple docstring'''




             self.run_and_check(
                 stage=_UpperCAmelCase		,    model=_UpperCAmelCase		,    distributed=_UpperCAmelCase		,    fpaa=_UpperCAmelCase		,    )



       @parameterized.expand(_UpperCAmelCase		,    name_func=_UpperCAmelCase   )
       def 					lowercase__    (							self      :  Optional[int]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  List[str]   )			->						Union[str, Any]:


             '''simple docstring'''




             self.run_and_check(
                 stage=_UpperCAmelCase		,    model=_UpperCAmelCase		,    distributed=_UpperCAmelCase		,    fpaa=_UpperCAmelCase		,    )



       @require_torch_multi_gpu
       @parameterized.expand(_UpperCAmelCase		,    name_func=_UpperCAmelCase   )
       def 					lowercase__    (							self      :  int		,    _UpperCAmelCase      :  int		,    _UpperCAmelCase      :  Optional[int]   )			->						Optional[int]:


             '''simple docstring'''




             self.run_and_check(
                 stage=_UpperCAmelCase		,    model=_UpperCAmelCase		,    distributed=_UpperCAmelCase		,    fpaa=_UpperCAmelCase		,    )



       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  Tuple   )			->						Optional[int]:


             '''simple docstring'''




             pass



       def 					lowercase__    (							self      :  Optional[Any]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  int = 10		,    _UpperCAmelCase      :  bool = True		,    _UpperCAmelCase      :  bool = True		,    _UpperCAmelCase      :  bool = True		,    )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							models[model]

             UpperCAmelCase_      			=							self.run_trainer(
                 stage=_UpperCAmelCase		,    model_name=_UpperCAmelCase		,    eval_steps=_UpperCAmelCase		,    num_train_epochs=1		,    distributed=_UpperCAmelCase		,    fpaa=_UpperCAmelCase		,    )

             self.do_checks(_UpperCAmelCase   )

             return output_dir



       def 					lowercase__    (							self      :  Union[str, Any]		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  str		,    _UpperCAmelCase      :  int = 10		,    _UpperCAmelCase      :  int = 1		,    _UpperCAmelCase      :  bool = True		,    _UpperCAmelCase      :  bool = True		,    )			->						Tuple:


             '''simple docstring'''




             UpperCAmelCase_      			=							self.get_auto_remove_tmp_dir("./xxx"		,    after=_UpperCAmelCase   )
             UpperCAmelCase_      			=							F"""
            --model_name_or_path {model_name}
            --dataset_name hf-internal-testing/librispeech_asr_dummy
            --dataset_config_name clean
            --train_split_name validation
            --validation_split_name validation
            --output_dir {output_dir}
            --num_train_epochs {str(_UpperCAmelCase   )}
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 2
            --evaluation_strategy steps
            --learning_rate 5e-4
            --warmup_steps 8
            --orthography timit
            --preprocessing_num_workers 1
            --group_by_length
            --freeze_feature_extractor
            --report_to none
            --save_steps 0
            --eval_steps {eval_steps}
            --report_to none
        """.split()

             if fpaa:
                   args.extend(["--fp16"]   )

             # currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
             # hence the separate config files
             UpperCAmelCase_      			=							F"""--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json""".split()
             UpperCAmelCase_      			=							[F"""{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py"""]
             UpperCAmelCase_      			=							self.get_launcher(_UpperCAmelCase   )

             UpperCAmelCase_      			=							launcher + script + args + ds_args
             # keep for quick debug
             # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
             execute_subprocess_async(_UpperCAmelCase		,    env=self.get_env()   )

             return output_dir







       def 					lowercase__    (							self      :  List[Any]		,    _UpperCAmelCase      :  int=False   )			->						Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_      			=							min(2		,    get_gpu_count()   ) if distributed else 1
             return F"""deepspeed --num_nodes 1 --num_gpus {num_gpus}""".split()




