
'''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()
lowercase__						=			logging.get_logger(__name__)

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

lowercase__						=			{
    '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,
}

lowercase__						=			list("ARNDCQEGHILKMFPSTWYV")

lowercase__						=			restypes + ['X']
lowercase__						=			restypes_with_x + ['<pad>', '<mask>', '<cls>', '<sep>', '<eos>']




def       _UpperCamelCase	(					)     ->       Any:


		'''simple docstring'''


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




def       _UpperCamelCase	(					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__ )     ->       int:


		'''simple docstring'''


		snake_case     :					Optional[Any]     =	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	(					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__   ,					SCREAMING_SNAKE_CASE__ )     ->       List[str]:


		'''simple docstring'''


		if model.startswith('''esmfold''' ):
				snake_case     :					Union[str, Any]     =	MODEL_MAPPING[model]()
		else:
				snake_case     :					Optional[Any]     =	MODEL_MAPPING[model]()
		esm.eval()  # disable dropout

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

		if is_folding_model:
				snake_case     :					str     =	esm.esm.alphabet
		snake_case     :					Any     =	tuple(alphabet.all_toks )
		snake_case     :					Optional[Any]     =	alphabet.mask_idx
		snake_case     :					Dict     =	alphabet.padding_idx

		if is_folding_model:
				snake_case     :					Any     =	esm.esm
		else:
				snake_case     :					Optional[Any]     =	esm

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

		if model.startswith('''esmfold''' ):
				snake_case     :					List[Any]     =	EsmForProteinFolding
		elif classification_head:
				snake_case     :					Tuple     =	EsmForSequenceClassification
		else:
				snake_case     :					Dict     =	EsmForMaskedLM
		snake_case     :					int     =	model_class(_SCREAMING_SNAKE_CASE )
		model.eval()

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

		if config.emb_layer_norm_before:
				snake_case     :					Any     =	original_esm_model.emb_layer_norm_before.weight
				snake_case     :					Dict     =	original_esm_model.emb_layer_norm_before.bias

		snake_case     :					List[Any]     =	original_esm_model.emb_layer_norm_after.weight
		snake_case     :					List[str]     =	original_esm_model.emb_layer_norm_after.bias

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

				# self attention
				snake_case     :					EsmSelfAttention     =	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) )
				)

				snake_case     :					Union[str, Any]     =	esm_layer.self_attn.q_proj.weight
				snake_case     :					Any     =	esm_layer.self_attn.q_proj.bias
				snake_case     :					List[Any]     =	esm_layer.self_attn.k_proj.weight
				snake_case     :					Union[str, Any]     =	esm_layer.self_attn.k_proj.bias
				snake_case     :					Optional[Any]     =	esm_layer.self_attn.v_proj.weight
				snake_case     :					List[str]     =	esm_layer.self_attn.v_proj.bias

				if getattr(esm_layer.self_attn   ,					'''rot_emb'''   ,					_SCREAMING_SNAKE_CASE ) 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.
						snake_case     :					Optional[int]     =	esm_layer.self_attn.rot_emb.inv_freq

				# LayerNorm changes for pre-activation
				snake_case     :					str     =	esm_layer.self_attn_layer_norm.weight
				snake_case     :					Dict     =	esm_layer.self_attn_layer_norm.bias
				snake_case     :					Optional[Any]     =	esm_layer.final_layer_norm.weight
				snake_case     :					List[str]     =	esm_layer.final_layer_norm.bias

				# self-attention output
				snake_case     :					EsmSelfOutput     =	layer.attention.output
				assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
				snake_case     :					int     =	esm_layer.self_attn.out_proj.weight
				snake_case     :					int     =	esm_layer.self_attn.out_proj.bias

				# intermediate
				snake_case     :					EsmIntermediate     =	layer.intermediate
				assert intermediate.dense.weight.shape == esm_layer.fca.weight.shape
				snake_case     :					int     =	esm_layer.fca.weight
				snake_case     :					Tuple     =	esm_layer.fca.bias

				# output
				snake_case     :					EsmOutput     =	layer.output
				assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
				snake_case     :					Tuple     =	esm_layer.fca.weight
				snake_case     :					int     =	esm_layer.fca.bias
				# end of layer

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

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

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

				if is_folding_model:
						snake_case     :					Optional[Any]     =	torch.max(torch.abs(our_output['''positions'''] - their_output['''positions'''] ) ).item()
						snake_case     :					Any     =	torch.allclose(our_output['''positions''']   ,					their_output['''positions''']   ,					atol=1E-5 )
				else:
						snake_case     :					Optional[int]     =	torch.max(torch.abs(our_output - their_output ) ).item()
						snake_case     :					Optional[int]     =	torch.allclose(_SCREAMING_SNAKE_CASE   ,					_SCREAMING_SNAKE_CASE   ,					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
						snake_case     :					List[Any]     =	model.predict_contacts(hf_tokens['''input_ids''']   ,					hf_tokens['''attention_mask'''] )
						snake_case     :					Tuple     =	esm.predict_contacts(hf_tokens['''input_ids'''] )
						snake_case     :					Any     =	torch.max(torch.abs(our_output - their_output ) ).item()
						snake_case     :					Union[str, Any]     =	torch.allclose(_SCREAMING_SNAKE_CASE   ,					_SCREAMING_SNAKE_CASE   ,					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(_SCREAMING_SNAKE_CASE ).mkdir(parents=_SCREAMING_SNAKE_CASE   ,					exist_ok=_SCREAMING_SNAKE_CASE )
				print(F'Saving model to {pytorch_dump_folder_path}' )
				model.save_pretrained(_SCREAMING_SNAKE_CASE )

				del esm  # Free up some memory before continuing

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

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


if __name__ == "__main__":
						lowercase__						=			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.")
						lowercase__						=			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'''
from __future__ import annotations

import unittest

from transformers import CTRLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers.models.ctrl.modeling_tf_ctrl import (
						    TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFCTRLForSequenceClassification,
						    TFCTRLLMHeadModel,
						    TFCTRLModel,
						)






class   snake_case__					(			__SCREAMING_SNAKE_CASE				):







		"""simple docstring"""

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






				"""simple docstring"""







				snake_case     :					List[str]     =	parent
				snake_case     :					str     =	13
				snake_case     :					Any     =	7
				snake_case     :					str     =	True
				snake_case     :					int     =	True
				snake_case     :					Optional[Any]     =	True
				snake_case     :					Optional[Any]     =	True
				snake_case     :					str     =	True
				snake_case     :					Any     =	99
				snake_case     :					int     =	32
				snake_case     :					str     =	2
				snake_case     :					Optional[Any]     =	4
				snake_case     :					str     =	37
				snake_case     :					Optional[int]     =	'''gelu'''
				snake_case     :					List[Any]     =	0.1
				snake_case     :					Dict     =	0.1
				snake_case     :					str     =	512
				snake_case     :					List[str]     =	16
				snake_case     :					Tuple     =	2
				snake_case     :					List[str]     =	0.02
				snake_case     :					Optional[int]     =	3
				snake_case     :					Optional[int]     =	4
				snake_case     :					List[str]     =	None
				snake_case     :					List[str]     =	self.vocab_size - 1



		def 			lowerCAmelCase							(			self							: Dict			)	-> Optional[Any]:






				"""simple docstring"""







				snake_case     :					Union[str, Any]     =	ids_tensor([self.batch_size, self.seq_length]     ,  self.vocab_size			)

				snake_case     :					Optional[Any]     =	None
				if self.use_input_mask:
						snake_case     :					List[str]     =	random_attention_mask([self.batch_size, self.seq_length]			)

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

				snake_case     :					Optional[int]     =	None
				if self.use_mc_token_ids:
						snake_case     :					List[str]     =	ids_tensor([self.batch_size, self.num_choices]     ,  self.seq_length			)

				snake_case     :					Dict     =	None
				snake_case     :					Any     =	None
				snake_case     :					str     =	None
				if self.use_labels:
						snake_case     :					Optional[Any]     =	ids_tensor([self.batch_size]     ,  self.type_sequence_label_size			)
						snake_case     :					Optional[int]     =	ids_tensor([self.batch_size, self.seq_length]     ,  self.num_labels			)
						snake_case     :					List[str]     =	ids_tensor([self.batch_size]     ,  self.num_choices			)

				snake_case     :					List[Any]     =	CTRLConfig(
				    vocab_size=self.vocab_size     ,  n_embd=self.hidden_size     ,  n_layer=self.num_hidden_layers     ,  n_head=self.num_attention_heads     ,  n_positions=self.max_position_embeddings     ,  pad_token_id=self.pad_token_id     ,  )

				snake_case     :					List[Any]     =	ids_tensor([self.num_hidden_layers, self.num_attention_heads]     ,  2			)

				return (
				    config,
				    input_ids,
				    input_mask,
				    head_mask,
				    token_type_ids,
				    mc_token_ids,
				    sequence_labels,
				    token_labels,
				    choice_labels,
				)



		def 			lowerCAmelCase							(			self							: Dict     ,  UpperCamelCase__							: int     ,  UpperCamelCase__							: List[str]     ,  UpperCamelCase__							: Optional[Any]     ,  UpperCamelCase__							: Optional[int]     ,  UpperCamelCase__							: int     ,  *UpperCamelCase__							: str			)	-> List[str]:






				"""simple docstring"""







				snake_case     :					Tuple     =	TFCTRLModel(config=UpperCamelCase__			)
				snake_case     :					Union[str, Any]     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
				snake_case     :					Optional[int]     =	model(UpperCamelCase__			)

				snake_case     :					Optional[Any]     =	[input_ids, None, input_mask]  # None is the input for 'past'
				snake_case     :					str     =	model(UpperCamelCase__			)

				snake_case     :					Optional[Any]     =	model(UpperCamelCase__			)

				self.parent.assertEqual(result.last_hidden_state.shape     ,  (self.batch_size, self.seq_length, self.hidden_size)			)



		def 			lowerCAmelCase							(			self							: Any     ,  UpperCamelCase__							: int     ,  UpperCamelCase__							: Union[str, Any]     ,  UpperCamelCase__							: str     ,  UpperCamelCase__							: Any     ,  UpperCamelCase__							: int     ,  *UpperCamelCase__							: List[str]			)	-> List[Any]:






				"""simple docstring"""







				snake_case     :					Union[str, Any]     =	TFCTRLLMHeadModel(config=UpperCamelCase__			)
				snake_case     :					int     =	{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
				snake_case     :					Any     =	model(UpperCamelCase__			)
				self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.seq_length, self.vocab_size)			)



		def 			lowerCAmelCase							(			self							: str     ,  UpperCamelCase__							: Tuple     ,  UpperCamelCase__							: Tuple     ,  UpperCamelCase__							: Dict     ,  UpperCamelCase__							: str     ,  UpperCamelCase__							: Tuple     ,  *UpperCamelCase__							: Tuple			)	-> Optional[Any]:






				"""simple docstring"""







				snake_case     :					int     =	self.num_labels
				snake_case     :					Optional[int]     =	ids_tensor([self.batch_size]     ,  self.type_sequence_label_size			)
				snake_case     :					Any     =	{
				    '''input_ids''': input_ids,
				    '''token_type_ids''': token_type_ids,
				    '''labels''': sequence_labels,
				}
				snake_case     :					Dict     =	TFCTRLForSequenceClassification(UpperCamelCase__			)
				snake_case     :					int     =	model(UpperCamelCase__			)
				self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.num_labels)			)



		def 			lowerCAmelCase							(			self							: int			)	-> str:






				"""simple docstring"""







				snake_case     :					Union[str, Any]     =	self.prepare_config_and_inputs()

				(
				    (
				    snake_case
				) ,(
				    snake_case
				) ,(
				    snake_case
				) ,(
				    snake_case
				) ,(
				    snake_case
				) ,(
				    snake_case
				) ,(
				    snake_case
				) ,(
				    snake_case
				) ,(
				    snake_case
				) ,
				)     :					str     =	config_and_inputs

				snake_case     :					Any     =	{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
				return config, inputs_dict







@require_tf
class   snake_case__					(			__SCREAMING_SNAKE_CASE						,				__SCREAMING_SNAKE_CASE						,				unittest.TestCase				):







		"""simple docstring"""

		lowerCamelCase           =	(TFCTRLModel, TFCTRLLMHeadModel, TFCTRLForSequenceClassification) if is_tf_available() else ()
		lowerCamelCase           =	(TFCTRLLMHeadModel,) if is_tf_available() else ()
		lowerCamelCase           =	(
		    {
		        """feature-extraction""": TFCTRLModel,
		        """text-classification""": TFCTRLForSequenceClassification,
		        """text-generation""": TFCTRLLMHeadModel,
		        """zero-shot""": TFCTRLForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		lowerCamelCase           =	False
		lowerCamelCase           =	False



		def 			lowerCAmelCase							(			self							: Union[str, Any]     ,  UpperCamelCase__							: Union[str, Any]     ,  UpperCamelCase__							: List[str]     ,  UpperCamelCase__							: Tuple     ,  UpperCamelCase__							: Union[str, Any]     ,  UpperCamelCase__							: Dict			)	-> str:






				"""simple docstring"""







				if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
						# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
						# `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
						# config could not be created.
						return True

				return False



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






				"""simple docstring"""







				snake_case     :					str     =	TFCTRLModelTester(self			)
				snake_case     :					Dict     =	ConfigTester(self     ,  config_class=UpperCamelCase__     ,  n_embd=37			)



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






				"""simple docstring"""







				self.config_tester.run_common_tests()



		def 			lowerCAmelCase							(			self							: Dict			)	-> Tuple:






				"""simple docstring"""







				snake_case     :					Optional[int]     =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_ctrl_model(*UpperCamelCase__			)



		def 			lowerCAmelCase							(			self							: Dict			)	-> List[str]:






				"""simple docstring"""







				snake_case     :					Union[str, Any]     =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_ctrl_lm_head(*UpperCamelCase__			)



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






				"""simple docstring"""







				snake_case     :					Dict     =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_ctrl_for_sequence_classification(*UpperCamelCase__			)



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






				"""simple docstring"""







				snake_case ,snake_case     :					Union[str, Any]     =	self.model_tester.prepare_config_and_inputs_for_common()
				snake_case     :					int     =	[TFCTRLLMHeadModel]
				snake_case     :					Tuple     =	[TFCTRLForSequenceClassification]

				for model_class in self.all_model_classes:
						snake_case     :					List[Any]     =	model_class(UpperCamelCase__			)
						model.build()  # may be needed for the get_bias() call below
						assert isinstance(model.get_input_embeddings()     ,  tf.keras.layers.Layer			)

						if model_class in list_lm_models:
								snake_case     :					Optional[Any]     =	model.get_output_embeddings()
								assert isinstance(UpperCamelCase__     ,  tf.keras.layers.Layer			)
								snake_case     :					Any     =	model.get_bias()
								assert isinstance(UpperCamelCase__     ,  UpperCamelCase__			)
								for k, v in name.items():
										assert isinstance(UpperCamelCase__     ,  tf.Variable			)
						elif model_class in list_other_models_with_output_ebd:
								snake_case     :					List[Any]     =	model.get_output_embeddings()
								assert isinstance(UpperCamelCase__     ,  tf.keras.layers.Layer			)
								snake_case     :					int     =	model.get_bias()
								assert name is None
						else:
								snake_case     :					Tuple     =	model.get_output_embeddings()
								assert x is None
								snake_case     :					Dict     =	model.get_bias()
								assert name is None



		@slow
		def 			lowerCAmelCase							(			self							: Optional[Any]			)	-> Optional[int]:






				"""simple docstring"""







				for model_name in TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						snake_case     :					str     =	TFCTRLModel.from_pretrained(UpperCamelCase__			)
						self.assertIsNotNone(UpperCamelCase__			)







@require_tf
class   snake_case__					(			unittest.TestCase				):







		"""simple docstring"""

		@slow
		def 			lowerCAmelCase							(			self							: Optional[int]			)	-> Union[str, Any]:






				"""simple docstring"""







				snake_case     :					Optional[Any]     =	TFCTRLLMHeadModel.from_pretrained('''ctrl'''			)
				snake_case     :					Optional[Any]     =	tf.convert_to_tensor([[1_1859, 0, 1611, 8]]     ,  dtype=tf.intaa			)  # Legal the president is
				snake_case     :					Dict     =	[
				    1_1859,
				    0,
				    1611,
				    8,
				    5,
				    150,
				    2_6449,
				    2,
				    19,
				    348,
				    469,
				    3,
				    2595,
				    48,
				    2_0740,
				    24_6533,
				    24_6533,
				    19,
				    30,
				    5,
				]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

				snake_case     :					int     =	model.generate(UpperCamelCase__     ,  do_sample=UpperCamelCase__			)
				self.assertListEqual(output_ids[0].numpy().tolist()     ,  UpperCamelCase__			)


