




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)
_lowerCamelCase :       Optional[Any]     					=    "\\n"

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

_lowerCamelCase :       Optional[Any]     					=    "\nCompute XTREME-S evaluation metric associated to each XTREME-S dataset.\nArgs:\n    predictions: list of predictions to score.\n        Each translation should be tokenized into a list of tokens.\n    references: list of lists of references for each translation.\n        Each reference should be tokenized into a list of tokens.\n    bleu_kwargs: optional dict of keywords to be passed when computing 'bleu'.\n        Keywords include Dict can be one of 'smooth_method', 'smooth_value', 'force', 'lowercase',\n        'tokenize', 'use_effective_order'.\n    wer_kwargs: optional dict of keywords to be passed when computing 'wer' and 'cer'.\n        Keywords include 'concatenate_texts'.\nReturns: depending on the XTREME-S task, one or several of:\n    \"accuracy\": Accuracy - for 'fleurs-lang_id', 'minds14'\n    \"f1\": F1 score - for 'minds14'\n    \"wer\": Word error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'\n    \"cer\": Character error rate - for 'mls', 'fleurs-asr', 'voxpopuli', 'babel'\n    \"bleu\": BLEU score according to the `sacrebleu` metric - for 'covost2'\nExamples:\n\n    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'mls')  # 'mls', 'voxpopuli', 'fleurs-asr' or 'babel'\n    >>> references = [\"it is sunny here\", \"paper and pen are essentials\"]\n    >>> predictions = [\"it's sunny\", \"paper pen are essential\"]\n    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)\n    >>> print({k: round(v, 2) for k, v in results.items()})\n    {'wer': 0.56, 'cer': 0.27}\n\n    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'covost2')\n    >>> references = [\"bonjour paris\", \"il est necessaire de faire du sport de temps en temp\"]\n    >>> predictions = [\"bonjour paris\", \"il est important de faire du sport souvent\"]\n    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)\n    >>> print({k: round(v, 2) for k, v in results.items()})\n    {'bleu': 31.65}\n\n    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'fleurs-lang_id')\n    >>> references = [0, 1, 0, 0, 1]\n    >>> predictions = [0, 1, 1, 0, 0]\n    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)\n    >>> print({k: round(v, 2) for k, v in results.items()})\n    {'accuracy': 0.6}\n\n    >>> xtreme_s_metric = datasets.load_metric('xtreme_s', 'minds14')\n    >>> references = [0, 1, 0, 0, 1]\n    >>> predictions = [0, 1, 1, 0, 0]\n    >>> results = xtreme_s_metric.compute(predictions=predictions, references=references)\n    >>> print({k: round(v, 2) for k, v in results.items()})\n    {'f1': 0.58, 'accuracy': 0.6}\n"

_lowerCamelCase :       Tuple     					=    ["fleurs-asr", "mls", "voxpopuli", "babel", "covost2", "fleurs-lang_id", "minds14"]
_lowerCamelCase :       int     					=    ""

try:
			from jiwer import transforms as tr

			_lowerCamelCase :       List[str]     					=    True
except ImportError:
			_lowerCamelCase :       Dict     					=    False

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




			class      SCREAMING_SNAKE_CASE__  (			tr.AbstractTransform      ):





				'''simple docstring'''



				def __init__(						self			:					Optional[int]    ,						lowercase			:					Tuple = " "			):


							'''simple docstring'''


							_snake_case						=    sentence_delimiter



				def 				A							(						self			:					List[str]    ,						lowercase			:					int			):


							'''simple docstring'''


							return list(lowercase			)





				def 				A							(						self			:					Tuple    ,						lowercase			:					Tuple			):


							'''simple docstring'''


							_snake_case						=    []
							for sent_idx, sentence in enumerate(lowercase			):
										chars.extend(self.process_string(lowercase			)			)
										if self.sentence_delimiter is not None and self.sentence_delimiter != "" and sent_idx < len(lowercase			) - 1:
													chars.append(self.sentence_delimiter			)
							return chars




			_lowerCamelCase :       str     					=    tr.Compose(
			    [tr.RemoveMultipleSpaces(), tr.Strip(), SentencesToListOfCharacters(SENTENCE_DELIMITER)]
			)
elif _jiwer_available:
			_lowerCamelCase :       Optional[Any]     					=    tr.Compose(
			    [
			        tr.RemoveMultipleSpaces(),
			        tr.Strip(),
			        tr.ReduceToSingleSentence(SENTENCE_DELIMITER),
			        tr.ReduceToListOfListOfChars(),
			    ]
			)
else:
			_lowerCamelCase :       Tuple     					=    None






def  a_	( __lowercase     :			Dict ,      __lowercase     :			int				)				->						List[Any]:
			return float((preds == labels).mean()				)






def  a_	( __lowercase     :			str ,      __lowercase     :			Tuple				)				->						int:
			return {
			    "f1": float(fa_score(y_true=__lowercase ,      y_pred=__lowercase ,      average='macro'				)				),
			    "accuracy": simple_accuracy(__lowercase ,      __lowercase				),
			}






def  a_	( __lowercase     :			int ,      __lowercase     :			Any ,      __lowercase     :			Optional[int]="exp" ,      __lowercase     :			Tuple=None ,      __lowercase     :			Union[str, Any]=False ,      __lowercase     :			Optional[int]=False ,      __lowercase     :			Optional[Any]=None ,      __lowercase     :			Union[str, Any]=False ,      )				->						List[Any]:
			_snake_case						=    [[label] for label in labels]
			_snake_case						=    list(__lowercase				)
			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(__lowercase				) != 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(__lowercase				)]
			_snake_case						=    scb.corpus_bleu(
			    __lowercase ,      __lowercase ,      smooth_method=__lowercase ,      smooth_value=__lowercase ,      force=__lowercase ,      lowercase=__lowercase ,      use_effective_order=__lowercase ,      **({'tokenize': tokenize} if tokenize else {}) ,      )
			return {"bleu": output.score}






def  a_	( __lowercase     :			Dict ,      __lowercase     :			Any ,      __lowercase     :			Tuple ,      __lowercase     :			str				)				->						List[str]:







			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(__lowercase ,      __lowercase				)['wer']

						_snake_case						=    compute_measures(__lowercase ,      __lowercase ,      truth_transform=__lowercase ,      hypothesis_transform=__lowercase				)['wer']
						return {"wer": wer, "cer": cer}
			else:

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

						return {"wer": compute_score(__lowercase ,      __lowercase ,      'wer'				), "cer": compute_score(__lowercase ,      __lowercase ,      'cer'				)}




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





	'''simple docstring'''



	def 				A							(						self			:					int			):


				'''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(lowercase			), 'references': datasets.Value(lowercase			)}			)    ,						codebase_urls=[]    ,						reference_urls=[]    ,						format='numpy'    ,						)





	def 				A							(						self			:					Optional[int]    ,						lowercase			:					str    ,						lowercase			:					List[Any]    ,						lowercase			:					Optional[int]=None    ,						lowercase			:					str=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(lowercase    ,						lowercase			)}
				elif self.config_name == "minds14":
							return fa_and_simple_accuracy(lowercase    ,						lowercase			)
				elif self.config_name == "covost2":
							_snake_case						=    bleu_kwargs.pop('smooth_method'    ,						'exp'			)
							_snake_case						=    bleu_kwargs.pop('smooth_value'    ,						lowercase			)
							_snake_case						=    bleu_kwargs.pop('force'    ,						lowercase			)
							_snake_case						=    bleu_kwargs.pop('lowercase'    ,						lowercase			)
							_snake_case						=    bleu_kwargs.pop('tokenize'    ,						lowercase			)
							_snake_case						=    bleu_kwargs.pop('use_effective_order'    ,						lowercase			)
							return bleu(
							    preds=lowercase    ,						labels=lowercase    ,						smooth_method=lowercase    ,						smooth_value=lowercase    ,						force=lowercase    ,						lowercase=lowercase    ,						tokenize=lowercase    ,						use_effective_order=lowercase    ,						)
				elif self.config_name in ["fleurs-asr", "mls", "voxpopuli", "babel"]:
							_snake_case						=    wer_kwargs.pop('concatenate_texts'    ,						lowercase			)
							return wer_and_cer(lowercase    ,						lowercase    ,						lowercase    ,						self.config_name			)
				else:
							raise KeyError(f'''You should supply a configuration name selected in {_CONFIG_NAMES}'''			)

import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


_lowerCamelCase :       Optional[Any]     					=    [nn.LayerNorm]

_lowerCamelCase :       Optional[int]     					=    logging.get_logger(__name__)

_lowerCamelCase :       Optional[int]     					=    version.parse(version.parse(torch.__version__).base_version)

_lowerCamelCase :       Dict     					=    parsed_torch_version_base >= version.parse('''2.0''')
_lowerCamelCase :       str     					=    parsed_torch_version_base >= version.parse('''1.12''')
_lowerCamelCase :       str     					=    parsed_torch_version_base >= version.parse('''1.11''')
_lowerCamelCase :       Dict     					=    parsed_torch_version_base < version.parse('''1.11''')






def  a_	( __lowercase     :			Optional[int] ,      __lowercase     :			Optional[int] ,      __lowercase     :			Optional[Any] ,      __lowercase     :			Union[str, Any] ,      self     :			Optional[int]				)				->						Dict:
			from torch import _softmax_backward_data

			if is_torch_less_than_1_11:
						return _softmax_backward_data(__lowercase ,      __lowercase ,      parent.dim ,      self				)
			else:
						return _softmax_backward_data(__lowercase ,      __lowercase ,      parent.dim ,      self.dtype				)






def  a_	( __lowercase     :			nn.Linear ,      __lowercase     :			torch.LongTensor ,      __lowercase     :			int = 0				)				->						nn.Linear:
			_snake_case						=    index.to(layer.weight.device				)
			_snake_case						=    layer.weight.index_select(__lowercase ,      __lowercase				).clone().detach()
			if layer.bias is not None:
						if dim == 1:
									_snake_case						=    layer.bias.clone().detach()
						else:
									_snake_case						=    layer.bias[index].clone().detach()
			_snake_case						=    list(layer.weight.size()				)
			_snake_case						=    len(__lowercase				)
			_snake_case						=    nn.Linear(new_size[1] ,      new_size[0] ,      bias=layer.bias is not None				).to(layer.weight.device				)
			_snake_case						=    False
			new_layer.weight.copy_(W.contiguous()				)
			_snake_case						=    True
			if layer.bias is not None:
						_snake_case						=    False
						new_layer.bias.copy_(b.contiguous()				)
						_snake_case						=    True
			return new_layer




class      SCREAMING_SNAKE_CASE__  (			nn.Module      ):





	'''simple docstring'''



	def __init__(						self			:					Optional[Any]    ,						lowercase			:					List[str]    ,						lowercase			:					Tuple			):


				'''simple docstring'''


				super().__init__()
				_snake_case						=    nf
				_snake_case						=    nn.Parameter(torch.empty(lowercase    ,						lowercase			)			)
				_snake_case						=    nn.Parameter(torch.zeros(lowercase			)			)
				nn.init.normal_(self.weight    ,						std=0.02			)





	def 				A							(						self			:					List[Any]    ,						lowercase			:					int			):


				'''simple docstring'''


				_snake_case						=    x.size()[:-1] + (self.nf,)
				_snake_case						=    torch.addmm(self.bias    ,						x.view(-1    ,						x.size(-1			)			)    ,						self.weight			)
				_snake_case						=    x.view(lowercase			)
				return x






def  a_	( __lowercase     :			ConvaD ,      __lowercase     :			torch.LongTensor ,      __lowercase     :			int = 1				)				->						ConvaD:
			_snake_case						=    index.to(layer.weight.device				)
			_snake_case						=    layer.weight.index_select(__lowercase ,      __lowercase				).clone().detach()
			if dim == 0:
						_snake_case						=    layer.bias.clone().detach()
			else:
						_snake_case						=    layer.bias[index].clone().detach()
			_snake_case						=    list(layer.weight.size()				)
			_snake_case						=    len(__lowercase				)
			_snake_case						=    ConvaD(new_size[1] ,      new_size[0]				).to(layer.weight.device				)
			_snake_case						=    False
			new_layer.weight.copy_(W.contiguous()				)
			_snake_case						=    True
			_snake_case						=    False
			new_layer.bias.copy_(b.contiguous()				)
			_snake_case						=    True
			return new_layer






def  a_	( __lowercase     :			Union[nn.Linear, ConvaD] ,      __lowercase     :			torch.LongTensor ,      __lowercase     :			Optional[int] = None				)				->						Union[nn.Linear, ConvaD]:
			if isinstance(__lowercase ,      nn.Linear				):
						return prune_linear_layer(__lowercase ,      __lowercase ,      dim=0 if dim is None else dim				)
			elif isinstance(__lowercase ,      __lowercase				):
						return prune_convad_layer(__lowercase ,      __lowercase ,      dim=1 if dim is None else dim				)
			else:
						raise ValueError(f'''Can\'t prune layer of class {layer.__class__}'''				)






def  a_	( __lowercase     :			Callable[..., torch.Tensor] ,      __lowercase     :			int ,      __lowercase     :			int ,      *__lowercase     :			List[str]				)				->						torch.Tensor:
			assert len(__lowercase				) > 0, f'''{input_tensors} has to be a tuple/list of tensors'''

			# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
			_snake_case						=    len(inspect.signature(__lowercase				).parameters				)
			if num_args_in_forward_chunk_fn != len(__lowercase				):
						raise ValueError(
						    f'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(__lowercase				)} input '''
						    'tensors are given'				)

			if chunk_size > 0:
						_snake_case						=    input_tensors[0].shape[chunk_dim]
						for input_tensor in input_tensors:
									if input_tensor.shape[chunk_dim] != tensor_shape:
												raise ValueError(
												    f'''All input tenors have to be of the same shape: {tensor_shape}, '''
												    f'''found shape {input_tensor.shape[chunk_dim]}'''				)

						if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
									raise ValueError(
									    f'''The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk '''
									    f'''size {chunk_size}'''				)

						_snake_case						=    input_tensors[0].shape[chunk_dim] // chunk_size

						# chunk input tensor into tuples
						_snake_case						=    tuple(input_tensor.chunk(__lowercase ,      dim=__lowercase				) for input_tensor in input_tensors				)
						# apply forward fn to every tuple
						_snake_case						=    tuple(forward_fn(*__lowercase				) for input_tensors_chunk in zip(*__lowercase				)				)
						# concatenate output at same dimension
						return torch.cat(__lowercase ,      dim=__lowercase				)

			return forward_fn(*__lowercase				)






def  a_	( __lowercase     :			List[int] ,      __lowercase     :			int ,      __lowercase     :			int ,      __lowercase     :			Set[int]				)				->						Tuple[Set[int], torch.LongTensor]:
			_snake_case						=    torch.ones(__lowercase ,      __lowercase				)
			_snake_case						=    set(__lowercase				) - already_pruned_heads  # Convert to set and remove already pruned heads
			for head in heads:
						# Compute how many pruned heads are before the head and move the index accordingly
						_snake_case						=    head - sum(1 if h < head else 0 for h in already_pruned_heads				)
						_snake_case						=    0
			_snake_case						=    mask.view(-1				).contiguous().eq(1				)
			_snake_case						=    torch.arange(len(__lowercase				)				)[mask].long()
			return heads, index






def  a_	( *__lowercase     :			Union[torch.Tensor, List[torch.Tensor]] ,      __lowercase     :			Optional[str] = None				)				->						Tuple[torch.Tensor, ...]:
			return torch.meshgrid(*__lowercase ,      indexing=__lowercase				)






def  a_	( __lowercase     :			torch.Tensor				)				->						Tuple[torch.device, int, int]:
			return tensor.device, storage_ptr(__lowercase				), storage_size(__lowercase				)