




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()
SCREAMING_SNAKE_CASE_    :       Any       =							logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_    :       List[Any]       =							[
    (
        '''protein1''',
        '''MNGTEGPNFYVPFSNATGVVRSPFEYPQYYLAEPWQFSMLAAYMFLLIVLGFPINFLTLYVTVQHKKLRTPLNYILLNLAVADLFMVLGGFTSTLYTSLHGYFVFGPTGCNLEGFFATLGGEIALWSLVVLAIERYVVVCKPMSNFRFGENHAIMGVAFTWVMALACAAPPLAGWSRYIPEGLQCSCGIDYYTLKPEVNNESFVIYMFVVHFTIPMIIIFFCYGQLVFTVKEAAAQQQESATTQKAEKEVTRMVIIMVIAFLICWVPYASVAFYIFTHQGSNFGPIFMTIPAFFAKSAAIYNPVIYIMMNKQFRNCMLTTICCGKNPLGDDEASATVSKTETSQVAPA''',
    ),
    ('''protein2''', '''MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLA'''),
    ('''protein3''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLAGG'''),
    ('''protein4''', '''MKTVRQERLKSI<mask>RILERSKEPVSGAQLAEELS<mask>SRQVIVQDIAYLRSLGYN<mask>VATPRGYVLA'''),
]

SCREAMING_SNAKE_CASE_    :       List[str]       =							{
    '''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,
}

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

SCREAMING_SNAKE_CASE_    :       Union[str, Any]       =							restypes + ['''X''']
SCREAMING_SNAKE_CASE_    :       Optional[Any]       =							restypes_with_x + ['''<pad>''', '''<mask>''', '''<cls>''', '''<sep>''', '''<eos>''']
def 							__lowerCAmelCase			():
    with TemporaryDirectory() as tempdir:
        a__  			=							"""\n""".join(UpperCAmelCase_		)
        a__  			=							Path(UpperCAmelCase_		) / """vocab.txt"""
        vocab_file.write_text(UpperCAmelCase_		)
        a__  			=							EsmTokenizer(vocab_file=str(UpperCAmelCase_		)		)
    a__  			=							0  # Overlaps with 'A' but that seems to be what they want
    return hf_tokenizer
def 							__lowerCAmelCase			(UpperCAmelCase_ ,     UpperCAmelCase_		):
    a__  			=							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 							__lowerCAmelCase			(UpperCAmelCase_ ,     UpperCAmelCase_ ,     UpperCAmelCase_ ,     UpperCAmelCase_ ,     UpperCAmelCase_		):
    if model.startswith("""esmfold"""		):
        a__  			=							MODEL_MAPPING[model]()
    else:
        a__			,		a__  			=							MODEL_MAPPING[model]()
    esm.eval()  # disable dropout

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

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

    if is_folding_model:
        a__  			=							esm.esm
    else:
        a__  			=							esm

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

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

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

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

    a__  			=							original_esm_model.emb_layer_norm_after.weight
    a__  			=							original_esm_model.emb_layer_norm_after.bias

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

        # self attention
        a__  			=							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)		)
        )

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

        if getattr(esm_layer.self_attn ,     """rot_emb""" ,     UpperCAmelCase_		) 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.
            a__  			=							esm_layer.self_attn.rot_emb.inv_freq

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

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

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

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

    if is_folding_model:
        a__  			=							esm.esm_s_combine.data
        a__  			=							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:
        a__  			=							esm.esm.classification_heads["""mnli"""].dense.weight
        a__  			=							esm.classification_heads["""mnli"""].dense.bias
        a__  			=							esm.classification_heads["""mnli"""].out_proj.weight
        a__  			=							esm.classification_heads["""mnli"""].out_proj.bias
    else:
        # LM Head
        a__  			=							esm.lm_head.dense.weight
        a__  			=							esm.lm_head.dense.bias
        a__  			=							esm.lm_head.layer_norm.weight
        a__  			=							esm.lm_head.layer_norm.bias
        a__  			=							esm.lm_head.weight
        a__  			=							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.
        a__  			=							SAMPLE_DATA[:2]
    else:
        a__  			=							SAMPLE_DATA

    if is_folding_model:
        a__  			=							get_esmfold_tokenizer()
        a__  			=							hf_tokenizer(
            [row[1] for row in sample_data] ,     return_tensors="""pt""" ,     padding=UpperCAmelCase_ ,     add_special_tokens=UpperCAmelCase_		)
        a__			,		a__			,		a__			,		a__			,		a__  			=							esmfold_encode_sequences([row[1] for row in sample_data]		)
        a__  			=							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.
        a__  			=							alphabet.get_batch_converter()
        a__			,		a__			,		a__  			=							batch_converter(UpperCAmelCase_		)
        # Prepare tokenizer and make sure it matches
        with TemporaryDirectory() as tempdir:
            a__  			=							"""\n""".join(alphabet.all_toks		)
            a__  			=							Path(UpperCAmelCase_		) / """vocab.txt"""
            vocab_file.write_text(UpperCAmelCase_		)
            a__  			=							EsmTokenizer(vocab_file=str(UpperCAmelCase_		)		)

        a__  			=							hf_tokenizer([row[1] for row in sample_data] ,     return_tensors="""pt""" ,     padding=UpperCAmelCase_		)
        a__  			=							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.
            a__  			=							esm.cuda().infer([row[1] for row in sample_data]		)
            a__  			=							model.cuda()(
                input_ids=hf_tokens["""input_ids"""].cuda() ,     attention_mask=hf_tokens["""attention_mask"""].cuda()		)
        else:
            a__  			=							model(**UpperCAmelCase_ ,     output_hidden_states=UpperCAmelCase_		)
            a__  			=							our_output["""logits"""]
            if classification_head:
                a__  			=							esm.model.classification_heads["""mnli"""](esm.extract_features(UpperCAmelCase_		)		)
            else:
                a__  			=							esm(hf_tokens["""input_ids"""] ,     repr_layers=list(range(999		)		)		)
                a__  			=							their_output["""logits"""]

        if is_folding_model:
            a__  			=							torch.max(torch.abs(our_output["""positions"""] - their_output["""positions"""]		)		).item()
            a__  			=							torch.allclose(our_output["""positions"""] ,     their_output["""positions"""] ,     atol=1E-5		)
        else:
            a__  			=							torch.max(torch.abs(our_output - their_output		)		).item()
            a__  			=							torch.allclose(UpperCAmelCase_ ,     UpperCAmelCase_ ,     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
            a__  			=							model.predict_contacts(hf_tokens["""input_ids"""] ,     hf_tokens["""attention_mask"""]		)
            a__  			=							esm.predict_contacts(hf_tokens["""input_ids"""]		)
            a__  			=							torch.max(torch.abs(our_output - their_output		)		).item()
            a__  			=							torch.allclose(UpperCAmelCase_ ,     UpperCAmelCase_ ,     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(UpperCAmelCase_		).mkdir(parents=UpperCAmelCase_ ,     exist_ok=UpperCAmelCase_		)
        print(F"""Saving model to {pytorch_dump_folder_path}"""		)
        model.save_pretrained(UpperCAmelCase_		)

        del esm  # Free up some memory before continuing

    print(F"""Saving tokenizer to {pytorch_dump_folder_path}"""		)
    hf_tokenizer.save_pretrained(UpperCAmelCase_		)

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


if __name__ == "__main__":
 SCREAMING_SNAKE_CASE_    :       List[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.''')
 SCREAMING_SNAKE_CASE_    :       Dict       =							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
 )












import math
from collections import OrderedDict

import torch
from packaging import version
from torch import Tensor, nn

from .utils import logging


SCREAMING_SNAKE_CASE_    :       Tuple       =							logging.get_logger(__name__)







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def __init__(   self			:							int     )      ->      Union[str, Any]:






         """simple docstring"""
         super().__init__()
         if version.parse(torch.__version__     ) < version.parse("""1.12.0"""     ):
             raise ImportError(
                 f"""You are using torch=={torch.__version__}, but torch>=1.12.0 is required to use """
                 """PytorchGELUTanh. Please upgrade torch."""     )






     def 					__SCREAMING_SNAKE_CASE		(   self			:							str    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return nn.functional.gelu(A    ,					approximate="""tanh"""     )







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def 					__SCREAMING_SNAKE_CASE		(   self			:							str    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi     ) * (input + 0.0_4_4_7_1_5 * torch.pow(A    ,					3.0     ))     ))







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def __init__(   self			:							Optional[Any]    ,					A			:							bool = False     )      ->      Optional[Any]:






         """simple docstring"""
         super().__init__()
         if use_gelu_python:
             a__  			=							self._gelu_python
         else:
             a__  			=							nn.functional.gelu

     def 					__SCREAMING_SNAKE_CASE		(   self			:							Dict    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return input * 0.5 * (1.0 + torch.erf(input / math.sqrt(2.0     )     ))






     def 					__SCREAMING_SNAKE_CASE		(   self			:							Union[str, Any]    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return self.act(A     )







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def 					__SCREAMING_SNAKE_CASE		(   self			:							Any    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return 0.5 * input * (1.0 + torch.tanh(input * 0.7_9_7_8_8_4_5_6_0_8 * (1.0 + 0.0_4_4_7_1_5 * input * input)     ))







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def 					__SCREAMING_SNAKE_CASE		(   self			:							List[Any]    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return input * torch.sigmoid(1.7_0_2 * input     )







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def __init__(   self			:							Tuple    ,					A			:							float    ,					A			:							float     )      ->      Optional[Any]:






         """simple docstring"""
         if min > max:
             raise ValueError(f"""min should be < max (got min: {min}, max: {max})"""     )

         super().__init__()
         a__  			=							min
         a__  			=							max






     def 					__SCREAMING_SNAKE_CASE		(   self			:							Optional[int]    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return torch.clip(gelu(A     )    ,					self.min    ,					self.max     )







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def __init__(   self			:							Tuple     )      ->      List[Any]:






         """simple docstring"""
         super().__init__()
         a__  			=							math.sqrt(2 / math.pi     )






     def 					__SCREAMING_SNAKE_CASE		(   self			:							int    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return 0.5 * input * (1 + torch.tanh(self.precomputed_constant * (input + 0.0_4_4_7_1_5 * torch.pow(A    ,					3     ))     ))







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def 					__SCREAMING_SNAKE_CASE		(   self			:							Optional[int]    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return nn.functional.silu(A     )







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def __init__(   self			:							Optional[Any]     )      ->      List[str]:






         """simple docstring"""
         super().__init__()
         if version.parse(torch.__version__     ) < version.parse("""1.9.0"""     ):
             a__  			=							self._mish_python
         else:
             a__  			=							nn.functional.mish

     def 					__SCREAMING_SNAKE_CASE		(   self			:							Any    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return input * torch.tanh(nn.functional.softplus(A     )     )






     def 					__SCREAMING_SNAKE_CASE		(   self			:							Union[str, Any]    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return self.act(A     )







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def 					__SCREAMING_SNAKE_CASE		(   self			:							Optional[Any]    ,					A			:							Tensor     )      ->      Tensor:






         """simple docstring"""
         return input







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def 					__SCREAMING_SNAKE_CASE		(   self			:							List[Any]    ,					A			:							str    ,					A			:							Dict=0.7_0_7_1_0_7    ,					A			:							Union[str, Any]=0.2_8_2_0_9_5     )      ->      int:






         """simple docstring"""
         a__  			=							(input - mu).div(sigma * math.sqrt(2.0     )     )
         return 0.5 * (1.0 + torch.erf(A     ))







class 		__UpperCAmelCase(		nn.Module			):





     """simple docstring"""



     def 					__SCREAMING_SNAKE_CASE		(   self			:							List[str]    ,					A			:							List[Any]     )      ->      Tuple:






         """simple docstring"""
         a__  			=							nn.functional.relu(A     )
         a__  			=							torch.square(A     )
         return squared







class 		__UpperCAmelCase(		lowercase__			):





     """simple docstring"""



     def __getitem__(   self			:							List[str]    ,					A			:							int     )      ->      Any:






         """simple docstring"""
         a__  			=							super().__getitem__(A     )
         a__			,		a__  			=							content if isinstance(A    ,					A     ) else (content, {})
         return cls(**A     )






SCREAMING_SNAKE_CASE_    :       List[Any]       =							{
    '''gelu''': GELUActivation,
    '''gelu_10''': (ClippedGELUActivation, {'''min''': -10, '''max''': 10}),
    '''gelu_fast''': FastGELUActivation,
    '''gelu_new''': NewGELUActivation,
    '''gelu_python''': (GELUActivation, {'''use_gelu_python''': True}),
    '''gelu_pytorch_tanh''': PytorchGELUTanh,
    '''gelu_accurate''': AccurateGELUActivation,
    '''laplace''': LaplaceActivation,
    '''linear''': LinearActivation,
    '''mish''': MishActivation,
    '''quick_gelu''': QuickGELUActivation,
    '''relu''': nn.ReLU,
    '''relu2''': ReLUSquaredActivation,
    '''relu6''': nn.ReLUa,
    '''sigmoid''': nn.Sigmoid,
    '''silu''': SiLUActivation,
    '''swish''': SiLUActivation,
    '''tanh''': nn.Tanh,
}
SCREAMING_SNAKE_CASE_    :       Any       =							ClassInstantier(ACTaCLS)
def 							__lowerCAmelCase			(UpperCAmelCase_		):
    if activation_string in ACTaFN:
        return ACTaFN[activation_string]
    else:
        raise KeyError(F"""function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()		)}"""		)


# For backwards compatibility with: from activations import gelu_python
SCREAMING_SNAKE_CASE_    :       Optional[int]       =							get_activation('''gelu_python''')
SCREAMING_SNAKE_CASE_    :       Union[str, Any]       =							get_activation('''gelu_new''')
SCREAMING_SNAKE_CASE_    :       str       =							get_activation('''gelu''')
SCREAMING_SNAKE_CASE_    :       Optional[Any]       =							get_activation('''gelu_fast''')
SCREAMING_SNAKE_CASE_    :       Union[str, Any]       =							get_activation('''quick_gelu''')
SCREAMING_SNAKE_CASE_    :       Tuple       =							get_activation('''silu''')
SCREAMING_SNAKE_CASE_    :       Union[str, Any]       =							get_activation('''mish''')







SCREAMING_SNAKE_CASE_    :       Tuple       =							get_activation('''linear''')






