

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

from ..utils.generic import ModelOutput






class 	UpperCAmelCase__		(	lowerCamelCase__	):






  """simple docstring"""


  def __init__(							self						:		int     ,       __lowerCamelCase						:		Optional[int]     ,       __lowerCamelCase						:		Tuple     ,       __lowerCamelCase						:		Dict )		->							Union[str, Any]:
       SCREAMING_SNAKE_CASE__   =   dataset
       SCREAMING_SNAKE_CASE__   =   process
       SCREAMING_SNAKE_CASE__   =   params






  def __len__(							self						:		Optional[Any] )		->							Optional[int]:
       return len(self.dataset )







  def __getitem__(							self						:		Optional[Any]     ,       __lowerCamelCase						:		Union[str, Any] )		->							str:
       SCREAMING_SNAKE_CASE__   =   self.dataset[i]
       SCREAMING_SNAKE_CASE__   =   self.process(__lowerCamelCase     ,       **self.params )
       return processed






class 	UpperCAmelCase__		(	lowerCamelCase__	):






  """simple docstring"""


  def __init__(							self						:		Optional[int]     ,       __lowerCamelCase						:		List[Any]     ,       __lowerCamelCase						:		Optional[int]     ,       __lowerCamelCase						:		Dict     ,       __lowerCamelCase						:		Tuple=None )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   loader
       SCREAMING_SNAKE_CASE__   =   infer
       SCREAMING_SNAKE_CASE__   =   params
       if loader_batch_size == 1:
            # Let's spare some time by deactivating altogether
            SCREAMING_SNAKE_CASE__   =   None
       SCREAMING_SNAKE_CASE__   =   loader_batch_size

       # Internal bookkeeping
       SCREAMING_SNAKE_CASE__   =   None
       SCREAMING_SNAKE_CASE__   =   None






  def __len__(							self						:		Tuple )		->							Dict:
       return len(self.loader )






  def __iter__(							self						:		Optional[Any] )		->							str:
       SCREAMING_SNAKE_CASE__   =   iter(self.loader )
       return self






  def 				lowercase_		(							self						:		Tuple )		->							Tuple:







       if isinstance(self._loader_batch_data     ,       torch.Tensor ):
            # Batch data is simple tensor, just fetch the slice
            SCREAMING_SNAKE_CASE__   =   self._loader_batch_data[self._loader_batch_index]
       else:
            # Batch data is assumed to be BaseModelOutput (or dict)
            SCREAMING_SNAKE_CASE__   =   {}
            for k, element in self._loader_batch_data.items():
                 if isinstance(__lowerCamelCase     ,       __lowerCamelCase ):
                      # Convert ModelOutput to tuple first
                      SCREAMING_SNAKE_CASE__   =   element.to_tuple()
                      if isinstance(element[0]     ,       torch.Tensor ):
                           SCREAMING_SNAKE_CASE__   =   tuple(el[self._loader_batch_index].unsqueeze(0 ) for el in element )
                      elif isinstance(element[0]     ,       np.ndarray ):
                           SCREAMING_SNAKE_CASE__   =   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(__lowerCamelCase     ,       __lowerCamelCase ):
                      # Those are stored as lists of tensors so need specific unbatching.
                      if isinstance(element[0]     ,       torch.Tensor ):
                           SCREAMING_SNAKE_CASE__   =   tuple(el[self._loader_batch_index].unsqueeze(0 ) for el in element )
                      elif isinstance(element[0]     ,       np.ndarray ):
                           SCREAMING_SNAKE_CASE__   =   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
                      SCREAMING_SNAKE_CASE__   =   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

                      SCREAMING_SNAKE_CASE__   =   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
                      SCREAMING_SNAKE_CASE__   =   np.expand_dims(element[self._loader_batch_index]     ,       0 )
                 else:
                      # This is typically a list, so no need to `unsqueeze`.
                      SCREAMING_SNAKE_CASE__   =   element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
            SCREAMING_SNAKE_CASE__   =   self._loader_batch_data.__class__(__lowerCamelCase )
       self._loader_batch_index += 1
       return result







  def 				lowercase_		(							self						:		str )		->							Tuple:







       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
       SCREAMING_SNAKE_CASE__   =   next(self.iterator )
       SCREAMING_SNAKE_CASE__   =   self.infer(__lowerCamelCase     ,       **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(__lowerCamelCase     ,       torch.Tensor ):
                 SCREAMING_SNAKE_CASE__   =   processed
            else:
                 SCREAMING_SNAKE_CASE__   =   list(processed.keys() )[0]
                 SCREAMING_SNAKE_CASE__   =   processed[key]
            if isinstance(__lowerCamelCase     ,       __lowerCamelCase ):
                 SCREAMING_SNAKE_CASE__   =   len(__lowerCamelCase )
            else:
                 SCREAMING_SNAKE_CASE__   =   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.
                 SCREAMING_SNAKE_CASE__   =   observed_batch_size
            # Setting internal index to unwrap the batch
            SCREAMING_SNAKE_CASE__   =   processed
            SCREAMING_SNAKE_CASE__   =   0
            return self.loader_batch_item()
       else:
            # We're not unrolling batches
            return processed






class 	UpperCAmelCase__		(	lowerCamelCase__	):






  """simple docstring"""


  def __init__(							self						:		str     ,       __lowerCamelCase						:		Dict     ,       __lowerCamelCase						:		Optional[int]     ,       __lowerCamelCase						:		Optional[int]     ,       __lowerCamelCase						:		Tuple=None )		->							str:
       super().__init__(__lowerCamelCase     ,       __lowerCamelCase     ,       __lowerCamelCase )






  def __iter__(							self						:		Any )		->							Optional[int]:
       SCREAMING_SNAKE_CASE__   =   iter(self.loader )
       SCREAMING_SNAKE_CASE__   =   None
       return self







  def 				lowercase_		(							self						:		List[Any] )		->							List[Any]:







       if self.subiterator is None:
            SCREAMING_SNAKE_CASE__   =   self.infer(next(self.iterator )     ,       **self.params )
       try:
            # Try to return next item
            SCREAMING_SNAKE_CASE__   =   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
            SCREAMING_SNAKE_CASE__   =   self.infer(next(self.iterator )     ,       **self.params )
            SCREAMING_SNAKE_CASE__   =   next(self.subiterator )
       return processed






class 	UpperCAmelCase__		(	lowerCamelCase__	):






  """simple docstring"""


  def __iter__(							self						:		int )		->							str:
       SCREAMING_SNAKE_CASE__   =   iter(self.loader )
       return self







  def 				lowercase_		(							self						:		Tuple )		->							int:
       SCREAMING_SNAKE_CASE__   =   False
       SCREAMING_SNAKE_CASE__   =   []
       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:
                 SCREAMING_SNAKE_CASE__   =   self.loader_batch_item()
                 SCREAMING_SNAKE_CASE__   =   item.pop('''is_last''' )
                 accumulator.append(__lowerCamelCase )
                 if is_last:
                      return accumulator

       while not is_last:
            SCREAMING_SNAKE_CASE__   =   self.infer(next(self.iterator )     ,       **self.params )
            if self.loader_batch_size is not None:
                 if isinstance(__lowerCamelCase     ,       torch.Tensor ):
                      SCREAMING_SNAKE_CASE__   =   processed
                 else:
                      SCREAMING_SNAKE_CASE__   =   list(processed.keys() )[0]
                      SCREAMING_SNAKE_CASE__   =   processed[key]
                 if isinstance(__lowerCamelCase     ,       __lowerCamelCase ):
                      SCREAMING_SNAKE_CASE__   =   len(__lowerCamelCase )
                 else:
                      SCREAMING_SNAKE_CASE__   =   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.
                      SCREAMING_SNAKE_CASE__   =   observed_batch_size
                 SCREAMING_SNAKE_CASE__   =   processed
                 SCREAMING_SNAKE_CASE__   =   0
                 while self._loader_batch_index < self.loader_batch_size:
                      SCREAMING_SNAKE_CASE__   =   self.loader_batch_item()
                      SCREAMING_SNAKE_CASE__   =   item.pop('''is_last''' )
                      accumulator.append(__lowerCamelCase )
                      if is_last:
                           return accumulator
            else:
                 SCREAMING_SNAKE_CASE__   =   processed
                 SCREAMING_SNAKE_CASE__   =   item.pop('''is_last''' )
                 accumulator.append(__lowerCamelCase )
       return accumulator






class 	UpperCAmelCase__		(	lowerCamelCase__	):






  """simple docstring"""


  def __init__(							self						:		str     ,       __lowerCamelCase						:		Any     ,       __lowerCamelCase						:		List[str] )		->							Optional[int]:
       SCREAMING_SNAKE_CASE__   =   dataset
       SCREAMING_SNAKE_CASE__   =   key






  def __len__(							self						:		str )		->							List[str]:
       return len(self.dataset )







  def __getitem__(							self						:		Any     ,       __lowerCamelCase						:		List[str] )		->							Any:
       return self.dataset[i][self.key]






class 	UpperCAmelCase__		(	lowerCamelCase__	):






  """simple docstring"""


  def __init__(							self						:		Dict     ,       __lowerCamelCase						:		Union[str, Any]     ,       __lowerCamelCase						:		Dict     ,       __lowerCamelCase						:		Union[str, Any] )		->							List[str]:
       SCREAMING_SNAKE_CASE__   =   dataset
       SCREAMING_SNAKE_CASE__   =   keya
       SCREAMING_SNAKE_CASE__   =   keya






  def __len__(							self						:		Optional[int] )		->							List[str]:
       return len(self.dataset )







  def __getitem__(							self						:		Optional[Any]     ,       __lowerCamelCase						:		Optional[int] )		->							List[str]:
       return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}





import datasets

from .evaluate import evaluate


_SCREAMING_SNAKE_CASE					:  Union[str, Any]             =      '''\
@inproceedings{Rajpurkar2016SQuAD10,
  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},
  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},
  booktitle={EMNLP},
  year={2016}
}
'''

_SCREAMING_SNAKE_CASE					:  Dict             =      '''
This metric wrap the official scoring script for version 1 of the Stanford Question Answering Dataset (SQuAD).

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by
crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,
from the corresponding reading passage, or the question might be unanswerable.
'''

_SCREAMING_SNAKE_CASE					:  str             =      '''
Computes SQuAD scores (F1 and EM).
Args:
    predictions: List of question-answers dictionaries with the following key-values:
        - \'id\': id of the question-answer pair as given in the references (see below)
        - \'prediction_text\': the text of the answer
    references: List of question-answers dictionaries with the following key-values:
        - \'id\': id of the question-answer pair (see above),
        - \'answers\': a Dict in the SQuAD dataset format
            {
                \'text\': list of possible texts for the answer, as a list of strings
                \'answer_start\': list of start positions for the answer, as a list of ints
            }
            Note that answer_start values are not taken into account to compute the metric.
Returns:
    \'exact_match\': Exact match (the normalized answer exactly match the gold answer)
    \'f1\': The F-score of predicted tokens versus the gold answer
Examples:

    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\'}]
    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]
    >>> squad_metric = datasets.load_metric("squad")
    >>> results = squad_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {\'exact_match\': 100.0, \'f1\': 100.0}
'''






@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION						,					_KWARGS_DESCRIPTION	)
class 	UpperCAmelCase__		(	datasets.Metric	):






  """simple docstring"""


  def 				lowercase_		(							self						:		List[Any] )		->							Optional[int]:
       return datasets.MetricInfo(
           description=_DESCRIPTION     ,       citation=_CITATION     ,       inputs_description=_KWARGS_DESCRIPTION     ,       features=datasets.Features(
               {
                   '''predictions''': {'''id''': datasets.Value('''string''' ), '''prediction_text''': datasets.Value('''string''' )},
                   '''references''': {
                       '''id''': datasets.Value('''string''' ),
                       '''answers''': datasets.features.Sequence(
                           {
                               '''text''': datasets.Value('''string''' ),
                               '''answer_start''': datasets.Value('''int32''' ),
                           } ),
                   },
               } )     ,       codebase_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']     ,       reference_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']     ,       )







  def 				lowercase_		(							self						:		Optional[Any]     ,       __lowerCamelCase						:		Dict     ,       __lowerCamelCase						:		List[Any] )		->							Optional[int]:
       SCREAMING_SNAKE_CASE__   =   {prediction['''id''']: prediction['''prediction_text'''] for prediction in predictions}
       SCREAMING_SNAKE_CASE__   =   [
           {
               '''paragraphs''': [
                   {
                       '''qas''': [
                           {
                               '''answers''': [{'''text''': answer_text} for answer_text in ref['''answers''']['''text''']],
                               '''id''': ref['''id'''],
                           }
                           for ref in references
                       ]
                   }
               ]
           }
       ]
       SCREAMING_SNAKE_CASE__   =   evaluate(dataset=__lowerCamelCase     ,       predictions=__lowerCamelCase )
       return score


