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

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

a__				:      Tuple    		=     {
    "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__				:      Any    		=     list('''ARNDCQEGHILKMFPSTWYV''')

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







def 				UpperCAmelCase_(    ):



    """simple docstring"""



    with TemporaryDirectory() as tempdir:
        SCREAMING_SNAKE_CASE				:				int				      = '''\n'''.join(a__   )
        SCREAMING_SNAKE_CASE				:				Any				      = Path(a__   ) / '''vocab.txt'''
        vocab_file.write_text(a__   )
        SCREAMING_SNAKE_CASE				:				Tuple				      = EsmTokenizer(vocab_file=str(a__   )   )
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = 0  # Overlaps with 'A' but that seems to be what they want
    return hf_tokenizer







def 				UpperCAmelCase_(    a__ ,						a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[int]				      = 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__   ):



    """simple docstring"""



    if model.startswith('''esmfold'''   ):
        SCREAMING_SNAKE_CASE				:				List[str]				      = MODEL_MAPPING[model]()
    else:
        SCREAMING_SNAKE_CASE      ,							SCREAMING_SNAKE_CASE				:				Dict				      = MODEL_MAPPING[model]()
    esm.eval()  # disable dropout

    if model.startswith('''esmfold'''   ):
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm.esm.embed_dim
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = esm.esm.num_layers
        SCREAMING_SNAKE_CASE				:				int				      = esm.esm.attention_heads
        SCREAMING_SNAKE_CASE				:				int				      = 4 * embed_dim
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = esm.esm.token_dropout
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = False  # This code path does not exist in ESM-2
        SCREAMING_SNAKE_CASE				:				Any				      = '''rotary'''
        SCREAMING_SNAKE_CASE				:				str				      = True
        SCREAMING_SNAKE_CASE				:				Dict				      = 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
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = esm.args.embed_dim
        SCREAMING_SNAKE_CASE				:				Tuple				      = esm.args.layers
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = esm.args.attention_heads
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm.args.ffn_embed_dim
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = esm.args.token_dropout
        SCREAMING_SNAKE_CASE				:				str				      = True if esm.emb_layer_norm_before else False
        SCREAMING_SNAKE_CASE				:				Tuple				      = '''absolute'''
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = False
        SCREAMING_SNAKE_CASE				:				Any				      = None
    else:
        # Indicates an ESM-2 model
        SCREAMING_SNAKE_CASE				:				Dict				      = esm.embed_dim
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = esm.num_layers
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm.attention_heads
        SCREAMING_SNAKE_CASE				:				int				      = 4 * embed_dim  # This is hardcoded in ESM-2
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm.token_dropout
        SCREAMING_SNAKE_CASE				:				List[Any]				      = False  # This code path does not exist in ESM-2
        SCREAMING_SNAKE_CASE				:				int				      = '''rotary'''
        SCREAMING_SNAKE_CASE				:				str				      = False
        SCREAMING_SNAKE_CASE				:				Dict				      = None

    if is_folding_model:
        SCREAMING_SNAKE_CASE				:				List[Any]				      = esm.esm.alphabet
    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = tuple(alphabet.all_toks   )
    SCREAMING_SNAKE_CASE				:				int				      = alphabet.mask_idx
    SCREAMING_SNAKE_CASE				:				str				      = alphabet.padding_idx

    if is_folding_model:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = esm.esm
    else:
        SCREAMING_SNAKE_CASE				:				Any				      = esm

    SCREAMING_SNAKE_CASE				:				List[str]				      = 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=1_026 ,						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:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = esm.classification_heads['''mnli'''].out_proj.weight.shape[0]
    print('''Our ESM config:''' ,						a__   )

    if model.startswith('''esmfold'''   ):
        SCREAMING_SNAKE_CASE				:				int				      = EsmForProteinFolding
    elif classification_head:
        SCREAMING_SNAKE_CASE				:				Any				      = EsmForSequenceClassification
    else:
        SCREAMING_SNAKE_CASE				:				int				      = EsmForMaskedLM
    SCREAMING_SNAKE_CASE				:				Tuple				      = model_class(a__   )
    model.eval()

    # Now let's copy all the weights.
    # Embeddings
    SCREAMING_SNAKE_CASE				:				List[str]				      = original_esm_model.embed_tokens.weight
    if position_embedding_type == "absolute":
        SCREAMING_SNAKE_CASE				:				str				      = original_esm_model.embed_positions.weight

    if config.emb_layer_norm_before:
        SCREAMING_SNAKE_CASE				:				List[Any]				      = original_esm_model.emb_layer_norm_before.weight
        SCREAMING_SNAKE_CASE				:				int				      = original_esm_model.emb_layer_norm_before.bias

    SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = original_esm_model.emb_layer_norm_after.weight
    SCREAMING_SNAKE_CASE				:				int				      = original_esm_model.emb_layer_norm_after.bias

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

        # self attention
        SCREAMING_SNAKE_CASE				:				List[str]				      = 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)   )
        )

        SCREAMING_SNAKE_CASE				:				int				      = esm_layer.self_attn.q_proj.weight
        SCREAMING_SNAKE_CASE				:				Tuple				      = esm_layer.self_attn.q_proj.bias
        SCREAMING_SNAKE_CASE				:				List[Any]				      = esm_layer.self_attn.k_proj.weight
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = esm_layer.self_attn.k_proj.bias
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = esm_layer.self_attn.v_proj.weight
        SCREAMING_SNAKE_CASE				:				Tuple				      = 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.
            SCREAMING_SNAKE_CASE				:				str				      = esm_layer.self_attn.rot_emb.inv_freq

        # LayerNorm changes for pre-activation
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = esm_layer.self_attn_layer_norm.weight
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = esm_layer.self_attn_layer_norm.bias
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm_layer.final_layer_norm.weight
        SCREAMING_SNAKE_CASE				:				int				      = esm_layer.final_layer_norm.bias

        # self-attention output
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = layer.attention.output
        assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
        SCREAMING_SNAKE_CASE				:				List[Any]				      = esm_layer.self_attn.out_proj.weight
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm_layer.self_attn.out_proj.bias

        # intermediate
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = layer.intermediate
        assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = esm_layer.fca.weight
        SCREAMING_SNAKE_CASE				:				int				      = esm_layer.fca.bias

        # output
        SCREAMING_SNAKE_CASE				:				int				      = layer.output
        assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
        SCREAMING_SNAKE_CASE				:				Any				      = esm_layer.fca.weight
        SCREAMING_SNAKE_CASE				:				str				      = esm_layer.fca.bias
        # end of layer

    if is_folding_model:
        SCREAMING_SNAKE_CASE				:				Dict				      = esm.esm_s_combine.data
        SCREAMING_SNAKE_CASE				:				int				      = 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:
        SCREAMING_SNAKE_CASE				:				Tuple				      = esm.esm.classification_heads['''mnli'''].dense.weight
        SCREAMING_SNAKE_CASE				:				Dict				      = esm.classification_heads['''mnli'''].dense.bias
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = esm.classification_heads['''mnli'''].out_proj.weight
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = esm.classification_heads['''mnli'''].out_proj.bias
    else:
        # LM Head
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm.lm_head.dense.weight
        SCREAMING_SNAKE_CASE				:				Tuple				      = esm.lm_head.dense.bias
        SCREAMING_SNAKE_CASE				:				str				      = esm.lm_head.layer_norm.weight
        SCREAMING_SNAKE_CASE				:				str				      = esm.lm_head.layer_norm.bias
        SCREAMING_SNAKE_CASE				:				List[str]				      = esm.lm_head.weight
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = 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.
        SCREAMING_SNAKE_CASE				:				List[str]				      = SAMPLE_DATA[:2]
    else:
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = SAMPLE_DATA

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

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

        if is_folding_model:
            SCREAMING_SNAKE_CASE				:				str				      = torch.max(torch.abs(our_output['''positions'''] - their_output['''positions''']   )   ).item()
            SCREAMING_SNAKE_CASE				:				str				      = torch.allclose(our_output['''positions'''] ,						their_output['''positions'''] ,						atol=1e-5   )
        else:
            SCREAMING_SNAKE_CASE				:				Optional[Any]				      = torch.max(torch.abs(our_output - their_output   )   ).item()
            SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = 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
            SCREAMING_SNAKE_CASE				:				List[Any]				      = model.predict_contacts(hf_tokens['''input_ids'''] ,						hf_tokens['''attention_mask''']   )
            SCREAMING_SNAKE_CASE				:				Tuple				      = esm.predict_contacts(hf_tokens['''input_ids''']   )
            SCREAMING_SNAKE_CASE				:				Any				      = torch.max(torch.abs(our_output - their_output   )   ).item()
            SCREAMING_SNAKE_CASE				:				Dict				      = 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__				:      Tuple    		=     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__				:      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
      )


import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import (
				    MODEL_FOR_PRETRAINING_MAPPING,
				    FunnelBaseModel,
				    FunnelForMaskedLM,
				    FunnelForMultipleChoice,
				    FunnelForPreTraining,
				    FunnelForQuestionAnswering,
				    FunnelForSequenceClassification,
				    FunnelForTokenClassification,
				    FunnelModel,
				)






class 							SCREAMING_SNAKE_CASE__     :

							"""simple docstring"""






							def __init__(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_=13  ,						SCREAMING_SNAKE_CASE_=7  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=True  ,						SCREAMING_SNAKE_CASE_=99  ,						SCREAMING_SNAKE_CASE_=[1, 1, 2]  ,						SCREAMING_SNAKE_CASE_=1  ,						SCREAMING_SNAKE_CASE_=32  ,						SCREAMING_SNAKE_CASE_=4  ,						SCREAMING_SNAKE_CASE_=8  ,						SCREAMING_SNAKE_CASE_=37  ,						SCREAMING_SNAKE_CASE_="gelu_new"  ,						SCREAMING_SNAKE_CASE_=0.1  ,						SCREAMING_SNAKE_CASE_=0.1  ,						SCREAMING_SNAKE_CASE_=0.0  ,						SCREAMING_SNAKE_CASE_=512  ,						SCREAMING_SNAKE_CASE_=3  ,						SCREAMING_SNAKE_CASE_=0.0_2  ,						SCREAMING_SNAKE_CASE_=3  ,						SCREAMING_SNAKE_CASE_=4  ,						SCREAMING_SNAKE_CASE_=None  ,						SCREAMING_SNAKE_CASE_=False  ,						)-> str:





													'''simple docstring'''






													__UpperCamelCase  	=      parent
													__UpperCamelCase  	=      batch_size
													__UpperCamelCase  	=      seq_length
													__UpperCamelCase  	=      is_training
													__UpperCamelCase  	=      use_input_mask
													__UpperCamelCase  	=      use_token_type_ids
													__UpperCamelCase  	=      use_labels
													__UpperCamelCase  	=      vocab_size
													__UpperCamelCase  	=      block_sizes
													__UpperCamelCase  	=      num_decoder_layers
													__UpperCamelCase  	=      d_model
													__UpperCamelCase  	=      n_head
													__UpperCamelCase  	=      d_head
													__UpperCamelCase  	=      d_inner
													__UpperCamelCase  	=      hidden_act
													__UpperCamelCase  	=      hidden_dropout
													__UpperCamelCase  	=      attention_dropout
													__UpperCamelCase  	=      activation_dropout
													__UpperCamelCase  	=      max_position_embeddings
													__UpperCamelCase  	=      type_vocab_size
													__UpperCamelCase  	=      2
													__UpperCamelCase  	=      num_labels
													__UpperCamelCase  	=      num_choices
													__UpperCamelCase  	=      scope
													__UpperCamelCase  	=      initializer_std

													# Used in the tests to check the size of the first attention layer
													__UpperCamelCase  	=      n_head
													# Used in the tests to check the size of the first hidden state
													__UpperCamelCase  	=      self.d_model
													# Used in the tests to check the number of output hidden states/attentions
													__UpperCamelCase  	=      sum(self.block_sizes      ) + (0 if base else self.num_decoder_layers)
													# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
													# the last hidden state of the first block (which is the first hidden state of the decoder).
													if not base:
																			__UpperCamelCase  	=      self.num_hidden_layers + 2



							def      A__		(	self      )-> Any:





													'''simple docstring'''






													__UpperCamelCase  	=      ids_tensor([self.batch_size, self.seq_length]  ,						self.vocab_size      )

													__UpperCamelCase  	=      None
													if self.use_input_mask:
																			__UpperCamelCase  	=      ids_tensor([self.batch_size, self.seq_length]  ,						vocab_size=2      )

													__UpperCamelCase  	=      None
													if self.use_token_type_ids:
																			__UpperCamelCase  	=      ids_tensor([self.batch_size, self.seq_length]  ,						self.type_vocab_size      )

													__UpperCamelCase  	=      None
													__UpperCamelCase  	=      None
													__UpperCamelCase  	=      None
													if self.use_labels:
																			__UpperCamelCase  	=      ids_tensor([self.batch_size]  ,						self.type_sequence_label_size      )
																			__UpperCamelCase  	=      ids_tensor([self.batch_size, self.seq_length]  ,						self.num_labels      )
																			__UpperCamelCase  	=      ids_tensor([self.batch_size]  ,						self.num_choices      )
																			__UpperCamelCase  	=      ids_tensor([self.batch_size, self.seq_length]  ,						1      )

													__UpperCamelCase  	=      self.get_config()

													return (
													    config,
													    input_ids,
													    token_type_ids,
													    input_mask,
													    sequence_labels,
													    token_labels,
													    choice_labels,
													    fake_token_labels,
													)



							def      A__		(	self      )-> List[str]:





													'''simple docstring'''






													return FunnelConfig(
													    vocab_size=self.vocab_size  ,						block_sizes=self.block_sizes  ,						num_decoder_layers=self.num_decoder_layers  ,						d_model=self.d_model  ,						n_head=self.n_head  ,						d_head=self.d_head  ,						d_inner=self.d_inner  ,						hidden_act=self.hidden_act  ,						hidden_dropout=self.hidden_dropout  ,						attention_dropout=self.attention_dropout  ,						activation_dropout=self.activation_dropout  ,						max_position_embeddings=self.max_position_embeddings  ,						type_vocab_size=self.type_vocab_size  ,						initializer_std=self.initializer_std  ,						)



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> str:





													'''simple docstring'''






													__UpperCamelCase  	=      FunnelModel(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.d_model)      )

													__UpperCamelCase  	=      False
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.d_model)      )

													__UpperCamelCase  	=      False
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.d_model)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      FunnelBaseModel(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, 2, self.d_model)      )

													__UpperCamelCase  	=      False
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, 3, self.d_model)      )

													__UpperCamelCase  	=      False
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, 2, self.d_model)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> List[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_labels
													__UpperCamelCase  	=      FunnelForPreTraining(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      FunnelForMaskedLM(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.vocab_size)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> List[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_labels
													__UpperCamelCase  	=      FunnelForSequenceClassification(SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_choices
													__UpperCamelCase  	=      FunnelForMultipleChoice(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      input_ids.unsqueeze(1      ).expand(-1  ,						self.num_choices  ,						-1      ).contiguous()
													__UpperCamelCase  	=      token_type_ids.unsqueeze(1      ).expand(-1  ,						self.num_choices  ,						-1      ).contiguous()
													__UpperCamelCase  	=      input_mask.unsqueeze(1      ).expand(-1  ,						self.num_choices  ,						-1      ).contiguous()
													__UpperCamelCase  	=      model(
													    SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_  ,						)
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_choices)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.num_labels
													__UpperCamelCase  	=      FunnelForTokenClassification(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						labels=SCREAMING_SNAKE_CASE_      )
													self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.num_labels)      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						)-> Tuple:





													'''simple docstring'''






													__UpperCamelCase  	=      FunnelForQuestionAnswering(config=SCREAMING_SNAKE_CASE_      )
													model.to(SCREAMING_SNAKE_CASE_      )
													model.eval()
													__UpperCamelCase  	=      model(
													    SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_  ,						start_positions=SCREAMING_SNAKE_CASE_  ,						end_positions=SCREAMING_SNAKE_CASE_  ,						)
													self.parent.assertEqual(result.start_logits.shape  ,						(self.batch_size, self.seq_length)      )
													self.parent.assertEqual(result.end_logits.shape  ,						(self.batch_size, self.seq_length)      )



							def      A__		(	self      )-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      self.prepare_config_and_inputs()
													(
													    (
													    __UpperCamelCase
													)	,				(
													    __UpperCamelCase
													)	,				(
													    __UpperCamelCase
													)	,				(
													    __UpperCamelCase
													)	,				(
													    __UpperCamelCase
													)	,				(
													    __UpperCamelCase
													)	,				(
													    __UpperCamelCase
													)	,				(
													    __UpperCamelCase
													)	,				
													)  	=      config_and_inputs
													__UpperCamelCase  	=      {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
													return config, inputs_dict






@require_torch
class 							SCREAMING_SNAKE_CASE__     (				SCREAMING_SNAKE_CASE_ ,		SCREAMING_SNAKE_CASE_ ,		unittest.TestCase     ):

							"""simple docstring"""






							_snake_case  		=	False
							_snake_case  		=	False
							_snake_case  		=	(
							    (
							        FunnelModel,
							        FunnelForMaskedLM,
							        FunnelForPreTraining,
							        FunnelForQuestionAnswering,
							        FunnelForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							_snake_case  		=	(
							    {
							        'feature-extraction': (FunnelBaseModel, FunnelModel),
							        'fill-mask': FunnelForMaskedLM,
							        'question-answering': FunnelForQuestionAnswering,
							        'text-classification': FunnelForSequenceClassification,
							        'token-classification': FunnelForTokenClassification,
							        'zero-shot': FunnelForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_=False      )-> str:





													'''simple docstring'''






													__UpperCamelCase  	=      super()._prepare_for_class(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						return_labels=SCREAMING_SNAKE_CASE_      )

													if return_labels:
																			if model_class in get_values(SCREAMING_SNAKE_CASE_      ):
																									__UpperCamelCase  	=      torch.zeros(
																									    (self.model_tester.batch_size, self.model_tester.seq_length)  ,						dtype=torch.long  ,						device=SCREAMING_SNAKE_CASE_      )
													return inputs_dict



							def      A__		(	self      )-> Tuple:





													'''simple docstring'''






													__UpperCamelCase  	=      FunnelModelTester(self      )
													__UpperCamelCase  	=      ConfigTester(self  ,						config_class=SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> str:





													'''simple docstring'''






													self.config_tester.run_common_tests()



							def      A__		(	self      )-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_pretraining(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Optional[int]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> int:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_question_answering(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_      )-> Optional[int]:





													'''simple docstring'''






													if hasattr(SCREAMING_SNAKE_CASE_  ,						'''weight'''      ) and module.weight is not None:
																			module.weight.data.fill_(3      )
													if hasattr(SCREAMING_SNAKE_CASE_  ,						'''bias'''      ) and module.bias is not None:
																			module.bias.data.fill_(3      )

													for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
																			if hasattr(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      ) and getattr(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      ) is not None:
																									__UpperCamelCase  	=      getattr(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )
																									weight.data.fill_(3      )






@require_torch
class 							SCREAMING_SNAKE_CASE__     (				SCREAMING_SNAKE_CASE_ ,		unittest.TestCase     ):

							"""simple docstring"""






							_snake_case  		=	False
							_snake_case  		=	False
							_snake_case  		=	(
							    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
							)



							def      A__		(	self      )-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      FunnelModelTester(self  ,						base=SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      ConfigTester(self  ,						config_class=SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> List[str]:





													'''simple docstring'''






													self.config_tester.run_common_tests()



							def      A__		(	self      )-> Tuple:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_base_model(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Dict:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_sequence_classification(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_multiple_choice(*SCREAMING_SNAKE_CASE_      )



							def      A__		(	self      )-> List[Any]:





													'''simple docstring'''






													__UpperCamelCase	,				__UpperCamelCase  	=      self.model_tester.prepare_config_and_inputs_for_common()
													__UpperCamelCase  	=      True

													for model_class in self.all_model_classes:
																			if model_class.__name__ == "FunnelBaseModel":
																									continue
																			__UpperCamelCase  	=      model_class(SCREAMING_SNAKE_CASE_      )
																			model.to(SCREAMING_SNAKE_CASE_      )
																			model.train()
																			__UpperCamelCase  	=      self._prepare_for_class(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_  ,						return_labels=SCREAMING_SNAKE_CASE_      )
																			__UpperCamelCase  	=      model(**SCREAMING_SNAKE_CASE_      ).loss
																			loss.backward()



							def      A__		(	self  ,						SCREAMING_SNAKE_CASE_      )-> str:





													'''simple docstring'''






													if hasattr(SCREAMING_SNAKE_CASE_  ,						'''weight'''      ) and module.weight is not None:
																			module.weight.data.fill_(3      )
													if hasattr(SCREAMING_SNAKE_CASE_  ,						'''bias'''      ) and module.bias is not None:
																			module.bias.data.fill_(3      )

													for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
																			if hasattr(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      ) and getattr(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      ) is not None:
																									__UpperCamelCase  	=      getattr(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )
																									weight.data.fill_(3      )






@require_torch
@require_sentencepiece
@require_tokenizers
class 							SCREAMING_SNAKE_CASE__     (				unittest.TestCase     ):

							"""simple docstring"""






							def      A__		(	self      )-> Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      13
													__UpperCamelCase  	=      7
													__UpperCamelCase  	=      torch.arange(0  ,						batch_size * sequence_length      ).long().reshape(SCREAMING_SNAKE_CASE_  ,						SCREAMING_SNAKE_CASE_      )
													__UpperCamelCase  	=      [0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
													__UpperCamelCase  	=      torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]      )

													__UpperCamelCase  	=      FunnelModel.from_pretrained('''sgugger/funnel-random-tiny'''      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )[0].abs()

													__UpperCamelCase  	=      torch.tensor(2_3_4_4.8_3_5_2      )
													__UpperCamelCase  	=      torch.tensor(0.8_0_5_2      )
													self.assertTrue(torch.allclose(output.sum()  ,						SCREAMING_SNAKE_CASE_  ,						atol=1E-4      )      )
													self.assertTrue(torch.allclose(output.mean()  ,						SCREAMING_SNAKE_CASE_  ,						atol=1E-4      )      )

													__UpperCamelCase  	=      torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]      )
													__UpperCamelCase  	=      model(SCREAMING_SNAKE_CASE_  ,						attention_mask=SCREAMING_SNAKE_CASE_  ,						token_type_ids=SCREAMING_SNAKE_CASE_      )[0].abs()

													__UpperCamelCase  	=      torch.tensor(2_3_4_3.8_4_2_5      )
													__UpperCamelCase  	=      torch.tensor(0.8_0_4_9      )
													self.assertTrue(torch.allclose(output.sum()  ,						SCREAMING_SNAKE_CASE_  ,						atol=1E-4      )      )
													self.assertTrue(torch.allclose(output.mean()  ,						SCREAMING_SNAKE_CASE_  ,						atol=1E-4      )      )



							@slow
							def      A__		(	self      )-> Union[str, Any]:





													'''simple docstring'''






													__UpperCamelCase  	=      FunnelTokenizer.from_pretrained('''huggingface/funnel-small'''      )
													__UpperCamelCase  	=      FunnelModel.from_pretrained('''huggingface/funnel-small'''      )
													__UpperCamelCase  	=      tokenizer('''Hello! I am the Funnel Transformer model.'''  ,						return_tensors='''pt'''      )
													__UpperCamelCase  	=      model(**SCREAMING_SNAKE_CASE_      )[0]

													__UpperCamelCase  	=      torch.tensor(2_3_5.7_2_4_6      )
													__UpperCamelCase  	=      torch.tensor(0.0_2_5_6      )
													self.assertTrue(torch.allclose(output.sum()  ,						SCREAMING_SNAKE_CASE_  ,						atol=1E-4      )      )
													self.assertTrue(torch.allclose(output.mean()  ,						SCREAMING_SNAKE_CASE_  ,						atol=1E-4      )      )




