





"""simple docstring"""


from typing import List

from packaging import version
from sklearn.metrics import fa_score

import datasets
from datasets.config import PY_VERSION


if PY_VERSION < version.parse("""3.8"""):
   import importlib_metadata
else:
   import importlib.metadata as importlib_metadata


# TODO(Patrick/Anton)
a__					:     Union[str, Any]							=				"""\
"""

a__					:     Optional[int]							=				"""\
    XTREME-S is a benchmark to evaluate universal cross-lingual speech representations in many languages.
    XTREME-S covers four task families: speech recognition, classification, speech-to-text translation and retrieval.
"""

a__					:     Union[str, Any]							=				"""
Compute XTREME-S evaluation metric associated to each XTREME-S dataset.
Args:
    predictions: list of predictions to score.
        Each translation should be tokenized into a list of tokens.
    references: list of lists of references for each translation.
        Each reference should be tokenized into a list of tokens.
    bleu_kwargs: optional dict of keywords to be passed when computing 'bleu'.
        Keywords include Dict can be one of 'smooth_method', 'smooth_value', 'force', 'lowercase',
        'tokenize', 'use_effective_order'.
    wer_kwargs: optional dict of keywords to be passed when computing 'wer' and 'cer'.
        Keywords include 'concatenate_texts'.
Returns: depending on the XTREME-S task, one or several of:
    \"accuracy\": Accuracy - for 'fleurs-lang_id', 'minds14'
    \"f1\": F1 score - for 'minds14'
    \"wer\": Word error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'
    \"cer\": Character error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'
    \"bleu\": BLEU score according to the `sacrebleu` metric - for 'covost2'
Examples:

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'mls')  # 'mls', 'voxpopuli', 'fleurs-asr' or 'babel'
    >>> references = [\"it is sunny here\", \"paper and pen are essentials\"]
    >>> predictions = [\"it's sunny\", \"paper pen are essential\"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'wer': 0.56, 'cer': 0.27}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'covost2')
    >>> references = [\"bonjour paris\", \"il est necessaire de faire du sport de temps en temp\"]
    >>> predictions = [\"bonjour paris\", \"il est important de faire du sport souvent\"]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'bleu': 31.65}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'fleurs-lang_id')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'accuracy': 0.6}

    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'minds14')
    >>> references = [0, 1, 0, 0, 1]
    >>> predictions = [0, 1, 1, 0, 0]
    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)
    >>> print({k: round(v, 2) for k, v in results.items()})
    {'f1': 0.58, 'accuracy': 0.6}
"""

a__					:     Union[str, Any]							=				["""fleurs-asr""", """mls""", """voxpopuli""", """babel""", """covost2""", """fleurs-lang_id""", """minds14"""]
a__					:     str							=				""""""

try:
   from jiwer import transforms as tr

   a__					:     int							=				True
except ImportError:
   a__					:     str							=				False

if _jiwer_available and version.parse(importlib_metadata.version("""jiwer""")) < version.parse("""2.3.0"""):



   class      lowercase_				(     tr.AbstractTransform							):

          def __init__(			self   ,					__SCREAMING_SNAKE_CASE = " "     ):






           '''simple docstring'''

           snake_case__    							=				sentence_delimiter

          def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE     ):






           '''simple docstring'''

           return list(__SCREAMING_SNAKE_CASE     )

          def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE     ):






           '''simple docstring'''

           snake_case__    							=				[]
           for sent_idx, sentence in enumerate(__SCREAMING_SNAKE_CASE     ):
            chars.extend(self.process_string(__SCREAMING_SNAKE_CASE     )     )
            if self.sentence_delimiter is not None and self.sentence_delimiter != "" and sent_idx < len(__SCREAMING_SNAKE_CASE     ) - 1:
             chars.append(self.sentence_delimiter     )
           return chars
   a__					:     Any							=				tr.Compose(
       [tr.RemoveMultipleSpaces(), tr.Strip(), SentencesToListOfCharacters(SENTENCE_DELIMITER)]
   )
elif _jiwer_available:
   a__					:     Optional[Any]							=				tr.Compose(
       [
           tr.RemoveMultipleSpaces(),
           tr.Strip(),
           tr.ReduceToSingleSentence(SENTENCE_DELIMITER),
           tr.ReduceToListOfListOfChars(),
       ]
   )
else:
   a__					:     Any							=				None





def 							a__			(  _UpperCAmelCase	,     _UpperCAmelCase	)    ->  str:

 '''simple docstring'''






 return float((preds == labels).mean()	)





def 							a__			(  _UpperCAmelCase	,     _UpperCAmelCase	)    ->  Union[str, Any]:

 '''simple docstring'''






 return {
     "f1": float(fa_score(y_true=_UpperCAmelCase	,     y_pred=_UpperCAmelCase	,     average="macro"	)	),
     "accuracy": simple_accuracy(_UpperCAmelCase	,     _UpperCAmelCase	),
 }





def 							a__			(  _UpperCAmelCase	,     _UpperCAmelCase	,     _UpperCAmelCase="exp"	,     _UpperCAmelCase=None	,     _UpperCAmelCase=False	,     _UpperCAmelCase=False	,     _UpperCAmelCase=None	,     _UpperCAmelCase=False	,     )    ->  Dict:

 '''simple docstring'''






 snake_case__    							=				[[label] for label in labels]
 snake_case__    							=				list(_UpperCAmelCase	)
 try:
  import sacrebleu as scb
 except ImportError:
  raise ValueError(
      "sacrebleu has to be installed in order to apply the bleu metric for covost2."
      "You can install it via `pip install sacrebleu`."	)

 if version.parse(scb.__version__	) < version.parse("1.4.12"	):
  raise ImportWarning(
      "To use `sacrebleu`, the module `sacrebleu>=1.4.12` is required, and the current version of `sacrebleu` doesn't match this condition.\n"
      "You can install it with `pip install \"sacrebleu>=1.4.12\"`."	)

 snake_case__    							=				len(labels[0]	)
 if any(len(_UpperCAmelCase	) != references_per_prediction for refs in labels	):
  raise ValueError("Sacrebleu requires the same number of references for each prediction"	)
 snake_case__    							=				[[refs[i] for refs in labels] for i in range(_UpperCAmelCase	)]
 snake_case__    							=				scb.corpus_bleu(
     _UpperCAmelCase	,     _UpperCAmelCase	,     smooth_method=_UpperCAmelCase	,     smooth_value=_UpperCAmelCase	,     force=_UpperCAmelCase	,     lowercase=_UpperCAmelCase	,     use_effective_order=_UpperCAmelCase	,     **({"tokenize": tokenize} if tokenize else {})	,     )
 return {"bleu": output.score}





def 							a__			(  _UpperCAmelCase	,     _UpperCAmelCase	,     _UpperCAmelCase	,     _UpperCAmelCase	)    ->  List[str]:

 '''simple docstring'''






 try:
  from jiwer import compute_measures
 except ImportError:
  raise ValueError(
      f"jiwer has to be installed in order to apply the wer metric for {config_name}."
      "You can install it via `pip install jiwer`."	)

 if concatenate_texts:
  snake_case__    							=				compute_measures(_UpperCAmelCase	,     _UpperCAmelCase	)["wer"]

  snake_case__    							=				compute_measures(_UpperCAmelCase	,     _UpperCAmelCase	,     truth_transform=_UpperCAmelCase	,     hypothesis_transform=_UpperCAmelCase	)["wer"]
  return {"wer": wer, "cer": cer}
 else:

  def compute_score(_UpperCAmelCase	,     _UpperCAmelCase	,     _UpperCAmelCase="wer"	):
   snake_case__    							=				0
   snake_case__    							=				0
   for prediction, reference in zip(_UpperCAmelCase	,     _UpperCAmelCase	):
    if score_type == "wer":
     snake_case__    							=				compute_measures(_UpperCAmelCase	,     _UpperCAmelCase	)
    elif score_type == "cer":
     snake_case__    							=				compute_measures(
         _UpperCAmelCase	,     _UpperCAmelCase	,     truth_transform=_UpperCAmelCase	,     hypothesis_transform=_UpperCAmelCase	)
    incorrect += measures["substitutions"] + measures["deletions"] + measures["insertions"]
    total += measures["substitutions"] + measures["deletions"] + measures["hits"]
   return incorrect / total

  return {"wer": compute_score(_UpperCAmelCase	,     _UpperCAmelCase	,     "wer"	), "cer": compute_score(_UpperCAmelCase	,     _UpperCAmelCase	,     "cer"	)}



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

       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        if self.config_name not in _CONFIG_NAMES:
         raise KeyError(F"You should supply a configuration name selected in {_CONFIG_NAMES}"     )

        snake_case__    							=				"int64" if self.config_name in ["fleurs-lang_id", "minds14"] else "string"

        return datasets.MetricInfo(
            description=_DESCRIPTION   ,					citation=_CITATION   ,					inputs_description=_KWARGS_DESCRIPTION   ,					features=datasets.Features(
                {"predictions": datasets.Value(__SCREAMING_SNAKE_CASE     ), "references": datasets.Value(__SCREAMING_SNAKE_CASE     )}     )   ,					codebase_urls=[]   ,					reference_urls=[]   ,					format="numpy"   ,					)

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE=None   ,					__SCREAMING_SNAKE_CASE=None     ):






        '''simple docstring'''

        snake_case__    							=				bleu_kwargs if bleu_kwargs is not None else {}
        snake_case__    							=				wer_kwargs if wer_kwargs is not None else {}

        if self.config_name == "fleurs-lang_id":
         return {"accuracy": simple_accuracy(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     )}
        elif self.config_name == "minds14":
         return fa_and_simple_accuracy(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     )
        elif self.config_name == "covost2":
         snake_case__    							=				bleu_kwargs.pop("smooth_method"   ,					"exp"     )
         snake_case__    							=				bleu_kwargs.pop("smooth_value"   ,					__SCREAMING_SNAKE_CASE     )
         snake_case__    							=				bleu_kwargs.pop("force"   ,					__SCREAMING_SNAKE_CASE     )
         snake_case__    							=				bleu_kwargs.pop("lowercase"   ,					__SCREAMING_SNAKE_CASE     )
         snake_case__    							=				bleu_kwargs.pop("tokenize"   ,					__SCREAMING_SNAKE_CASE     )
         snake_case__    							=				bleu_kwargs.pop("use_effective_order"   ,					__SCREAMING_SNAKE_CASE     )
         return bleu(
             preds=__SCREAMING_SNAKE_CASE   ,					labels=__SCREAMING_SNAKE_CASE   ,					smooth_method=__SCREAMING_SNAKE_CASE   ,					smooth_value=__SCREAMING_SNAKE_CASE   ,					force=__SCREAMING_SNAKE_CASE   ,					lowercase=__SCREAMING_SNAKE_CASE   ,					tokenize=__SCREAMING_SNAKE_CASE   ,					use_effective_order=__SCREAMING_SNAKE_CASE   ,					)
        elif self.config_name in ["fleurs-asr", "mls", "voxpopuli", "babel"]:
         snake_case__    							=				wer_kwargs.pop("concatenate_texts"   ,					__SCREAMING_SNAKE_CASE     )
         return wer_and_cer(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					self.config_name     )
        else:
         raise KeyError(F"You should supply a configuration name selected in {_CONFIG_NAMES}"     )








"""simple docstring"""


from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax.numpy as jnp
from scipy import integrate

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    broadcast_to_shape_from_left,
)



@flax.struct.dataclass
class      lowercase_				:
       snake_case__				:							CommonSchedulerState

       # setable values
       snake_case__				:							jnp.ndarray
       snake_case__				:							jnp.ndarray
       snake_case__				:							jnp.ndarray
       snake_case__				:							Optional[int]											=     None

       # running values
       snake_case__				:							Optional[jnp.ndarray]											=     None

       @classmethod
       def 		_SCREAMING_SNAKE_CASE(			cls   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        return cls(common=__SCREAMING_SNAKE_CASE   ,					init_noise_sigma=__SCREAMING_SNAKE_CASE   ,					timesteps=__SCREAMING_SNAKE_CASE   ,					sigmas=__SCREAMING_SNAKE_CASE     )



@dataclass
class      lowercase_				(     __lowerCAmelCase							):
       snake_case__				:							LMSDiscreteSchedulerState



class      lowercase_				(     __lowerCAmelCase				,__lowerCAmelCase							):
       snake_case__				:							Any											=     [e.name for e in FlaxKarrasDiffusionSchedulers]

       snake_case__				:							jnp.dtype

       @property
       def 		_SCREAMING_SNAKE_CASE(			self     ):






        '''simple docstring'''

        return True

       @register_to_config
       def __init__(			self   ,					__SCREAMING_SNAKE_CASE = 1000   ,					__SCREAMING_SNAKE_CASE = 0.0001   ,					__SCREAMING_SNAKE_CASE = 0.02   ,					__SCREAMING_SNAKE_CASE = "linear"   ,					__SCREAMING_SNAKE_CASE = None   ,					__SCREAMING_SNAKE_CASE = "epsilon"   ,					__SCREAMING_SNAKE_CASE = jnp.floataa   ,					):






        '''simple docstring'''

        snake_case__    							=				dtype

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE = None     ):






        '''simple docstring'''

        if common is None:
         snake_case__    							=				CommonSchedulerState.create(self     )

        snake_case__    							=				jnp.arange(0   ,					self.config.num_train_timesteps     ).round()[::-1]
        snake_case__    							=				((1 - common.alphas_cumprod) / common.alphas_cumprod) ** 0.5

        # standard deviation of the initial noise distribution
        snake_case__    							=				sigmas.max()

        return LMSDiscreteSchedulerState.create(
            common=__SCREAMING_SNAKE_CASE   ,					init_noise_sigma=__SCREAMING_SNAKE_CASE   ,					timesteps=__SCREAMING_SNAKE_CASE   ,					sigmas=__SCREAMING_SNAKE_CASE   ,					)

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        ((snake_case__)							,   )    							=				jnp.where(state.timesteps == timestep   ,					size=1     )
        snake_case__    							=				step_index[0]

        snake_case__    							=				state.sigmas[step_index]
        snake_case__    							=				sample / ((sigma**2 + 1) ** 0.5)
        return sample

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ):






        '''simple docstring'''

        def lms_derivative(__SCREAMING_SNAKE_CASE     ):
         snake_case__    							=				1.0
         for k in range(__SCREAMING_SNAKE_CASE     ):
          if current_order == k:
           continue
          prod *= (tau - state.sigmas[t - k]) / (state.sigmas[t - current_order] - state.sigmas[t - k])
         return prod

        snake_case__    							=				integrate.quad(__SCREAMING_SNAKE_CASE   ,					state.sigmas[t]   ,					state.sigmas[t + 1]   ,					epsrel=1e-4     )[0]

        return integrated_coeff

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE = ()     ):






        '''simple docstring'''

        snake_case__    							=				jnp.linspace(self.config.num_train_timesteps - 1   ,					0   ,					__SCREAMING_SNAKE_CASE   ,					dtype=self.dtype     )

        snake_case__    							=				jnp.floor(__SCREAMING_SNAKE_CASE     ).astype(jnp.intaa     )
        snake_case__    							=				jnp.ceil(__SCREAMING_SNAKE_CASE     ).astype(jnp.intaa     )

        snake_case__    							=				jnp.mod(__SCREAMING_SNAKE_CASE   ,					1.0     )

        snake_case__    							=				((1 - state.common.alphas_cumprod) / state.common.alphas_cumprod) ** 0.5
        snake_case__    							=				(1 - frac) * sigmas[low_idx] + frac * sigmas[high_idx]
        snake_case__    							=				jnp.concatenate([sigmas, jnp.array([0.0]   ,					dtype=self.dtype     )]     )

        snake_case__    							=				timesteps.astype(jnp.intaa     )

        # initial running values
        snake_case__    							=				jnp.zeros((0,) + shape   ,					dtype=self.dtype     )

        return state.replace(
            timesteps=__SCREAMING_SNAKE_CASE   ,					sigmas=__SCREAMING_SNAKE_CASE   ,					num_inference_steps=__SCREAMING_SNAKE_CASE   ,					derivatives=__SCREAMING_SNAKE_CASE   ,					)

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE = 4   ,					__SCREAMING_SNAKE_CASE = True   ,					):






        '''simple docstring'''

        if state.num_inference_steps is None:
         raise ValueError(
             "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"     )

        snake_case__    							=				state.sigmas[timestep]

        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
        if self.config.prediction_type == "epsilon":
         snake_case__    							=				sample - sigma * model_output
        elif self.config.prediction_type == "v_prediction":
         # * c_out + input * c_skip
         snake_case__    							=				model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))
        else:
         raise ValueError(
             F"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`"     )

        # 2. Convert to an ODE derivative
        snake_case__    							=				(sample - pred_original_sample) / sigma
        snake_case__    							=				state.replace(derivatives=jnp.append(state.derivatives   ,					__SCREAMING_SNAKE_CASE     )     )
        if len(state.derivatives     ) > order:
         snake_case__    							=				state.replace(derivatives=jnp.delete(state.derivatives   ,					0     )     )

        # 3. Compute linear multistep coefficients
        snake_case__    							=				min(timestep + 1   ,					__SCREAMING_SNAKE_CASE     )
        snake_case__    							=				[self.get_lms_coefficient(__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE     ) for curr_order in range(__SCREAMING_SNAKE_CASE     )]

        # 4. Compute previous sample based on the derivatives path
        snake_case__    							=				sample + sum(
            coeff * derivative for coeff, derivative in zip(__SCREAMING_SNAKE_CASE   ,					reversed(state.derivatives     )     )     )

        if not return_dict:
         return (prev_sample, state)

        return FlaxLMSSchedulerOutput(prev_sample=__SCREAMING_SNAKE_CASE   ,					state=__SCREAMING_SNAKE_CASE     )

       def 		_SCREAMING_SNAKE_CASE(			self   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					__SCREAMING_SNAKE_CASE   ,					):






        '''simple docstring'''

        snake_case__    							=				state.sigmas[timesteps].flatten()
        snake_case__    							=				broadcast_to_shape_from_left(__SCREAMING_SNAKE_CASE   ,					noise.shape     )

        snake_case__    							=				original_samples + noise * sigma

        return noisy_samples

       def __len__(			self     ):






        '''simple docstring'''

        return self.config.num_train_timesteps

