



'''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)
lowercase__     :  Any    					=			'\\n'

lowercase__     :  Dict    					=			'\\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'

lowercase__     :  List[str]    					=			'\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'

lowercase__     :  List[Any]    					=			['fleurs-asr', 'mls', 'voxpopuli', 'babel', 'covost2', 'fleurs-lang_id', 'minds14']
lowercase__     :  Any    					=			''

try:
      from jiwer import transforms as tr

      lowercase__     :  Optional[int]    					=			True
except ImportError:
      lowercase__     :  List[Any]    					=			False

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

      class       lowerCamelCase_							(	tr.AbstractTransform  ):






         '''simple docstring'''

         def __init__(  self     ,						__UpperCamelCase = " " ):
             UpperCamelCase__      =			sentence_delimiter


         def 	_UpperCamelCase       (  self     ,						__UpperCamelCase ):
             return list(__UpperCamelCase )







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

      lowercase__     :  int    					=			tr.Compose(
          [tr.RemoveMultipleSpaces(), tr.Strip(), SentencesToListOfCharacters(SENTENCE_DELIMITER)]
      )
elif _jiwer_available:
      lowercase__     :  str    					=			tr.Compose(
          [
              tr.RemoveMultipleSpaces(),
              tr.Strip(),
              tr.ReduceToSingleSentence(SENTENCE_DELIMITER),
              tr.ReduceToListOfListOfChars(),
          ]
      )
else:
      lowercase__     :  Dict    					=			None





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_	):
    return float((preds == labels).mean()	)





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





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_="exp"					,  lowerCAmelCase_=None					,  lowerCAmelCase_=False					,  lowerCAmelCase_=False					,  lowerCAmelCase_=None					,  lowerCAmelCase_=False					,  ):
    # xtreme-s can only have one label
    UpperCamelCase__      =			[[label] for label in labels]
    UpperCamelCase__      =			list(lowerCAmelCase_	)
    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\"`."	)

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





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_	):
    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:
        UpperCamelCase__      =			compute_measures(lowerCAmelCase_					,  lowerCAmelCase_	)["wer"]

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

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

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





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






   '''simple docstring'''

   def 	_UpperCamelCase       (  self ):
       if self.config_name not in _CONFIG_NAMES:
           raise KeyError(f'You should supply a configuration name selected in {_CONFIG_NAMES}' )

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







   def 	_UpperCamelCase       (  self     ,						__UpperCamelCase     ,						__UpperCamelCase     ,						__UpperCamelCase=None     ,						__UpperCamelCase=None ):
       UpperCamelCase__      =			bleu_kwargs if bleu_kwargs is not None else {}
       UpperCamelCase__      =			wer_kwargs if wer_kwargs is not None else {}

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









'''simple docstring'''







import argparse

from tax import checkpoints

from transformers import FlaxTaForConditionalGeneration, TaConfig





def 						snake_case__ (lowerCAmelCase_					,  lowerCAmelCase_					,  lowerCAmelCase_	):
    UpperCamelCase__      =			TaConfig.from_pretrained(lowerCAmelCase_	)
    UpperCamelCase__      =			FlaxTaForConditionalGeneration(config=lowerCAmelCase_	)
    UpperCamelCase__      =			checkpoints.load_tax_checkpoint(lowerCAmelCase_	)

    UpperCamelCase__      =			"wi_0" in tax_model["target"]["encoder"]["layers_0"]["mlp"]

    # Encoder
    for layer_index in range(config.num_layers	):
        UpperCamelCase__      =			F'layers_{str(lowerCAmelCase_	)}'

        # Self-Attention
        UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["attention"]["key"]["kernel"]
        UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["attention"]["out"]["kernel"]
        UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["attention"]["query"]["kernel"]
        UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["attention"]["value"]["kernel"]

        # Layer Normalization
        UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["pre_attention_layer_norm"]["scale"]

        if split_mlp_wi:
            UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["mlp"]["wi_0"]["kernel"]
            UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["mlp"]["wi_1"]["kernel"]
        else:
            UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["mlp"]["wi"]["kernel"]

        UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["mlp"]["wo"]["kernel"]

        # Layer Normalization
        UpperCamelCase__      =			tax_model["target"]["encoder"][layer_name]["pre_mlp_layer_norm"]["scale"]

        # Assigning
        UpperCamelCase__      =			tax_attention_key
        UpperCamelCase__      =			tax_attention_out
        UpperCamelCase__      =			tax_attention_query
        UpperCamelCase__      =			tax_attention_value

        UpperCamelCase__      =			tax_attention_layer_norm

        if split_mlp_wi:
            UpperCamelCase__      =			tax_mlp_wi_a
            UpperCamelCase__      =			tax_mlp_wi_a
        else:
            UpperCamelCase__      =			tax_mlp_wi

        UpperCamelCase__      =			tax_mlp_wo
        UpperCamelCase__      =			tax_mlp_layer_norm

    # Only for layer 0:
    UpperCamelCase__      =			tax_model["target"]["encoder"]["relpos_bias"]["rel_embedding"].T
    UpperCamelCase__      =			tax_encoder_rel_embedding

    # Assigning
    UpperCamelCase__      =			tax_model["target"]["encoder"]["encoder_norm"]["scale"]
    UpperCamelCase__      =			tax_encoder_norm

    # Decoder
    for layer_index in range(config.num_decoder_layers	):
        UpperCamelCase__      =			F'layers_{str(lowerCAmelCase_	)}'

        # Self-Attention
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["self_attention"]["key"]["kernel"]
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["self_attention"]["out"]["kernel"]
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["self_attention"]["query"]["kernel"]
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["self_attention"]["value"]["kernel"]

        # Layer Normalization
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["pre_self_attention_layer_norm"][
            "scale"
        ]

        # Encoder-Decoder-Attention
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["key"][
            "kernel"
        ]
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["out"][
            "kernel"
        ]
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["query"][
            "kernel"
        ]
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["value"][
            "kernel"
        ]

        # Layer Normalization
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["pre_cross_attention_layer_norm"]["scale"]

        # MLP
        if split_mlp_wi:
            UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["mlp"]["wi_0"]["kernel"]
            UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["mlp"]["wi_1"]["kernel"]
        else:
            UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["mlp"]["wi"]["kernel"]

        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["mlp"]["wo"]["kernel"]

        # Layer Normalization
        UpperCamelCase__      =			tax_model["target"]["decoder"][layer_name]["pre_mlp_layer_norm"]["scale"]

        # Assigning
        UpperCamelCase__      =			tax_attention_key
        UpperCamelCase__      =			tax_attention_out
        UpperCamelCase__      =			tax_attention_query
        UpperCamelCase__      =			tax_attention_value

        UpperCamelCase__      =			tax_pre_attention_layer_norm

        UpperCamelCase__      =			tax_enc_dec_attention_key
        UpperCamelCase__      =			tax_enc_dec_attention_out
        UpperCamelCase__      =			tax_enc_dec_attention_query
        UpperCamelCase__      =			tax_enc_dec_attention_value

        UpperCamelCase__      =			tax_cross_layer_norm

        if split_mlp_wi:
            UpperCamelCase__      =			tax_mlp_wi_a
            UpperCamelCase__      =			tax_mlp_wi_a
        else:
            UpperCamelCase__      =			tax_mlp_wi

        UpperCamelCase__      =			tax_mlp_wo

        UpperCamelCase__      =			txa_mlp_layer_norm

    # Decoder Normalization
    UpperCamelCase__      =			tax_model["target"]["decoder"]["decoder_norm"]["scale"]
    UpperCamelCase__      =			txa_decoder_norm

    # Only for layer 0:
    UpperCamelCase__      =			tax_model["target"]["decoder"]["relpos_bias"]["rel_embedding"].T
    UpperCamelCase__      =			tax_decoder_rel_embedding

    # Token Embeddings
    UpperCamelCase__      =			tax_model["target"]["token_embedder"]["embedding"]
    UpperCamelCase__      =			txa_token_embeddings

    # LM Head (only in v1.1 checkpoints)
    if "logits_dense" in tax_model["target"]["decoder"]:
        UpperCamelCase__      =			tax_model["target"]["decoder"]["logits_dense"]["kernel"]

    flax_model.save_pretrained(lowerCAmelCase_	)
    print("T5X Model was sucessfully converted!"	)


if __name__ == "__main__":
      lowercase__     :  Any    					=			argparse.ArgumentParser()
      # Required parameters
      parser.add_argument(
          '--t5x_checkpoint_path', default=None, type=str, required=True, help='Path the TX5 checkpoint.'
      )
      parser.add_argument('--config_name', default=None, type=str, required=True, help='Config name of T5 model.')
      parser.add_argument(
          '--flax_dump_folder_path', default=None, type=str, required=True, help='Path to the output FLAX model.'
      )
      lowercase__     :  int    					=			parser.parse_args()
      convert_tax_checkpoint_to_flax(args.tax_checkpoint_path, args.config_name, args.flax_dump_folder_path)




