"""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()
UpperCamelCase_							:							Optional[int]				 =		logging.get_logger(__name__)

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

UpperCamelCase_							:							Dict				 =		{
    """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,
}

UpperCamelCase_							:							int				 =		list("""ARNDCQEGHILKMFPSTWYV""")

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


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


def 	a__   (		UpperCAmelCase__							:      Optional[int]		,      UpperCAmelCase__							:      Optional[int]      )->						Union[str, Any]:
						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 	a__   (		UpperCAmelCase__							:      str		,      UpperCAmelCase__							:      str		,      UpperCAmelCase__							:      bool		,      UpperCAmelCase__							:      str		,      UpperCAmelCase__							:      str      )->						str:
						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(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
												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=UpperCAmelCase__		,      hidden_size=UpperCAmelCase__		,      num_hidden_layers=UpperCAmelCase__		,      num_attention_heads=UpperCAmelCase__		,      intermediate_size=UpperCAmelCase__		,      max_position_embeddings=1_0_2_6		,      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:
												UpperCAmelCase__       						=		esm.classification_heads["""mnli"""].out_proj.weight.shape[0]
						print("""Our ESM config:"""		,      UpperCAmelCase__      )

						if model.startswith("""esmfold"""      ):
												UpperCAmelCase__       						=		EsmForProteinFolding
						elif classification_head:
												UpperCAmelCase__       						=		EsmForSequenceClassification
						else:
												UpperCAmelCase__       						=		EsmForMaskedLM
						UpperCAmelCase__       						=		model_class(UpperCAmelCase__      )
						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"""		,      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.
																		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=UpperCAmelCase__		,      add_special_tokens=UpperCAmelCase__      )
												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(UpperCAmelCase__      )
												# Prepare tokenizer and make sure it matches
												with TemporaryDirectory() as tempdir:
																		UpperCAmelCase__       						=		"""\n""".join(alphabet.all_toks      )
																		UpperCAmelCase__       						=		Path(UpperCAmelCase__      ) / """vocab.txt"""
																		vocab_file.write_text(UpperCAmelCase__      )
																		UpperCAmelCase__       						=		EsmTokenizer(vocab_file=str(UpperCAmelCase__      )      )

												UpperCAmelCase__       						=		hf_tokenizer([row[1] for row in sample_data]		,      return_tensors="""pt"""		,      padding=UpperCAmelCase__      )
												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(**UpperCAmelCase__		,      output_hidden_states=UpperCAmelCase__      )
																		UpperCAmelCase__       						=		our_output["""logits"""]
																		if classification_head:
																								UpperCAmelCase__       						=		esm.model.classification_heads["""mnli"""](esm.extract_features(UpperCAmelCase__      )      )
																		else:
																								UpperCAmelCase__       						=		esm(hf_tokens["""input_ids"""]		,      repr_layers=list(range(9_9_9      )      )      )
																								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(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
																		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(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__":
						UpperCamelCase_							:							Optional[int]				 =		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.""")
						UpperCamelCase_							:							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
						)


"""simple docstring"""







import unittest

from transformers import LayoutLMConfig, is_torch_available
from transformers.testing_utils import 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 (
						    LayoutLMForMaskedLM,
						    LayoutLMForQuestionAnswering,
						    LayoutLMForSequenceClassification,
						    LayoutLMForTokenClassification,
						    LayoutLMModel,
						)


class       _lowercase		:

				def __init__(  self							:     Tuple							,     _lowercase							:     Union[str, Any]							,     _lowercase							:     List[Any]=1_3							,     _lowercase							:     Dict=7							,     _lowercase							:     Tuple=True							,     _lowercase							:     Optional[Any]=True							,     _lowercase							:     Union[str, Any]=True							,     _lowercase							:     Any=True							,     _lowercase							:     Dict=9_9							,     _lowercase							:     Optional[int]=3_2							,     _lowercase							:     Optional[int]=5							,     _lowercase							:     Tuple=4							,     _lowercase							:     Dict=3_7							,     _lowercase							:     Optional[int]="gelu"							,     _lowercase							:     Union[str, Any]=0.1							,     _lowercase							:     Dict=0.1							,     _lowercase							:     Optional[int]=5_1_2							,     _lowercase							:     List[str]=1_6							,     _lowercase							:     List[str]=2							,     _lowercase							:     Optional[Any]=0.02							,     _lowercase							:     List[Any]=3							,     _lowercase							:     Tuple=4							,     _lowercase							:     List[Any]=None							,     _lowercase							:     List[str]=1_0_0_0							,     ):


										'''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__       						=		hidden_size
										UpperCAmelCase__       						=		num_hidden_layers
										UpperCAmelCase__       						=		num_attention_heads
										UpperCAmelCase__       						=		intermediate_size
										UpperCAmelCase__       						=		hidden_act
										UpperCAmelCase__       						=		hidden_dropout_prob
										UpperCAmelCase__       						=		attention_probs_dropout_prob
										UpperCAmelCase__       						=		max_position_embeddings
										UpperCAmelCase__       						=		type_vocab_size
										UpperCAmelCase__       						=		type_sequence_label_size
										UpperCAmelCase__       						=		initializer_range
										UpperCAmelCase__       						=		num_labels
										UpperCAmelCase__       						=		num_choices
										UpperCAmelCase__       						=		scope
										UpperCAmelCase__       						=		range_bbox

				def 							A			(  self							:     Tuple  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		ids_tensor([self.batch_size, self.seq_length]							,     self.vocab_size  )

										UpperCAmelCase__       						=		ids_tensor([self.batch_size, self.seq_length, 4]							,     self.range_bbox  )
										# Ensure that bbox is legal
										for i in range(bbox.shape[0]  ):
																for j in range(bbox.shape[1]  ):
																						if bbox[i, j, 3] < bbox[i, j, 1]:
																												UpperCAmelCase__       						=		bbox[i, j, 3]
																												UpperCAmelCase__       						=		bbox[i, j, 1]
																												UpperCAmelCase__       						=		t
																						if bbox[i, j, 2] < bbox[i, j, 0]:
																												UpperCAmelCase__       						=		bbox[i, j, 2]
																												UpperCAmelCase__       						=		bbox[i, j, 0]
																												UpperCAmelCase__       						=		t

										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__       						=		self.get_config()

										return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels

				def 							A			(  self							:     str  ):


										'''simple docstring'''
										return LayoutLMConfig(
										    vocab_size=self.vocab_size							,     hidden_size=self.hidden_size							,     num_hidden_layers=self.num_hidden_layers							,     num_attention_heads=self.num_attention_heads							,     intermediate_size=self.intermediate_size							,     hidden_act=self.hidden_act							,     hidden_dropout_prob=self.hidden_dropout_prob							,     attention_probs_dropout_prob=self.attention_probs_dropout_prob							,     max_position_embeddings=self.max_position_embeddings							,     type_vocab_size=self.type_vocab_size							,     initializer_range=self.initializer_range							,     )

				def 							A			(  self							:     Dict							,     _lowercase							:     List[Any]							,     _lowercase							:     Union[str, Any]							,     _lowercase							:     int							,     _lowercase							:     List[Any]							,     _lowercase							:     Optional[int]							,     _lowercase							:     Dict							,     _lowercase							:     Optional[int]							,     _lowercase							:     List[str]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		LayoutLMModel(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										UpperCAmelCase__       						=		model(_lowercase							,     _lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase  )
										UpperCAmelCase__       						=		model(_lowercase							,     _lowercase							,     token_type_ids=_lowercase  )
										UpperCAmelCase__       						=		model(_lowercase							,     _lowercase  )
										self.parent.assertEqual(result.last_hidden_state.shape							,     (self.batch_size, self.seq_length, self.hidden_size)  )
										self.parent.assertEqual(result.pooler_output.shape							,     (self.batch_size, self.hidden_size)  )

				def 							A			(  self							:     List[str]							,     _lowercase							:     Optional[int]							,     _lowercase							:     Tuple							,     _lowercase							:     Dict							,     _lowercase							:     int							,     _lowercase							:     Tuple							,     _lowercase							:     int							,     _lowercase							:     Optional[int]							,     _lowercase							:     Optional[int]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		LayoutLMForMaskedLM(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										UpperCAmelCase__       						=		model(_lowercase							,     _lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase							,     labels=_lowercase  )
										self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.seq_length, self.vocab_size)  )

				def 							A			(  self							:     Union[str, Any]							,     _lowercase							:     Optional[Any]							,     _lowercase							:     Union[str, Any]							,     _lowercase							:     Dict							,     _lowercase							:     Tuple							,     _lowercase							:     List[Any]							,     _lowercase							:     str							,     _lowercase							:     Optional[int]							,     _lowercase							:     str  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.num_labels
										UpperCAmelCase__       						=		LayoutLMForSequenceClassification(_lowercase  )
										model.to(_lowercase  )
										model.eval()
										UpperCAmelCase__       						=		model(
										    _lowercase							,     _lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase							,     labels=_lowercase  )
										self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.num_labels)  )

				def 							A			(  self							:     Dict							,     _lowercase							:     Dict							,     _lowercase							:     Any							,     _lowercase							:     Union[str, Any]							,     _lowercase							:     int							,     _lowercase							:     List[str]							,     _lowercase							:     Optional[Any]							,     _lowercase							:     Any							,     _lowercase							:     int  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.num_labels
										UpperCAmelCase__       						=		LayoutLMForTokenClassification(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										UpperCAmelCase__       						=		model(_lowercase							,     _lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase							,     labels=_lowercase  )
										self.parent.assertEqual(result.logits.shape							,     (self.batch_size, self.seq_length, self.num_labels)  )

				def 							A			(  self							:     List[Any]							,     _lowercase							:     Dict							,     _lowercase							:     Union[str, Any]							,     _lowercase							:     Tuple							,     _lowercase							:     Union[str, Any]							,     _lowercase							:     Optional[int]							,     _lowercase							:     Tuple							,     _lowercase							:     Optional[int]							,     _lowercase							:     Union[str, Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		LayoutLMForQuestionAnswering(config=_lowercase  )
										model.to(_lowercase  )
										model.eval()
										UpperCAmelCase__       						=		model(
										    _lowercase							,     bbox=_lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase							,     start_positions=_lowercase							,     end_positions=_lowercase							,     )
										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							:     int  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.prepare_config_and_inputs()
										(
										    (
										    UpperCAmelCase__
										)						,   (
										    UpperCAmelCase__
										)						,   (
										    UpperCAmelCase__
										)						,   (
										    UpperCAmelCase__
										)						,   (
										    UpperCAmelCase__
										)						,   (
										    UpperCAmelCase__
										)						,   (
										    UpperCAmelCase__
										)						,   (
										    UpperCAmelCase__
										)						,   
										)       						=		config_and_inputs
										UpperCAmelCase__       						=		{
										    """input_ids""": input_ids,
										    """bbox""": bbox,
										    """token_type_ids""": token_type_ids,
										    """attention_mask""": input_mask,
										}
										return config, inputs_dict


@require_torch
class       _lowercase		( __lowerCamelCase   ,   __lowerCamelCase   ,   unittest.TestCase	):
				lowerCamelCase__							=							(
				    (
				        LayoutLMModel,
				        LayoutLMForMaskedLM,
				        LayoutLMForSequenceClassification,
				        LayoutLMForTokenClassification,
				        LayoutLMForQuestionAnswering,
				    )
				    if is_torch_available()
				    else None
				)
				lowerCamelCase__							=							(
				    {
				        'document-question-answering': LayoutLMForQuestionAnswering,
				        'feature-extraction': LayoutLMModel,
				        'fill-mask': LayoutLMForMaskedLM,
				        'text-classification': LayoutLMForSequenceClassification,
				        'token-classification': LayoutLMForTokenClassification,
				        'zero-shot': LayoutLMForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				lowerCamelCase__							=							True

				def 							A			(  self							:     Union[str, Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		LayoutLMModelTester(self  )
										UpperCAmelCase__       						=		ConfigTester(self							,     config_class=_lowercase							,     hidden_size=3_7  )

				def 							A			(  self							:     str  ):


										'''simple docstring'''
										self.config_tester.run_common_tests()

				def 							A			(  self							:     Any  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowercase  )

				def 							A			(  self							:     Optional[Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																UpperCAmelCase__       						=		type
																self.model_tester.create_and_check_model(*_lowercase  )

				def 							A			(  self							:     Optional[int]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*_lowercase  )

				def 							A			(  self							:     Optional[Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*_lowercase  )

				def 							A			(  self							:     List[Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*_lowercase  )

				def 							A			(  self							:     Union[str, Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*_lowercase  )

				@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""  )
				def 							A			(  self							:     str  ):


										'''simple docstring'''
										pass


def 	a__   (		)->						Any:
						# Here we prepare a batch of 2 sequences to test a LayoutLM forward pass on:
						# fmt: off
						UpperCAmelCase__       						=		torch.tensor([[1_0_1,1_0_1_9,1_0_1_4,1_0_1_6,1_0_3_7,1_2_8_4_9,4_7_4_7,1_0_0_4,1_4_2_4_6,2_2_7_8,5_4_3_9,4_5_2_4,5_0_0_2,2_9_3_0,2_1_9_3,2_9_3_0,4_3_4_1,3_2_0_8,1_0_0_5,1_0_5_5,2_1_7_1,2_8_4_8,1_1_3_0_0,3_5_3_1,1_0_2],[1_0_1,4_0_7_0,4_0_3_4,7_0_2_0,1_0_2_4,3_0_5_8,1_0_1_5,1_0_1_3,2_8_6_1,1_0_1_3,6_0_7_0,1_9_2_7_4,2_7_7_2,6_2_0_5,2_7_8_1_4,1_6_1_4_7,1_6_1_4_7,4_3_4_3,2_0_4_7,1_0_2_8_3,1_0_9_6_9,1_4_3_8_9,1_0_1_2,2_3_3_8,1_0_2]]		,      device=UpperCAmelCase__      )  # noqa: E231
						UpperCAmelCase__       						=		torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],]		,      device=UpperCAmelCase__      )  # noqa: E231
						UpperCAmelCase__       						=		torch.tensor([[[0,0,0,0],[4_2_3,2_3_7,4_4_0,2_5_1],[4_2_7,2_7_2,4_4_1,2_8_7],[4_1_9,1_1_5,4_3_7,1_2_9],[9_6_1,8_8_5,9_9_2,9_1_2],[2_5_6,3_8,3_3_0,5_8],[2_5_6,3_8,3_3_0,5_8],[3_3_6,4_2,3_5_3,5_7],[3_6_0,3_9,4_0_1,5_6],[3_6_0,3_9,4_0_1,5_6],[4_1_1,3_9,4_7_1,5_9],[4_7_9,4_1,5_2_8,5_9],[5_3_3,3_9,6_3_0,6_0],[6_7,1_1_3,1_3_4,1_3_1],[1_4_1,1_1_5,2_0_9,1_3_2],[6_8,1_4_9,1_3_3,1_6_6],[1_4_1,1_4_9,1_8_7,1_6_4],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[1_9_5,1_4_8,2_8_7,1_6_5],[2_9_5,1_4_8,3_4_9,1_6_5],[4_4_1,1_4_9,4_9_2,1_6_6],[4_9_7,1_4_9,5_4_6,1_6_4],[6_4,2_0_1,1_2_5,2_1_8],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]],[[0,0,0,0],[6_6_2,1_5_0,7_5_4,1_6_6],[6_6_5,1_9_9,7_4_2,2_1_1],[5_1_9,2_1_3,5_5_4,2_2_8],[5_1_9,2_1_3,5_5_4,2_2_8],[1_3_4,4_3_3,1_8_7,4_5_4],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[1_3_0,4_6_7,2_0_4,4_8_0],[3_1_4,4_6_9,3_7_6,4_8_2],[5_0_4,6_8_4,5_8_2,7_0_6],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[9_4_1,8_2_5,9_7_3,9_0_0],[6_1_0,7_4_9,6_5_2,7_6_5],[1_3_0,6_5_9,1_6_8,6_7_2],[1_7_6,6_5_7,2_3_7,6_7_2],[2_3_8,6_5_7,3_1_2,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[4_4_3,6_5_3,6_2_8,6_7_2],[7_1_6,3_0_1,8_2_5,3_1_7],[1_0_0_0,1_0_0_0,1_0_0_0,1_0_0_0]]]		,      device=UpperCAmelCase__      )  # noqa: E231
						UpperCAmelCase__       						=		torch.tensor([[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0],[0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0]]		,      device=UpperCAmelCase__      )  # noqa: E231
						# these are sequence labels (i.e. at the token level)
						UpperCAmelCase__       						=		torch.tensor([[-1_0_0,1_0,1_0,1_0,9,1,-1_0_0,7,7,-1_0_0,7,7,4,2,5,2,8,8,-1_0_0,-1_0_0,5,0,3,2,-1_0_0],[-1_0_0,1_2,1_2,1_2,-1_0_0,1_2,1_0,-1_0_0,-1_0_0,-1_0_0,-1_0_0,1_0,1_2,9,-1_0_0,-1_0_0,-1_0_0,1_0,1_0,1_0,9,1_2,-1_0_0,1_0,-1_0_0]]		,      device=UpperCAmelCase__      )  # noqa: E231
						# fmt: on

						return input_ids, attention_mask, bbox, token_type_ids, labels







@require_torch
class       _lowercase		( unittest.TestCase	):

				@slow
				def 							A			(  self							:     List[Any]  ):


										'''simple docstring'''
										UpperCAmelCase__       						=		LayoutLMModel.from_pretrained("""microsoft/layoutlm-base-uncased"""  ).to(_lowercase  )

										UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__       						=		prepare_layoutlm_batch_inputs()

										# forward pass
										UpperCAmelCase__       						=		model(input_ids=_lowercase							,     bbox=_lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase  )

										# test the sequence output on [0, :3, :3]
										UpperCAmelCase__       						=		torch.tensor(
										    [[0.1785, -0.1947, -0.0425], [-0.3254, -0.2807, 0.2553], [-0.5391, -0.3322, 0.3364]]							,     device=_lowercase							,     )

										self.assertTrue(torch.allclose(outputs.last_hidden_state[0, :3, :3]							,     _lowercase							,     atol=1e-3  )  )

										# test the pooled output on [1, :3]
										UpperCAmelCase__       						=		torch.tensor([-0.6580, -0.0214, 0.8552]							,     device=_lowercase  )

										self.assertTrue(torch.allclose(outputs.pooler_output[1, :3]							,     _lowercase							,     atol=1e-3  )  )

				@slow
				def 							A			(  self							:     Optional[Any]  ):


										'''simple docstring'''
										# initialize model with randomly initialized sequence classification head
										UpperCAmelCase__       						=		LayoutLMForSequenceClassification.from_pretrained("""microsoft/layoutlm-base-uncased"""							,     num_labels=2  ).to(
										    _lowercase  )

										UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__       						=		prepare_layoutlm_batch_inputs()

										# forward pass
										UpperCAmelCase__       						=		model(
										    input_ids=_lowercase							,     bbox=_lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase							,     labels=torch.tensor([1, 1]							,     device=_lowercase  )							,     )

										# test whether we get a loss as a scalar
										UpperCAmelCase__       						=		outputs.loss
										UpperCAmelCase__       						=		torch.Size([]  )
										self.assertEqual(loss.shape							,     _lowercase  )

										# test the shape of the logits
										UpperCAmelCase__       						=		outputs.logits
										UpperCAmelCase__       						=		torch.Size((2, 2)  )
										self.assertEqual(logits.shape							,     _lowercase  )

				@slow
				def 							A			(  self							:     Any  ):


										'''simple docstring'''
										# initialize model with randomly initialized token classification head
										UpperCAmelCase__       						=		LayoutLMForTokenClassification.from_pretrained("""microsoft/layoutlm-base-uncased"""							,     num_labels=1_3  ).to(
										    _lowercase  )

										UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__       						=		prepare_layoutlm_batch_inputs()

										# forward pass
										UpperCAmelCase__       						=		model(
										    input_ids=_lowercase							,     bbox=_lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase							,     labels=_lowercase  )

										# test the loss calculation to be around 2.65
										# expected_loss = torch.tensor(2.65, device=torch_device)

										# The loss is currently somewhat random and can vary between 0.1-0.3 atol.
										# self.assertTrue(torch.allclose(outputs.loss, expected_loss, atol=0.1))

										# test the shape of the logits
										UpperCAmelCase__       						=		outputs.logits
										UpperCAmelCase__       						=		torch.Size((2, 2_5, 1_3)  )
										self.assertEqual(logits.shape							,     _lowercase  )

				@slow
				def 							A			(  self							:     Optional[int]  ):


										'''simple docstring'''
										# initialize model with randomly initialized token classification head
										UpperCAmelCase__       						=		LayoutLMForQuestionAnswering.from_pretrained("""microsoft/layoutlm-base-uncased"""  ).to(_lowercase  )

										UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__						,   UpperCAmelCase__       						=		prepare_layoutlm_batch_inputs()

										# forward pass
										UpperCAmelCase__       						=		model(input_ids=_lowercase							,     bbox=_lowercase							,     attention_mask=_lowercase							,     token_type_ids=_lowercase  )

										# test the shape of the logits
										UpperCAmelCase__       						=		torch.Size((2, 2_5)  )
										self.assertEqual(outputs.start_logits.shape							,     _lowercase  )
										self.assertEqual(outputs.end_logits.shape							,     _lowercase  )

