"""simple docstring"""







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

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

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


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

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

_a	:	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,
}

_a	:	int=					list("ARNDCQEGHILKMFPSTWYV")

_a	:	Union[str, Any]=					restypes + ["""X"""]
_a	:	str=					restypes_with_x + ["""<pad>""", """<mask>""", """<cls>""", """<sep>""", """<eos>"""]
def     __UpperCAmelCase							(       )				->     Dict:

	'''simple docstring'''





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

	'''simple docstring'''





	__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							(       UpperCAmelCase_       :							str     ,				UpperCAmelCase_       :							str     ,				UpperCAmelCase_       :							bool     ,				UpperCAmelCase_       :							str     ,				UpperCAmelCase_       :							str    )				->     str:

	'''simple docstring'''





	if model.startswith('esmfold'    ):
		__snake_case     :				Dict 						=					MODEL_MAPPING[model]()
	else:
		__snake_case						,							__snake_case     :				Optional[int] 						=					MODEL_MAPPING[model]()
	esm.eval()  # disable dropout

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

	if is_folding_model:
		__snake_case     :				str 						=					esm.esm.alphabet
	__snake_case     :				Union[str, Any] 						=					tuple(alphabet.all_toks    )
	__snake_case     :				Dict 						=					alphabet.mask_idx
	__snake_case     :				List[Any] 						=					alphabet.padding_idx

	if is_folding_model:
		__snake_case     :				Union[str, Any] 						=					esm.esm
	else:
		__snake_case     :				int 						=					esm

	__snake_case     :				str 						=					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=10_26     ,				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:
		__snake_case     :				Dict 						=					esm.classification_heads['mnli'].out_proj.weight.shape[0]
	print('Our ESM config:'     ,				UpperCAmelCase__    )

	if model.startswith('esmfold'    ):
		__snake_case     :				Dict 						=					EsmForProteinFolding
	elif classification_head:
		__snake_case     :				Optional[Any] 						=					EsmForSequenceClassification
	else:
		__snake_case     :				Optional[int] 						=					EsmForMaskedLM
	__snake_case     :				str 						=					model_class(UpperCAmelCase__    )
	model.eval()

	# Now let's copy all the weights.
	# Embeddings
	__snake_case     :				List[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     :				Union[str, Any] 						=					original_esm_model.emb_layer_norm_before.bias

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

	for i in range(config.num_hidden_layers    ):
		# Encoder: start of layer
		__snake_case     :				Optional[int] 						=					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     :				Tuple 						=					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     :				List[Any] 						=					esm_layer.self_attn.q_proj.bias
		__snake_case     :				Any 						=					esm_layer.self_attn.k_proj.weight
		__snake_case     :				Optional[int] 						=					esm_layer.self_attn.k_proj.bias
		__snake_case     :				Optional[int] 						=					esm_layer.self_attn.v_proj.weight
		__snake_case     :				str 						=					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.
			__snake_case     :				Union[str, Any] 						=					esm_layer.self_attn.rot_emb.inv_freq

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

		# self-attention output
		__snake_case     :				Tuple 						=					layer.attention.output
		assert self_output.dense.weight.shape == esm_layer.self_attn.out_proj.weight.shape
		__snake_case     :				List[str] 						=					esm_layer.self_attn.out_proj.weight
		__snake_case     :				List[str] 						=					esm_layer.self_attn.out_proj.bias

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

		# output
		__snake_case     :				Dict 						=					layer.output
		assert bert_output.dense.weight.shape == esm_layer.fca.weight.shape
		__snake_case     :				Any 						=					esm_layer.fca.weight
		__snake_case     :				Any 						=					esm_layer.fca.bias
		# end of layer

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

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

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

		if is_folding_model:
			__snake_case     :				Tuple 						=					torch.max(torch.abs(our_output['positions'] - their_output['positions']    )    ).item()
			__snake_case     :				Union[str, 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     :				int 						=					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
			__snake_case     :				Tuple 						=					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     :				str 						=					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__":
	_a	:	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.")
	_a	:	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"""







from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    get_resize_output_image_size,
    normalize,
    pad,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging
from ...utils.import_utils import is_vision_available


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


if is_vision_available():
						import PIL


class       _lowercase		( __lowerCamelCase	):
				lowerCamelCase__							=							['pixel_values']

				def __init__(  self							:     str							,     _lowercase							:     bool = True							,     _lowercase							:     Dict[str, int] = None							,     _lowercase							:     PILImageResampling = PILImageResampling.BILINEAR							,     _lowercase							:     bool = True							,     _lowercase							:     bool = False							,     _lowercase							:     bool = True							,     _lowercase							:     bool = True							,     _lowercase							:     Union[int, float] = 1 / 2_5_5							,     _lowercase							:     bool = True							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     **_lowercase							:     List[str]							,     ):


										'''simple docstring'''
										super().__init__(**_lowercase  )

										UpperCAmelCase__       						=		size if size is not None else {"""height""": 2_5_6_0, """width""": 1_9_2_0}
										if isinstance(_lowercase							,     (tuple, list)  ):
																# The previous feature extractor size parameter was in (width, height) format
																UpperCAmelCase__       						=		size[::-1]
										UpperCAmelCase__       						=		get_size_dict(_lowercase  )

										UpperCAmelCase__       						=		do_resize
										UpperCAmelCase__       						=		size
										UpperCAmelCase__       						=		resample
										UpperCAmelCase__       						=		do_thumbnail
										UpperCAmelCase__       						=		do_align_long_axis
										UpperCAmelCase__       						=		do_pad
										UpperCAmelCase__       						=		do_rescale
										UpperCAmelCase__       						=		rescale_factor
										UpperCAmelCase__       						=		do_normalize
										UpperCAmelCase__       						=		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
										UpperCAmelCase__       						=		image_std if image_std is not None else IMAGENET_STANDARD_STD

				def 							A			(  self							:     str							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None  ):


										'''simple docstring'''
										UpperCAmelCase__						,   UpperCAmelCase__       						=		get_image_size(_lowercase  )
										UpperCAmelCase__						,   UpperCAmelCase__       						=		size["""height"""], size["""width"""]

										if (output_width < output_height and input_width > input_height) or (
										    output_width > output_height and input_width < input_height
										):
																UpperCAmelCase__       						=		np.rotaa(_lowercase							,     3  )

										if data_format is not None:
																UpperCAmelCase__       						=		to_channel_dimension_format(_lowercase							,     _lowercase  )

										return image

				def 							A			(  self							:     Tuple							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     bool = False							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     ):


										'''simple docstring'''
										UpperCAmelCase__						,   UpperCAmelCase__       						=		size["""height"""], size["""width"""]
										UpperCAmelCase__						,   UpperCAmelCase__       						=		get_image_size(_lowercase  )

										UpperCAmelCase__       						=		output_width - input_width
										UpperCAmelCase__       						=		output_height - input_height

										if random_padding:
																UpperCAmelCase__       						=		np.random.randint(low=0							,     high=delta_height + 1  )
																UpperCAmelCase__       						=		np.random.randint(low=0							,     high=delta_width + 1  )
										else:
																UpperCAmelCase__       						=		delta_height // 2
																UpperCAmelCase__       						=		delta_width // 2

										UpperCAmelCase__       						=		delta_height - pad_top
										UpperCAmelCase__       						=		delta_width - pad_left

										UpperCAmelCase__       						=		((pad_top, pad_bottom), (pad_left, pad_right))
										return pad(_lowercase							,     _lowercase							,     data_format=_lowercase  )

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


										'''simple docstring'''
										logger.info("""pad is deprecated and will be removed in version 4.27. Please use pad_image instead."""  )
										return self.pad_image(*_lowercase							,     **_lowercase  )

				def 							A			(  self							:     Optional[int]							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     PILImageResampling = PILImageResampling.BICUBIC							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Dict							,     ):


										'''simple docstring'''
										UpperCAmelCase__						,   UpperCAmelCase__       						=		get_image_size(_lowercase  )
										UpperCAmelCase__						,   UpperCAmelCase__       						=		size["""height"""], size["""width"""]

										# We always resize to the smallest of either the input or output size.
										UpperCAmelCase__       						=		min(_lowercase							,     _lowercase  )
										UpperCAmelCase__       						=		min(_lowercase							,     _lowercase  )

										if height == input_height and width == input_width:
																return image

										if input_height > input_width:
																UpperCAmelCase__       						=		int(input_width * height / input_height  )
										elif input_width > input_height:
																UpperCAmelCase__       						=		int(input_height * width / input_width  )

										return resize(
										    _lowercase							,     size=(height, width)							,     resample=_lowercase							,     reducing_gap=2.0							,     data_format=_lowercase							,     **_lowercase  )

				def 							A			(  self							:     List[str]							,     _lowercase							:     np.ndarray							,     _lowercase							:     Dict[str, int]							,     _lowercase							:     PILImageResampling = PILImageResampling.BICUBIC							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Optional[Any]							,     ):


										'''simple docstring'''
										UpperCAmelCase__       						=		get_size_dict(_lowercase  )
										UpperCAmelCase__       						=		min(size["""height"""]							,     size["""width"""]  )
										UpperCAmelCase__       						=		get_resize_output_image_size(_lowercase							,     size=_lowercase							,     default_to_square=_lowercase  )
										UpperCAmelCase__       						=		resize(_lowercase							,     size=_lowercase							,     resample=_lowercase							,     data_format=_lowercase							,     **_lowercase  )
										return resized_image

				def 							A			(  self							:     int							,     _lowercase							:     np.ndarray							,     _lowercase							:     Union[int, float]							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Union[str, Any]							,     ):


										'''simple docstring'''
										return rescale(_lowercase							,     scale=_lowercase							,     data_format=_lowercase							,     **_lowercase  )

				def 							A			(  self							:     Optional[int]							,     _lowercase							:     np.ndarray							,     _lowercase							:     Union[float, List[float]]							,     _lowercase							:     Union[float, List[float]]							,     _lowercase							:     Optional[Union[str, ChannelDimension]] = None							,     **_lowercase							:     Tuple							,     ):


										'''simple docstring'''
										return normalize(_lowercase							,     mean=_lowercase							,     std=_lowercase							,     data_format=_lowercase							,     **_lowercase  )

				def 							A			(  self							:     List[str]							,     _lowercase							:     ImageInput							,     _lowercase							:     bool = None							,     _lowercase							:     Dict[str, int] = None							,     _lowercase							:     PILImageResampling = None							,     _lowercase							:     bool = None							,     _lowercase							:     bool = None							,     _lowercase							:     bool = None							,     _lowercase							:     bool = False							,     _lowercase							:     bool = None							,     _lowercase							:     float = None							,     _lowercase							:     bool = None							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     _lowercase							:     Optional[Union[float, List[float]]] = None							,     _lowercase							:     Optional[Union[str, TensorType]] = None							,     _lowercase							:     Optional[ChannelDimension] = ChannelDimension.FIRST							,     **_lowercase							:     Dict							,     ):


										'''simple docstring'''
										UpperCAmelCase__       						=		do_resize if do_resize is not None else self.do_resize
										UpperCAmelCase__       						=		size if size is not None else self.size
										if isinstance(_lowercase							,     (tuple, list)  ):
																# Previous feature extractor had size in (width, height) format
																UpperCAmelCase__       						=		size[::-1]
										UpperCAmelCase__       						=		get_size_dict(_lowercase  )
										UpperCAmelCase__       						=		resample if resample is not None else self.resample
										UpperCAmelCase__       						=		do_thumbnail if do_thumbnail is not None else self.do_thumbnail
										UpperCAmelCase__       						=		do_align_long_axis if do_align_long_axis is not None else self.do_align_long_axis
										UpperCAmelCase__       						=		do_pad if do_pad is not None else self.do_pad
										UpperCAmelCase__       						=		do_rescale if do_rescale is not None else self.do_rescale
										UpperCAmelCase__       						=		rescale_factor if rescale_factor is not None else self.rescale_factor
										UpperCAmelCase__       						=		do_normalize if do_normalize is not None else self.do_normalize
										UpperCAmelCase__       						=		image_mean if image_mean is not None else self.image_mean
										UpperCAmelCase__       						=		image_std if image_std is not None else self.image_std

										UpperCAmelCase__       						=		make_list_of_images(_lowercase  )

										if not valid_images(_lowercase  ):
																raise ValueError(
																    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
																    """torch.Tensor, tf.Tensor or jax.ndarray."""  )

										if do_resize and size is None:
																raise ValueError("""Size must be specified if do_resize is True."""  )

										if do_rescale and rescale_factor is None:
																raise ValueError("""Rescale factor must be specified if do_rescale is True."""  )

										if do_pad and size is None:
																raise ValueError("""Size must be specified if do_pad is True."""  )

										if do_normalize and (image_mean is None or image_std is None):
																raise ValueError("""Image mean and std must be specified if do_normalize is True."""  )

										# All transformations expect numpy arrays.
										UpperCAmelCase__       						=		[to_numpy_array(_lowercase  ) for image in images]

										if do_align_long_axis:
																UpperCAmelCase__       						=		[self.align_long_axis(_lowercase							,     size=_lowercase  ) for image in images]

										if do_resize:
																UpperCAmelCase__       						=		[self.resize(image=_lowercase							,     size=_lowercase							,     resample=_lowercase  ) for image in images]

										if do_thumbnail:
																UpperCAmelCase__       						=		[self.thumbnail(image=_lowercase							,     size=_lowercase  ) for image in images]

										if do_pad:
																UpperCAmelCase__       						=		[self.pad_image(image=_lowercase							,     size=_lowercase							,     random_padding=_lowercase  ) for image in images]

										if do_rescale:
																UpperCAmelCase__       						=		[self.rescale(image=_lowercase							,     scale=_lowercase  ) for image in images]

										if do_normalize:
																UpperCAmelCase__       						=		[self.normalize(image=_lowercase							,     mean=_lowercase							,     std=_lowercase  ) for image in images]

										UpperCAmelCase__       						=		[to_channel_dimension_format(_lowercase							,     _lowercase  ) for image in images]

										UpperCAmelCase__       						=		{"""pixel_values""": images}
										return BatchFeature(data=_lowercase							,     tensor_type=_lowercase  )

