





# Lint as: python3
import sys
from collections.abc import Mapping
from typing import TYPE_CHECKING, Dict, Optional

import numpy as np
import pyarrow as pa

from .. import config
from ..utils.logging import get_logger
from ..utils.py_utils import map_nested
from .formatting import TensorFormatter


if TYPE_CHECKING:
      import jax
      import jaxlib

_lowerCAmelCase 				=     get_logger()

_lowerCAmelCase 				=     None

class    _UpperCAmelCase							(		TensorFormatter[Mapping, '''jax.Array''', Mapping]     ):
      def __init__(		self      ,       a__=None      ,       a__=None      ,       **a__       ):
         super().__init__(features=a__       )
         import jax
         from jaxlib.xla_client import Device

         if isinstance(a__      ,       a__       ):
            raise ValueError(
                F"""Expected {device} to be a `str` not {type(a__       )}, as `jaxlib.xla_extension.Device` """
                """is not serializable neither with `pickle` nor with `dill`. Instead you can surround """
                """the device with `str()` to get its string identifier that will be internally mapped """
                """to the actual `jaxlib.xla_extension.Device`."""       )
         A_     :				Union[str, Any]			      =	device if isinstance(a__      ,       a__       ) else str(jax.devices()[0]       )
         # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
         # with `pickle` nor with `dill`, so we need to use a global variable instead
         global DEVICE_MAPPING
         if DEVICE_MAPPING is None:
            A_     :				int			      =	self._map_devices_to_str()
         if self.device not in list(DEVICE_MAPPING.keys()       ):
            logger.warning(
                F"""Device with string identifier {self.device} not listed among the available """
                F"""devices: {list(DEVICE_MAPPING.keys()       )}, so falling back to the default """
                F"""device: {str(jax.devices()[0]       )}."""       )
            A_     :				Tuple			      =	str(jax.devices()[0]       )
         A_     :				str			      =	jnp_array_kwargs
      @staticmethod
      def 		_lowerCamelCase      (		):
         import jax

         return {str(a__       ): device for device in jax.devices()}
      def 		_lowerCamelCase      (		self      ,       a__       ):
         import jax
         import jax.numpy as jnp

         if isinstance(a__      ,       a__       ) and column:
            if all(
                isinstance(a__      ,       jax.Array       ) and x.shape == column[0].shape and x.dtype == column[0].dtype for x in column       ):
               return jnp.stack(a__      ,       axis=0       )
         return column
      def 		_lowerCamelCase      (		self      ,       a__       ):
         import jax
         import jax.numpy as jnp

         if isinstance(a__      ,       (str, bytes, type(a__       ))       ):
            return value
         elif isinstance(a__      ,       (np.character, np.ndarray)       ) and np.issubdtype(value.dtype      ,       np.character       ):
            return value.tolist()

         A_     :				Optional[Any]			      =	{}

         if isinstance(a__      ,       (np.number, np.ndarray)       ) and np.issubdtype(value.dtype      ,       np.integer       ):
            # the default int precision depends on the jax config
            # see https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#double-64bit-precision
            if jax.config.jax_enable_xaa:
               A_     :				Optional[Any]			      =	{"""dtype""": jnp.intaa}
            else:
               A_     :				Optional[Any]			      =	{"""dtype""": jnp.intaa}
         elif isinstance(a__      ,       (np.number, np.ndarray)       ) and np.issubdtype(value.dtype      ,       np.floating       ):
            A_     :				str			      =	{"""dtype""": jnp.floataa}
         elif config.PIL_AVAILABLE and "PIL" in sys.modules:
            import PIL.Image

            if isinstance(a__      ,       PIL.Image.Image       ):
               A_     :				Optional[int]			      =	np.asarray(a__       )

        # using global variable since `jaxlib.xla_extension.Device` is not serializable neither
        # with `pickle` nor with `dill`, so we need to use a global variable instead
         global DEVICE_MAPPING
         if DEVICE_MAPPING is None:
            A_     :				str			      =	self._map_devices_to_str()

         with jax.default_device(DEVICE_MAPPING[self.device]       ):
            # calling jnp.array on a np.ndarray does copy the data
            # see https://github.com/google/jax/issues/4486
            return jnp.array(a__      ,       **{**default_dtype, **self.jnp_array_kwargs}       )
      def 		_lowerCamelCase      (		self      ,       a__       ):
         import jax

         # support for torch, tf, jax etc.
         if config.TORCH_AVAILABLE and "torch" in sys.modules:
            import torch

            if isinstance(a__      ,       torch.Tensor       ):
               return self._tensorize(data_struct.detach().cpu().numpy()[()]       )
         if hasattr(a__      ,       """__array__"""       ) and not isinstance(a__      ,       jax.Array       ):
            A_     :				Optional[Any]			      =	data_struct.__array__()
         # support for nested types like struct of list of struct
         if isinstance(a__      ,       np.ndarray       ):
            if data_struct.dtype == object:  # jax arrays cannot be instantied from an array of objects
               return self._consolidate([self.recursive_tensorize(a__       ) for substruct in data_struct]       )
         elif isinstance(a__      ,       (list, tuple)       ):
            return self._consolidate([self.recursive_tensorize(a__       ) for substruct in data_struct]       )
         return self._tensorize(a__       )
      def 		_lowerCamelCase      (		self      ,       a__       ):
         return map_nested(self._recursive_tensorize      ,       a__      ,       map_list=a__       )
      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				List[str]			      =	self.numpy_arrow_extractor().extract_row(a__       )
         A_     :				Dict			      =	self.python_features_decoder.decode_row(a__       )
         return self.recursive_tensorize(a__       )
      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				int			      =	self.numpy_arrow_extractor().extract_column(a__       )
         A_     :				List[str]			      =	self.python_features_decoder.decode_column(a__      ,       pa_table.column_names[0]       )
         A_     :				List[Any]			      =	self.recursive_tensorize(a__       )
         A_     :				Optional[int]			      =	self._consolidate(a__       )
         return column






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				List[str]			      =	self.numpy_arrow_extractor().extract_batch(a__       )
         A_     :				Optional[int]			      =	self.python_features_decoder.decode_batch(a__       )
         A_     :				Any			      =	self.recursive_tensorize(a__       )
         for column_name in batch:
            A_     :				int			      =	self._consolidate(batch[column_name]       )
         return batch













import datasets

from .evaluate import evaluate


_lowerCAmelCase 				=     """\
@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}
}
"""

_lowerCAmelCase 				=     """
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.
"""

_lowerCAmelCase 				=     """
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     ):
      def 		_lowerCamelCase      (		self       ):
         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 		_lowerCamelCase      (		self      ,       a__      ,       a__       ):
         A_     :				Union[str, Any]			      =	{prediction["""id"""]: prediction["""prediction_text"""] for prediction in predictions}
         A_     :				Optional[Any]			      =	[
             {
                 """paragraphs""": [
                     {
                         """qas""": [
                             {
                                 """answers""": [{"""text""": answer_text} for answer_text in ref["""answers"""]["""text"""]],
                                 """id""": ref["""id"""],
                             }
                             for ref in references
                         ]
                     }
                 ]
             }
         ]
         A_     :				Union[str, Any]			      =	evaluate(dataset=a__      ,       predictions=a__       )
         return score






