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)
UpperCamelCase    				=		'\\n'

UpperCamelCase    				=		'\\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'

UpperCamelCase    				=		'\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'

UpperCamelCase    				=		['fleurs-asr', 'mls', 'voxpopuli', 'babel', 'covost2', 'fleurs-lang_id', 'minds14']
UpperCamelCase    				=		''

try:
      from jiwer import transforms as tr

      UpperCamelCase    				=		True
except ImportError:
      UpperCamelCase    				=		False

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



      class     snake_case_			(						tr.AbstractTransform			):






         def __init__( self							:	Optional[Any]      ,		lowercase_							:	Optional[Any] = " "    )   ->		int:
             lowercase__					:  str          =							sentence_delimiter






         def  __UpperCamelCase	( self							:	Any      ,		lowercase_							:	Union[str, Any]    )   ->		List[str]:
             return list(lowercase_    )







         def  __UpperCamelCase	( self							:	Union[str, Any]      ,		lowercase_							:	Union[str, Any]    )   ->		Dict:
             lowercase__					:  int          =							[]
             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



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

def      lowercase_							(     _lowerCamelCase				:			Optional[Any]						,      _lowerCamelCase				:			Union[str, Any]):
    return float((preds == labels).mean())

def      lowercase_							(     _lowerCamelCase				:			str						,      _lowerCamelCase				:			Union[str, Any]):
    return {
        "f1": float(fa_score(y_true=__A						,      y_pred=__A						,      average="macro")),
        "accuracy": simple_accuracy(__A						,      __A),
    }

def      lowercase_							(     _lowerCamelCase				:			Dict						,      _lowerCamelCase				:			Optional[int]						,      _lowerCamelCase				:			List[Any]="exp"						,      _lowerCamelCase				:			Optional[Any]=None						,      _lowerCamelCase				:			Tuple=False						,      _lowerCamelCase				:			Optional[Any]=False						,      _lowerCamelCase				:			Union[str, Any]=None						,      _lowerCamelCase				:			Any=False						,      ):
    lowercase__					:  List[Any]          =							[[label] for label in labels]
    lowercase__					:  List[str]          =							list(__A)
    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\"`.")

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

def      lowercase_							(     _lowerCamelCase				:			Union[str, Any]						,      _lowerCamelCase				:			str						,      _lowerCamelCase				:			Tuple						,      _lowerCamelCase				:			Dict):



    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:
        lowercase__					:  Union[str, Any]          =							compute_measures(__A						,      __A)["wer"]

        lowercase__					:  Union[str, Any]          =							compute_measures(__A						,      __A						,      truth_transform=__A						,      hypothesis_transform=__A)["wer"]
        return {"wer": wer, "cer": cer}
    else:

        def compute_score(_lowerCamelCase				:			str						,      _lowerCamelCase				:			int						,      _lowerCamelCase				:			Dict="wer"):
            lowercase__					:  Optional[int]          =							0
            lowercase__					:  Dict          =							0
            for prediction, reference in zip(__A						,      __A):
                if score_type == "wer":
                    lowercase__					:  int          =							compute_measures(__A						,      __A)
                elif score_type == "cer":
                    lowercase__					:  Optional[int]          =							compute_measures(
                        __A						,      __A						,      truth_transform=__A						,      hypothesis_transform=__A)
                incorrect += measures["substitutions"] + measures["deletions"] + measures["insertions"]
                total += measures["substitutions"] + measures["deletions"] + measures["hits"]
            return incorrect / total

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



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






   def  __UpperCamelCase	( self							:	Optional[Any]    )   ->		Optional[int]:
       if self.config_name not in _CONFIG_NAMES:
           raise KeyError(F'''You should supply a configuration name selected in {_CONFIG_NAMES}'''    )

       lowercase__					:  Tuple          =							"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  __UpperCamelCase	( self							:	Union[str, Any]      ,		lowercase_							:	Union[str, Any]      ,		lowercase_							:	Tuple      ,		lowercase_							:	str=None      ,		lowercase_							:	Union[str, Any]=None    )   ->		Optional[int]:
       lowercase__					:  Optional[int]          =							bleu_kwargs if bleu_kwargs is not None else {}
       lowercase__					:  Optional[int]          =							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":
           lowercase__					:  str          =							bleu_kwargs.pop("smooth_method"      ,		"exp"    )
           lowercase__					:  Optional[Any]          =							bleu_kwargs.pop("smooth_value"      ,		lowercase_    )
           lowercase__					:  int          =							bleu_kwargs.pop("force"      ,		lowercase_    )
           lowercase__					:  Optional[Any]          =							bleu_kwargs.pop("lowercase"      ,		lowercase_    )
           lowercase__					:  int          =							bleu_kwargs.pop("tokenize"      ,		lowercase_    )
           lowercase__					:  int          =							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"]:
           lowercase__					:  Dict          =							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}'''    )







'''simple docstring'''







import argparse
import pathlib
from pathlib import Path
from tempfile import TemporaryDirectory

import esm as esm_module
import torch
from esm.esmfold.va.misc import batch_encode_sequences as esmfold_encode_sequences
from esm.esmfold.va.pretrained import esmfold_va

from transformers.models.esm.configuration_esm import EsmConfig, EsmFoldConfig
from transformers.models.esm.modeling_esm import (
    EsmForMaskedLM,
    EsmForSequenceClassification,
    EsmIntermediate,
    EsmLayer,
    EsmOutput,
    EsmSelfAttention,
    EsmSelfOutput,
)
from transformers.models.esm.modeling_esmfold import EsmForProteinFolding
from transformers.models.esm.tokenization_esm import EsmTokenizer
from transformers.utils import logging


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

a__    :   Optional[int]								=			[
    (
        'protein1',
        'MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA',
    ),
    ('protein2', 'MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'),
    ('protein3', 'MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'),
    ('protein4', 'MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'),
]

a__    :   Optional[int]								=			{
    'esm1b_t33_650M_UR50S': esm_module.pretrained.esm1b_t33_650M_UR50S,
    'esm1v_t33_650M_UR90S_1': esm_module.pretrained.esm1v_t33_650M_UR90S_1,
    'esm1v_t33_650M_UR90S_2': esm_module.pretrained.esm1v_t33_650M_UR90S_2,
    'esm1v_t33_650M_UR90S_3': esm_module.pretrained.esm1v_t33_650M_UR90S_3,
    'esm1v_t33_650M_UR90S_4': esm_module.pretrained.esm1v_t33_650M_UR90S_4,
    'esm1v_t33_650M_UR90S_5': esm_module.pretrained.esm1v_t33_650M_UR90S_5,
    'esm2_t48_15B_UR50D': esm_module.pretrained.esm2_t48_15B_UR50D,
    'esm2_t36_3B_UR50D': esm_module.pretrained.esm2_t36_3B_UR50D,
    'esm2_t33_650M_UR50D': esm_module.pretrained.esm2_t33_650M_UR50D,
    'esm2_t30_150M_UR50D': esm_module.pretrained.esm2_t30_150M_UR50D,
    'esm2_t12_35M_UR50D': esm_module.pretrained.esm2_t12_35M_UR50D,
    'esm2_t6_8M_UR50D': esm_module.pretrained.esm2_t6_8M_UR50D,
    'esmfold_v1': esmfold_va,
}

a__    :   List[Any]								=			list('ARNDCQEGHILKMFPSTWYV')

a__    :   Any								=			restypes + ['X']
a__    :   List[str]								=			restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']







def _UpperCamelCase			(	)				->						str:


      '''simple docstring'''



      with TemporaryDirectory() as tempdir:
            UpperCamelCase__      =							"\n".join(__A       )
            UpperCamelCase__      =							Path(__A       ) / "vocab.txt"
            vocab_file.write_text(__A       )
            UpperCamelCase__      =							EsmTokenizer(vocab_file=str(__A       )       )
      UpperCamelCase__      =							0  # Overlaps with 'A' but that seems to be what they want
      return hf_tokenizer







def _UpperCamelCase			(	__A ,     __A       )				->						List[str]:


      '''simple docstring'''



      UpperCamelCase__      =							our_module.load_state_dict(original_module.state_dict()       )
      if status.missing_keys:
            raise ValueError(F'''Missing keys: {status.missing_keys}'''       )
      if status.unexpected_keys:
            raise ValueError(F'''Unexpected keys: {status.unexpected_keys}'''       )







def _UpperCamelCase			(	__A ,     __A ,     __A ,     __A ,     __A       )				->						int:


      '''simple docstring'''



      if model.startswith("esmfold"       ):
            UpperCamelCase__      =							MODEL_MAPPING[model]()
      else:
            UpperCamelCase__   ,      UpperCamelCase__      =							MODEL_MAPPING[model]()
      esm.eval()  # disable dropout

      if model.startswith("esmfold"       ):
            UpperCamelCase__      =							esm.esm.embed_dim
            UpperCamelCase__      =							esm.esm.num_layers
            UpperCamelCase__      =							esm.esm.attention_heads
            UpperCamelCase__      =							4 * embed_dim
            UpperCamelCase__      =							esm.esm.token_dropout
            UpperCamelCase__      =							False  # This code path does not exist in ESM-2
            UpperCamelCase__      =							"rotary"
            UpperCamelCase__      =							True
            UpperCamelCase__      =							EsmFoldConfig()
            for key, val in esm.cfg.items():
                  if hasattr(__A ,     __A       ) and key != "trunk":
                        setattr(__A ,     __A ,     __A       )
            for key, val in esm.cfg.trunk.items():
                  if hasattr(esmfold_config.trunk ,     __A       ) and key != "structure_module":
                        setattr(esmfold_config.trunk ,     __A ,     __A       )
            for key, val in esm.cfg.trunk.structure_module.items():
                  if hasattr(esmfold_config.trunk.structure_module ,     __A       ):
                        setattr(esmfold_config.trunk.structure_module ,     __A ,     __A       )
      elif hasattr(__A ,     "args"       ):
            # Indicates an ESM-1b or ESM-1v model
            UpperCamelCase__      =							esm.args.embed_dim
            UpperCamelCase__      =							esm.args.layers
            UpperCamelCase__      =							esm.args.attention_heads
            UpperCamelCase__      =							esm.args.ffn_embed_dim
            UpperCamelCase__      =							esm.args.token_dropout
            UpperCamelCase__      =							True if esm.emb_layer_norm_before else False
            UpperCamelCase__      =							"absolute"
            UpperCamelCase__      =							False
            UpperCamelCase__      =							None
      else:
            # Indicates an ESM-2 model
            UpperCamelCase__      =							esm.embed_dim
            UpperCamelCase__      =							esm.num_layers
            UpperCamelCase__      =							esm.attention_heads
            UpperCamelCase__      =							4 * embed_dim  # This is hardcoded in ESM-2
            UpperCamelCase__      =							esm.token_dropout
            UpperCamelCase__      =							False  # This code path does not exist in ESM-2
            UpperCamelCase__      =							"rotary"
            UpperCamelCase__      =							False
            UpperCamelCase__      =							None

      if is_folding_model:
            UpperCamelCase__      =							esm.esm.alphabet
      UpperCamelCase__      =							tuple(alphabet.all_toks       )
      UpperCamelCase__      =							alphabet.mask_idx
      UpperCamelCase__      =							alphabet.padding_idx

      if is_folding_model:
            UpperCamelCase__      =							esm.esm
      else:
            UpperCamelCase__      =							esm

      UpperCamelCase__      =							EsmConfig(
          vocab_size=original_esm_model.embed_tokens.num_embeddings ,     mask_token_id=__A ,     hidden_size=__A ,     num_hidden_layers=__A ,     num_attention_heads=__A ,     intermediate_size=__A ,     max_position_embeddings=1026 ,     layer_norm_eps=1E-5 ,     attention_probs_dropout_prob=0.0 ,     hidden_dropout_prob=0.0 ,     pad_token_id=__A ,     emb_layer_norm_before=__A ,     token_dropout=__A ,     position_embedding_type=__A ,     is_folding_model=__A ,     esmfold_config=__A ,     vocab_list=__A ,     )
      if classification_head:
            UpperCamelCase__      =							esm.classification_heads["mnli"].out_proj.weight.shape[0]
      print("Our ESM config:" ,     __A       )

      if model.startswith("esmfold"       ):
            UpperCamelCase__      =							EsmForProteinFolding
      elif classification_head:
            UpperCamelCase__      =							EsmForSequenceClassification
      else:
            UpperCamelCase__      =							EsmForMaskedLM
      UpperCamelCase__      =							model_class(__A       )
      model.eval()

      # Now let's copy all the weights.
      # Embeddings
      UpperCamelCase__      =							original_esm_model.embed_tokens.weight
      if position_embedding_type == "absolute":
            UpperCamelCase__      =							original_esm_model.embed_positions.weight

      if config.emb_layer_norm_before:
            UpperCamelCase__      =							original_esm_model.emb_layer_norm_before.weight
            UpperCamelCase__      =							original_esm_model.emb_layer_norm_before.bias

      UpperCamelCase__      =							original_esm_model.emb_layer_norm_after.weight
      UpperCamelCase__      =							original_esm_model.emb_layer_norm_after.bias

      for i in range(config.num_hidden_layers       ):
            # Encoder: start of layer
            UpperCamelCase__      =							model.esm.encoder.layer[i]
            # esm_layer: TransformerSentenceEncoderLayer = original_esm_model.layers[i]
            UpperCamelCase__      =							original_esm_model.layers[i]

            # self attention
            UpperCamelCase__      =							layer.attention.self
            assert (
                esm_layer.self_attn.k_proj.weight.data.shape
                == esm_layer.self_attn.q_proj.weight.data.shape
                == esm_layer.self_attn.v_proj.weight.data.shape
                == torch.Size((config.hidden_size, config.hidden_size)       )
            )

            UpperCamelCase__      =							esm_layer.self_attn.q_proj.weight
            UpperCamelCase__      =							esm_layer.self_attn.q_proj.bias
            UpperCamelCase__      =							esm_layer.self_attn.k_proj.weight
            UpperCamelCase__      =							esm_layer.self_attn.k_proj.bias
            UpperCamelCase__      =							esm_layer.self_attn.v_proj.weight
            UpperCamelCase__      =							esm_layer.self_attn.v_proj.bias

            if getattr(esm_layer.self_attn ,     "rot_emb" ,     __A       ) is not None:
                  # Matt: Although inv_freq is not a trainable weight, it is computed at model init and cached.
                  # During the training of ESM-2 the model was converted to float16 precision, which also converts
                  # the inv_freq tensor, and the loss of precision remains even if the model is loaded later as float32.
                  # If we recompute inv_freq without this loss of precision then we will get subtly different rotary
                  # embeddings, which are enough to cause significant discrepancies in model outputs. To avoid this,
                  # we make sure the new model copies the data from the old inv_freq.
                  UpperCamelCase__      =							esm_layer.self_attn.rot_emb.inv_freq

            # LayerNorm changes for pre-activation
            UpperCamelCase__      =							esm_layer.self_attn_layer_norm.weight
            UpperCamelCase__      =							esm_layer.self_attn_layer_norm.bias
            UpperCamelCase__      =							esm_layer.final_layer_norm.weight
            UpperCamelCase__      =							esm_layer.final_layer_norm.bias

            # self-attention output
            UpperCamelCase__      =							layer.attention.output
            assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
            UpperCamelCase__      =							esm_layer.self_attn.out_proj.weight
            UpperCamelCase__      =							esm_layer.self_attn.out_proj.bias

            # intermediate
            UpperCamelCase__      =							layer.intermediate
            assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
            UpperCamelCase__      =							esm_layer.fca.weight
            UpperCamelCase__      =							esm_layer.fca.bias

            # output
            UpperCamelCase__      =							layer.output
            assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
            UpperCamelCase__      =							esm_layer.fca.weight
            UpperCamelCase__      =							esm_layer.fca.bias
            # end of layer

      if is_folding_model:
            UpperCamelCase__      =							esm.esm_s_combine.data
            UpperCamelCase__      =							esm.afa_to_esm.data
            transfer_and_check_weights(esm.embedding ,     model.embedding       )
            transfer_and_check_weights(esm.esm_s_mlp ,     model.esm_s_mlp       )
            transfer_and_check_weights(esm.trunk ,     model.trunk       )
            transfer_and_check_weights(esm.distogram_head ,     model.distogram_head       )
            transfer_and_check_weights(esm.ptm_head ,     model.ptm_head       )
            transfer_and_check_weights(esm.lm_head ,     model.lm_head       )
            transfer_and_check_weights(esm.lddt_head ,     model.lddt_head       )

      elif classification_head:
            UpperCamelCase__      =							esm.esm.classification_heads["mnli"].dense.weight
            UpperCamelCase__      =							esm.classification_heads["mnli"].dense.bias
            UpperCamelCase__      =							esm.classification_heads["mnli"].out_proj.weight
            UpperCamelCase__      =							esm.classification_heads["mnli"].out_proj.bias
      else:
            # LM Head
            UpperCamelCase__      =							esm.lm_head.dense.weight
            UpperCamelCase__      =							esm.lm_head.dense.bias
            UpperCamelCase__      =							esm.lm_head.layer_norm.weight
            UpperCamelCase__      =							esm.lm_head.layer_norm.bias
            UpperCamelCase__      =							esm.lm_head.weight
            UpperCamelCase__      =							esm.lm_head.bias

      # Contact prediction head
      transfer_and_check_weights(esm.contact_head ,     model.esm.contact_head       )

      # Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
      if is_folding_model:
            # Folding models aren't trained on masked inputs and don't like mask tokens.
            UpperCamelCase__      =							SAMPLE_DATA[:2]
      else:
            UpperCamelCase__      =							SAMPLE_DATA

      if is_folding_model:
            UpperCamelCase__      =							get_esmfold_tokenizer()
            UpperCamelCase__      =							hf_tokenizer(
                [row[1] for row in sample_data] ,     return_tensors="pt" ,     padding=__A ,     add_special_tokens=__A       )
            UpperCamelCase__   ,      UpperCamelCase__   ,      UpperCamelCase__   ,      UpperCamelCase__   ,      UpperCamelCase__      =							esmfold_encode_sequences([row[1] for row in sample_data]       )
            UpperCamelCase__      =							torch.all(hf_tokens["input_ids"] == esmfold_aas       ) and torch.all(
                hf_tokens["attention_mask"] == esmfold_mask       )
      else:
            # Let's check that we get the same results.
            UpperCamelCase__      =							alphabet.get_batch_converter()
            UpperCamelCase__   ,      UpperCamelCase__   ,      UpperCamelCase__      =							batch_converter(__A       )
            # Prepare tokenizer and make sure it matches
            with TemporaryDirectory() as tempdir:
                  UpperCamelCase__      =							"\n".join(alphabet.all_toks       )
                  UpperCamelCase__      =							Path(__A       ) / "vocab.txt"
                  vocab_file.write_text(__A       )
                  UpperCamelCase__      =							EsmTokenizer(vocab_file=str(__A       )       )

            UpperCamelCase__      =							hf_tokenizer([row[1] for row in sample_data] ,     return_tensors="pt" ,     padding=__A       )
            UpperCamelCase__      =							torch.all(hf_tokens["input_ids"] == batch_tokens       )

      print("Do both models tokenizers output the same tokens?" ,     "🔥" if success else "💩"       )
      if not success:
            raise Exception("Tokenization does not match!"       )

      with torch.no_grad():
            if is_folding_model:
                  # Let's test the model in parts
                  # ESMFold always converts the ESM stem to float16, which requires float16 ops
                  # that don't exist on CPU. Therefore, to test it we need to run it on GPU. However,
                  # ESMFold is what we in the community call a "big boy" and so we desperately avoid putting both the
                  # original and the converted model on the GPU at the same time.
                  UpperCamelCase__      =							esm.cuda().infer([row[1] for row in sample_data]       )
                  UpperCamelCase__      =							model.cuda()(
                      input_ids=hf_tokens["input_ids"].cuda() ,     attention_mask=hf_tokens["attention_mask"].cuda()       )
            else:
                  UpperCamelCase__      =							model(**__A ,     output_hidden_states=__A       )
                  UpperCamelCase__      =							our_output["logits"]
                  if classification_head:
                        UpperCamelCase__      =							esm.model.classification_heads["mnli"](esm.extract_features(__A       )       )
                  else:
                        UpperCamelCase__      =							esm(hf_tokens["input_ids"] ,     repr_layers=list(range(999       )       )       )
                        UpperCamelCase__      =							their_output["logits"]

            if is_folding_model:
                  UpperCamelCase__      =							torch.max(torch.abs(our_output["positions"] - their_output["positions"]       )       ).item()
                  UpperCamelCase__      =							torch.allclose(our_output["positions"] ,     their_output["positions"] ,     atol=1E-5       )
            else:
                  UpperCamelCase__      =							torch.max(torch.abs(our_output - their_output       )       ).item()
                  UpperCamelCase__      =							torch.allclose(__A ,     __A ,     atol=1E-5       )

            print(F'''max_absolute_diff = {max_absolute_diff}'''       )  # ~ 1e-5
            print("Do both models output the same tensors?" ,     "🔥" if success else "💩"       )

            if not success:
                  raise Exception("Something went wRoNg"       )

            if not is_folding_model:
                  # Let's check contact prediction too
                  UpperCamelCase__      =							model.predict_contacts(hf_tokens["input_ids"] ,     hf_tokens["attention_mask"]       )
                  UpperCamelCase__      =							esm.predict_contacts(hf_tokens["input_ids"]       )
                  UpperCamelCase__      =							torch.max(torch.abs(our_output - their_output       )       ).item()
                  UpperCamelCase__      =							torch.allclose(__A ,     __A ,     atol=1E-5       )

                  print("Contact prediction testing:"       )
                  print(F'''max_absolute_diff = {max_absolute_diff}'''       )  # ~ 1e-5
                  print("Do both models output the same tensors?" ,     "🔥" if success else "💩"       )

                  if not success:
                        raise Exception("Something went wRoNg"       )

            pathlib.Path(__A       ).mkdir(parents=__A ,     exist_ok=__A       )
            print(F'''Saving model to {pytorch_dump_folder_path}'''       )
            model.save_pretrained(__A       )

            del esm  # Free up some memory before continuing

      print(F'''Saving tokenizer to {pytorch_dump_folder_path}'''       )
      hf_tokenizer.save_pretrained(__A       )

      if push_to_repo:
            model.push_to_hub(repo_id=__A ,     use_auth_token=__A       )
            hf_tokenizer.push_to_hub(repo_id=__A ,     use_auth_token=__A       )


if __name__ == "__main__":
  a__    :   str								=			argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      '--pytorch_dump_folder_path', type=str, required=True, help='Path to the output PyTorch model.'
  )
  parser.add_argument(
      '--classification_head', action='store_true', help='Whether to convert a final classification head.'
  )
  parser.add_argument('--model', default=None, type=str, required=True, help='Name of model to convert.')
  parser.add_argument('--push_to_repo', type=str, help='Repo to upload to (including username!).')
  parser.add_argument('--auth_token', type=str, help='HuggingFace auth token.')
  a__    :   Optional[Any]								=			parser.parse_args()
  convert_esm_checkpoint_to_pytorch(
      args.model, args.pytorch_dump_folder_path, args.classification_head, args.push_to_repo, args.auth_token
  )





