

import json
import logging
import os
import re
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import datasets
import numpy as np
import torch
import torchaudio
from packaging import version
from torch import nn

import transformers
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaProcessor,
    is_apex_available,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process


if is_apex_available():
      from apex import amp


if version.parse(version.parse(torch.__version__).base_version) >= version.parse("""1.6"""):
      __lowercase				      =						True
      from torch.cuda.amp import autocast

__lowercase				      =						logging.getLogger(__name__)







def 							_lowerCamelCase   (			SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=None   ):
      '''simple docstring'''







      return field(default_factory=lambda: default		,				metadata=SCREAMING_SNAKE_CASE   )

@dataclass
class    _lowercase				:
   _lowercase     :       str  					=    field(
       metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}			)
   _lowercase     :       Optional[str]  					=    field(
       default=__lowerCamelCase,metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'},)
   _lowercase     :       Optional[bool]  					=    field(
       default=__lowerCamelCase,metadata={'help': 'Whether to freeze the feature extractor layers of the model.'}			)
   _lowercase     :       Optional[float]  					=    field(
       default=0.1,metadata={'help': 'The dropout ratio for the attention probabilities.'}			)
   _lowercase     :       Optional[float]  					=    field(
       default=0.1,metadata={'help': 'The dropout ratio for activations inside the fully connected layer.'}			)
   _lowercase     :       Optional[float]  					=    field(
       default=0.1,metadata={
           'help': 'The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler.'
       },)
   _lowercase     :       Optional[float]  					=    field(
       default=0.1,metadata={'help': 'The dropout probabilitiy for all 1D convolutional layers in feature extractor.'},)
   _lowercase     :       Optional[float]  					=    field(
       default=0.05,metadata={
           'help': (
               'Propability of each feature vector along the time axis to be chosen as the start of the vector'
               'span to be masked. Approximately ``mask_time_prob * sequence_length // mask_time_length`` feature'
               'vectors will be masked along the time axis. This is only relevant if ``apply_spec_augment is True``.'
           )
       },)
   _lowercase     :       Optional[float]  					=    field(default=0.0,metadata={'help': 'The LayerDrop probability.'}			)



@dataclass
class    _lowercase				:
   _lowercase     :       Optional[str]  					=    field(
       default=__lowerCamelCase,metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'}			)
   _lowercase     :       Optional[str]  					=    field(
       default='train+validation',metadata={
           'help': 'The name of the training data set split to use (via the datasets library). Defaults to \'train\''
       },)
   _lowercase     :       bool  					=    field(
       default=__lowerCamelCase,metadata={'help': 'Overwrite the cached preprocessed datasets or not.'}			)
   _lowercase     :       Optional[int]  					=    field(
       default=__lowerCamelCase,metadata={'help': 'The number of processes to use for the preprocessing.'},)
   _lowercase     :       Optional[int]  					=    field(
       default=__lowerCamelCase,metadata={
           'help': (
               'For debugging purposes or quicker training, truncate the number of training examples to this '
               'value if set.'
           )
       },)
   _lowercase     :       Optional[int]  					=    field(
       default=__lowerCamelCase,metadata={
           'help': (
               'For debugging purposes or quicker training, truncate the number of validation examples to this '
               'value if set.'
           )
       },)
   _lowercase     :       List[str]  					=    list_field(
       default=[',', '?', '.', '!', '-', ';', ':', '""', '%', '\'', '"', '�'],metadata={'help': 'A list of characters to remove from the transcripts.'},)



@dataclass
class    _lowercase				:
   _lowercase     :       WavaVecaProcessor
   _lowercase     :       Union[bool, str]  					=    True
   _lowercase     :       Optional[int]  					=    None
   _lowercase     :       Optional[int]  					=    None
   _lowercase     :       Optional[int]  					=    None
   _lowercase     :       Optional[int]  					=    None

   def __call__(				self	:    int			,						lowerCamelCase__	:    List[Dict[str, Union[List[int], torch.Tensor]]]       ) ->			Dict[str, torch.Tensor]:
         """simple docstring"""





         A_							       =		[{'''input_values''': feature['''input_values''']} for feature in features]
         A_							       =		[{'''input_ids''': feature['''labels''']} for feature in features]

         A_							       =		self.processor.pad(
             lowerCamelCase__			,						padding=self.padding			,						max_length=self.max_length			,						pad_to_multiple_of=self.pad_to_multiple_of			,						return_tensors='''pt'''			,						)
         A_							       =		self.processor.pad(
             labels=lowerCamelCase__			,						padding=self.padding			,						max_length=self.max_length_labels			,						pad_to_multiple_of=self.pad_to_multiple_of_labels			,						return_tensors='''pt'''			,						)

         # replace padding with -100 to ignore loss correctly
         A_							       =		labels_batch['''input_ids'''].masked_fill(labels_batch.attention_mask.ne(1       )			,						-1_0_0       )

         A_							       =		labels

         return batch



class    _lowercase				(					__lowerCamelCase			):

   def     UpperCamelCase					(				self	:    Optional[Any]			,						lowerCamelCase__	:    nn.Module			,						lowerCamelCase__	:    Dict[str, Union[torch.Tensor, Any]]       ) ->			torch.Tensor:
         """simple docstring"""





         model.train()
         A_							       =		self._prepare_inputs(lowerCamelCase__       )

         if self.use_amp:
               with autocast():
                     A_							       =		self.compute_loss(lowerCamelCase__			,						lowerCamelCase__       )
         else:
               A_							       =		self.compute_loss(lowerCamelCase__			,						lowerCamelCase__       )

         if self.args.n_gpu > 1:
               if model.module.config.ctc_loss_reduction == "mean":
                     A_							       =		loss.mean()
               elif model.module.config.ctc_loss_reduction == "sum":
                     A_							       =		loss.sum() / (inputs['''labels'''] >= 0).sum()
               else:
                     raise ValueError(F"{model.config.ctc_loss_reduction} is not valid. Choose one of ['mean', 'sum']"       )

         if self.args.gradient_accumulation_steps > 1:
               A_							       =		loss / self.args.gradient_accumulation_steps

         if self.use_amp:
               self.scaler.scale(lowerCamelCase__       ).backward()
         elif self.use_apex:
               with amp.scale_loss(lowerCamelCase__			,						self.optimizer       ) as scaled_loss:
                     scaled_loss.backward()
         elif self.deepspeed:
               self.deepspeed.backward(lowerCamelCase__       )
         else:
               loss.backward()

         return loss.detach()







def 							_lowerCamelCase   (			):
      '''simple docstring'''







      A_							       =		HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)   )
      if len(sys.argv   ) == 2 and sys.argv[1].endswith('''.json'''   ):
            # If we pass only one argument to the script and it's the path to a json file,
            # let's parse it to get our arguments.
            A_       ,A_       ,A_							       =		parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]   )   )
      else:
            A_       ,A_       ,A_							       =		parser.parse_args_into_dataclasses()

      # Detecting last checkpoint.
      A_							       =		None
      if os.path.isdir(training_args.output_dir   ) and training_args.do_train and not training_args.overwrite_output_dir:
            A_							       =		get_last_checkpoint(training_args.output_dir   )
            if last_checkpoint is None and len(os.listdir(training_args.output_dir   )   ) > 0:
                  raise ValueError(
                      f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                      '''Use --overwrite_output_dir to overcome.'''   )
            elif last_checkpoint is not None:
                  logger.info(
                      f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                      '''the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'''   )

    # Setup logging
      logging.basicConfig(
          format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''		,				datefmt='''%m/%d/%Y %H:%M:%S'''		,				handlers=[logging.StreamHandler(sys.stdout   )]		,				)
      logger.setLevel(logging.INFO if is_main_process(training_args.local_rank   ) else logging.WARN   )

      # Log on each process the small summary:
      logger.warning(
          f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
          + f"distributed training: {bool(training_args.local_rank != -1   )}, 16-bits training: {training_args.fpaa}"   )
      # Set the verbosity to info of the Transformers logger (on main process only):
      if is_main_process(training_args.local_rank   ):
            transformers.utils.logging.set_verbosity_info()
      logger.info('''Training/evaluation parameters %s'''		,				SCREAMING_SNAKE_CASE   )

      # Set seed before initializing model.
      set_seed(training_args.seed   )

      # Get the datasets:
      A_							       =		datasets.load_dataset(
          '''common_voice'''		,				data_args.dataset_config_name		,				split=data_args.train_split_name   )
      A_							       =		datasets.load_dataset('''common_voice'''		,				data_args.dataset_config_name		,				split='''test'''   )

      # Create and save tokenizer
      A_							       =		f"[{''.join(data_args.chars_to_ignore   )}]"

      def remove_special_characters(SCREAMING_SNAKE_CASE   ):
            A_							       =		re.sub(SCREAMING_SNAKE_CASE		,				''''''		,				batch['''sentence''']   ).lower() + ''' '''
            return batch

      A_							       =		train_dataset.map(SCREAMING_SNAKE_CASE		,				remove_columns=['''sentence''']   )
      A_							       =		eval_dataset.map(SCREAMING_SNAKE_CASE		,				remove_columns=['''sentence''']   )

      def extract_all_chars(SCREAMING_SNAKE_CASE   ):
            A_							       =		''' '''.join(batch['''text''']   )
            A_							       =		list(set(SCREAMING_SNAKE_CASE   )   )
            return {"vocab": [vocab], "all_text": [all_text]}

      A_							       =		train_dataset.map(
          SCREAMING_SNAKE_CASE		,				batched=SCREAMING_SNAKE_CASE		,				batch_size=-1		,				keep_in_memory=SCREAMING_SNAKE_CASE		,				remove_columns=train_dataset.column_names		,				)
      A_							       =		train_dataset.map(
          SCREAMING_SNAKE_CASE		,				batched=SCREAMING_SNAKE_CASE		,				batch_size=-1		,				keep_in_memory=SCREAMING_SNAKE_CASE		,				remove_columns=eval_dataset.column_names		,				)

      A_							       =		list(set(vocab_train['''vocab'''][0]   ) | set(vocab_test['''vocab'''][0]   )   )
      A_							       =		{v: k for k, v in enumerate(SCREAMING_SNAKE_CASE   )}
      A_							       =		vocab_dict[''' ''']
      del vocab_dict[" "]
      A_							       =		len(SCREAMING_SNAKE_CASE   )
      A_							       =		len(SCREAMING_SNAKE_CASE   )

      with open('''vocab.json'''		,				'''w'''   ) as vocab_file:
            json.dump(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE   )

      # Load pretrained model and tokenizer
      #
      # Distributed training:
      # The .from_pretrained methods guarantee that only one local process can concurrently
      # download model & vocab.
      A_							       =		WavaVecaCTCTokenizer(
          '''vocab.json'''		,				unk_token='''[UNK]'''		,				pad_token='''[PAD]'''		,				word_delimiter_token='''|'''		,				)
      A_							       =		WavaVecaFeatureExtractor(
          feature_size=1		,				sampling_rate=16000		,				padding_value=0.0		,				do_normalize=SCREAMING_SNAKE_CASE		,				return_attention_mask=SCREAMING_SNAKE_CASE   )
      A_							       =		WavaVecaProcessor(feature_extractor=SCREAMING_SNAKE_CASE		,				tokenizer=SCREAMING_SNAKE_CASE   )
      A_							       =		WavaVecaForCTC.from_pretrained(
          model_args.model_name_or_path		,				cache_dir=model_args.cache_dir		,				activation_dropout=model_args.activation_dropout		,				attention_dropout=model_args.attention_dropout		,				hidden_dropout=model_args.hidden_dropout		,				feat_proj_dropout=model_args.feat_proj_dropout		,				mask_time_prob=model_args.mask_time_prob		,				gradient_checkpointing=training_args.gradient_checkpointing		,				layerdrop=model_args.layerdrop		,				ctc_loss_reduction='''mean'''		,				pad_token_id=processor.tokenizer.pad_token_id		,				vocab_size=len(processor.tokenizer   )		,				)

      if data_args.max_train_samples is not None:
            A_							       =		min(len(SCREAMING_SNAKE_CASE   )		,				data_args.max_train_samples   )
            A_							       =		train_dataset.select(range(SCREAMING_SNAKE_CASE   )   )

      if data_args.max_val_samples is not None:
            A_							       =		eval_dataset.select(range(data_args.max_val_samples   )   )

      A_							       =		torchaudio.transforms.Resample(48000		,				16000   )

      # Preprocessing the datasets.
      # We need to read the aduio files as arrays and tokenize the targets.
      def speech_file_to_array_fn(SCREAMING_SNAKE_CASE   ):
            A_       ,A_							       =		torchaudio.load(batch['''path''']   )
            A_							       =		resampler(SCREAMING_SNAKE_CASE   ).squeeze().numpy()
            A_							       =		16000
            A_							       =		batch['''text''']
            return batch

      A_							       =		train_dataset.map(
          SCREAMING_SNAKE_CASE		,				remove_columns=train_dataset.column_names		,				num_proc=data_args.preprocessing_num_workers		,				)
      A_							       =		eval_dataset.map(
          SCREAMING_SNAKE_CASE		,				remove_columns=eval_dataset.column_names		,				num_proc=data_args.preprocessing_num_workers		,				)

      def prepare_dataset(SCREAMING_SNAKE_CASE   ):
            # check that all files have the correct sampling rate
            assert (
                len(set(batch['''sampling_rate''']   )   ) == 1
            ), f"Make sure all inputs have the same sampling rate of {processor.feature_extractor.sampling_rate}."

            A_							       =		processor(
                audio=batch['''speech''']		,				text=batch['''target_text''']		,				sampling_rate=batch['''sampling_rate'''][0]   )
            batch.update(SCREAMING_SNAKE_CASE   )
            return batch

      A_							       =		train_dataset.map(
          SCREAMING_SNAKE_CASE		,				remove_columns=train_dataset.column_names		,				batch_size=training_args.per_device_train_batch_size		,				batched=SCREAMING_SNAKE_CASE		,				num_proc=data_args.preprocessing_num_workers		,				)
      A_							       =		eval_dataset.map(
          SCREAMING_SNAKE_CASE		,				remove_columns=eval_dataset.column_names		,				batch_size=training_args.per_device_train_batch_size		,				batched=SCREAMING_SNAKE_CASE		,				num_proc=data_args.preprocessing_num_workers		,				)

      # Metric
      A_							       =		datasets.load_metric('''wer'''   )

      def compute_metrics(SCREAMING_SNAKE_CASE   ):
            A_							       =		pred.predictions
            A_							       =		np.argmax(SCREAMING_SNAKE_CASE		,				axis=-1   )

            A_							       =		processor.tokenizer.pad_token_id

            A_							       =		processor.batch_decode(SCREAMING_SNAKE_CASE   )
            # we do not want to group tokens when computing the metrics
            A_							       =		processor.batch_decode(pred.label_ids		,				group_tokens=SCREAMING_SNAKE_CASE   )

            A_							       =		wer_metric.compute(predictions=SCREAMING_SNAKE_CASE		,				references=SCREAMING_SNAKE_CASE   )

            return {"wer": wer}

      if model_args.freeze_feature_extractor:
            model.freeze_feature_extractor()

      # Data collator
      A_							       =		DataCollatorCTCWithPadding(processor=SCREAMING_SNAKE_CASE		,				padding=SCREAMING_SNAKE_CASE   )

      # Initialize our Trainer
      A_							       =		CTCTrainer(
          model=SCREAMING_SNAKE_CASE		,				data_collator=SCREAMING_SNAKE_CASE		,				args=SCREAMING_SNAKE_CASE		,				compute_metrics=SCREAMING_SNAKE_CASE		,				train_dataset=train_dataset if training_args.do_train else None		,				eval_dataset=eval_dataset if training_args.do_eval else None		,				tokenizer=processor.feature_extractor		,				)

      # Training
      if training_args.do_train:
            if last_checkpoint is not None:
                  A_							       =		last_checkpoint
            elif os.path.isdir(model_args.model_name_or_path   ):
                  A_							       =		model_args.model_name_or_path
            else:
                  A_							       =		None

            # Save the feature_extractor and the tokenizer
            if is_main_process(training_args.local_rank   ):
                  processor.save_pretrained(training_args.output_dir   )

            A_							       =		trainer.train(resume_from_checkpoint=SCREAMING_SNAKE_CASE   )
            trainer.save_model()

            A_							       =		train_result.metrics
            A_							       =		(
                data_args.max_train_samples if data_args.max_train_samples is not None else len(SCREAMING_SNAKE_CASE   )
            )
            A_							       =		min(SCREAMING_SNAKE_CASE		,				len(SCREAMING_SNAKE_CASE   )   )

            trainer.log_metrics('''train'''		,				SCREAMING_SNAKE_CASE   )
            trainer.save_metrics('''train'''		,				SCREAMING_SNAKE_CASE   )
            trainer.save_state()

      # Evaluation
      A_							       =		{}
      if training_args.do_eval:
            logger.info('''*** Evaluate ***'''   )
            A_							       =		trainer.evaluate()
            A_							       =		data_args.max_val_samples if data_args.max_val_samples is not None else len(SCREAMING_SNAKE_CASE   )
            A_							       =		min(SCREAMING_SNAKE_CASE		,				len(SCREAMING_SNAKE_CASE   )   )

            trainer.log_metrics('''eval'''		,				SCREAMING_SNAKE_CASE   )
            trainer.save_metrics('''eval'''		,				SCREAMING_SNAKE_CASE   )

      return results


if __name__ == "__main__":
      main()








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


__lowercase				      =						logging.get_logger(__name__)

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

__lowercase				      =						{
    """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"""
        ),
    },
}
__lowercase				      =						{
    """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"""
        ),
    },
}
__lowercase				      =						{
    """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"""
        ),
    },
}

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


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

class    _lowercase				(					__lowerCamelCase			):
   _lowercase     :       List[Any]  					=    VOCAB_FILES_NAMES
   _lowercase     :       Tuple  					=    CONTEXT_ENCODER_PRETRAINED_VOCAB_FILES_MAP
   _lowercase     :       int  					=    CONTEXT_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   _lowercase     :       Optional[Any]  					=    CONTEXT_ENCODER_PRETRAINED_INIT_CONFIGURATION



class    _lowercase				(					__lowerCamelCase			):
   _lowercase     :       Union[str, Any]  					=    VOCAB_FILES_NAMES
   _lowercase     :       Any  					=    QUESTION_ENCODER_PRETRAINED_VOCAB_FILES_MAP
   _lowercase     :       Optional[int]  					=    QUESTION_ENCODER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   _lowercase     :       Dict  					=    QUESTION_ENCODER_PRETRAINED_INIT_CONFIGURATION



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

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


__lowercase				      =						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(__lowerCamelCase			)
class    _lowercase				:

   def __call__(				self	:    str			,						lowerCamelCase__	:    Optional[int]			,						lowerCamelCase__	:    Optional[str] = None			,						lowerCamelCase__	:    Optional[str] = None			,						lowerCamelCase__	:    Union[bool, str] = False			,						lowerCamelCase__	:    Union[bool, str] = False			,						lowerCamelCase__	:    Optional[int] = None			,						lowerCamelCase__	:    Optional[Union[str, TensorType]] = None			,						lowerCamelCase__	:    Optional[bool] = None			,						**lowerCamelCase__	:    Optional[int]			,						) ->			BatchEncoding:
         """simple docstring"""





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

   def     UpperCamelCase					(				self	:    int			,						lowerCamelCase__	:    BatchEncoding			,						lowerCamelCase__	:    DPRReaderOutput			,						lowerCamelCase__	:    int = 1_6			,						lowerCamelCase__	:    int = 6_4			,						lowerCamelCase__	:    int = 4			,						) ->			List[DPRSpanPrediction]:
         """simple docstring"""





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

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





   def     UpperCamelCase					(				self	:    str			,						lowerCamelCase__	:    List[int]			,						lowerCamelCase__	:    List[int]			,						lowerCamelCase__	:    int			,						lowerCamelCase__	:    int			,						) ->			List[DPRSpanPrediction]:
         """simple docstring"""





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



@add_end_docstrings(__lowerCamelCase			)
class    _lowercase				(					__lowerCamelCase,__lowerCamelCase			):
   _lowercase     :       Any  					=    VOCAB_FILES_NAMES
   _lowercase     :       Any  					=    READER_PRETRAINED_VOCAB_FILES_MAP
   _lowercase     :       Union[str, Any]  					=    READER_PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   _lowercase     :       Union[str, Any]  					=    READER_PRETRAINED_INIT_CONFIGURATION
   _lowercase     :       List[Any]  					=    ['input_ids', 'attention_mask']





