

"""simple docstring"""





import argparse

import torch
from datasets import load_dataset
from donut import DonutModel

from transformers import (
    DonutImageProcessor,
    DonutProcessor,
    DonutSwinConfig,
    DonutSwinModel,
    MBartConfig,
    MBartForCausalLM,
    VisionEncoderDecoderModel,
    XLMRobertaTokenizerFast,
)

def     __SCREAMING_SNAKE_CASE    ( A_					):
	lowerCAmelCase__      :		int    							=      model.config

	lowerCAmelCase__      :		Optional[Any]    							=      DonutSwinConfig(
	    image_size=original_config.input_size		,	patch_size=4		,	depths=original_config.encoder_layer		,	num_heads=[4, 8, 16, 32]		,	window_size=original_config.window_size		,	embed_dim=1_28		,	)
	lowerCAmelCase__      :		Dict    							=      MBartConfig(
	    is_decoder=__UpperCAmelCase		,	is_encoder_decoder=__UpperCAmelCase		,	add_cross_attention=__UpperCAmelCase		,	decoder_layers=original_config.decoder_layer		,	max_position_embeddings=original_config.max_position_embeddings		,	vocab_size=len(
	        model.decoder.tokenizer					)		,	scale_embedding=__UpperCAmelCase		,	add_final_layer_norm=__UpperCAmelCase		,	)

	return encoder_config, decoder_config

def     __SCREAMING_SNAKE_CASE    ( A_					):
	if "encoder.model" in name:
		lowerCAmelCase__      :		List[Any]    							=      name.replace('''encoder.model'''		,	'''encoder'''					)
	if "decoder.model" in name:
		lowerCAmelCase__      :		List[Any]    							=      name.replace('''decoder.model'''		,	'''decoder'''					)
	if "patch_embed.proj" in name:
		lowerCAmelCase__      :		List[Any]    							=      name.replace('''patch_embed.proj'''		,	'''embeddings.patch_embeddings.projection'''					)
	if "patch_embed.norm" in name:
		lowerCAmelCase__      :		Dict    							=      name.replace('''patch_embed.norm'''		,	'''embeddings.norm'''					)
	if name.startswith('''encoder'''					):
		if "layers" in name:
			lowerCAmelCase__      :		Union[str, Any]    							=      """encoder.""" + name
		if "attn.proj" in name:
			lowerCAmelCase__      :		str    							=      name.replace('''attn.proj'''		,	'''attention.output.dense'''					)
		if "attn" in name and "mask" not in name:
			lowerCAmelCase__      :		Any    							=      name.replace('''attn'''		,	'''attention.self'''					)
		if "norm1" in name:
			lowerCAmelCase__      :		List[Any]    							=      name.replace('''norm1'''		,	'''layernorm_before'''					)
		if "norm2" in name:
			lowerCAmelCase__      :		Union[str, Any]    							=      name.replace('''norm2'''		,	'''layernorm_after'''					)
		if "mlp.fc1" in name:
			lowerCAmelCase__      :		List[Any]    							=      name.replace('''mlp.fc1'''		,	'''intermediate.dense'''					)
		if "mlp.fc2" in name:
			lowerCAmelCase__      :		Any    							=      name.replace('''mlp.fc2'''		,	'''output.dense'''					)

		if name == "encoder.norm.weight":
			lowerCAmelCase__      :		int    							=      """encoder.layernorm.weight"""
		if name == "encoder.norm.bias":
			lowerCAmelCase__      :		List[Any]    							=      """encoder.layernorm.bias"""

	return name

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_					):
	for key in orig_state_dict.copy().keys():
		lowerCAmelCase__      :		Optional[Any]    							=      orig_state_dict.pop(__UpperCAmelCase					)

		if "qkv" in key:
			lowerCAmelCase__      :		Any    							=      key.split('''.'''					)
			lowerCAmelCase__      :		Optional[Any]    							=      int(key_split[3]					)
			lowerCAmelCase__      :		int    							=      int(key_split[5]					)
			lowerCAmelCase__      :		List[Any]    							=      model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

			if "weight" in key:
				lowerCAmelCase__      :		int    							=      val[:dim, :]
				lowerCAmelCase__      :		int    							=      val[dim : dim * 2, :]
				lowerCAmelCase__      :		int    							=      val[-dim:, :]
			else:
				lowerCAmelCase__      :		List[Any]    							=      val[:dim]
				lowerCAmelCase__      :		Tuple    							=      val[dim : dim * 2]
				lowerCAmelCase__      :		List[str]    							=      val[-dim:]
		elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]:
			# HuggingFace implementation doesn't use attn_mask buffer
			# and model doesn't use final LayerNorms for the encoder
			pass
		else:
			lowerCAmelCase__      :		Tuple    							=      val

	return orig_state_dict

def     __SCREAMING_SNAKE_CASE    ( A_		,	A_=None		,	A_=False					):
	# load original model
	lowerCAmelCase__      :		Union[str, Any]    							=      DonutModel.from_pretrained(__UpperCAmelCase					).eval()

	# load HuggingFace model
	lowerCAmelCase__      :		Union[str, Any]    							=      get_configs(__UpperCAmelCase					)
	lowerCAmelCase__      :		List[str]    							=      DonutSwinModel(__UpperCAmelCase					)
	lowerCAmelCase__      :		Union[str, Any]    							=      MBartForCausalLM(__UpperCAmelCase					)
	lowerCAmelCase__      :		Dict    							=      VisionEncoderDecoderModel(encoder=__UpperCAmelCase		,	decoder=__UpperCAmelCase					)
	model.eval()

	lowerCAmelCase__      :		Union[str, Any]    							=      original_model.state_dict()
	lowerCAmelCase__      :		int    							=      convert_state_dict(__UpperCAmelCase		,	__UpperCAmelCase					)
	model.load_state_dict(__UpperCAmelCase					)

	# verify results on scanned document
	lowerCAmelCase__      :		Optional[Any]    							=      load_dataset('''hf-internal-testing/example-documents'''					)
	lowerCAmelCase__      :		Union[str, Any]    							=      dataset["""test"""][0]["""image"""].convert('''RGB'''					)

	lowerCAmelCase__      :		int    							=      XLMRobertaTokenizerFast.from_pretrained(__UpperCAmelCase		,	from_slow=__UpperCAmelCase					)
	lowerCAmelCase__      :		str    							=      DonutImageProcessor(
	    do_align_long_axis=original_model.config.align_long_axis		,	size=original_model.config.input_size[::-1]					)
	lowerCAmelCase__      :		Any    							=      DonutProcessor(__UpperCAmelCase		,	__UpperCAmelCase					)
	lowerCAmelCase__      :		Union[str, Any]    							=      processor(__UpperCAmelCase		,	return_tensors='''pt'''					).pixel_values

	if model_name == "naver-clova-ix/donut-base-finetuned-docvqa":
		lowerCAmelCase__      :		Any    							=      """<s_docvqa><s_question>{user_input}</s_question><s_answer>"""
		lowerCAmelCase__      :		Optional[Any]    							=      """When is the coffee break?"""
		lowerCAmelCase__      :		Optional[Any]    							=      task_prompt.replace('''{user_input}'''		,	__UpperCAmelCase					)
	elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip":
		lowerCAmelCase__      :		List[str]    							=      """<s_rvlcdip>"""
	elif model_name in [
	    "naver-clova-ix/donut-base-finetuned-cord-v1",
	    "naver-clova-ix/donut-base-finetuned-cord-v1-2560",
	]:
		lowerCAmelCase__      :		str    							=      """<s_cord>"""
	elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2":
		lowerCAmelCase__      :		List[Any]    							=      """s_cord-v2>"""
	elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket":
		lowerCAmelCase__      :		Optional[Any]    							=      """<s_zhtrainticket>"""
	elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]:
		# use a random prompt
		lowerCAmelCase__      :		List[Any]    							=      """hello world"""
	else:
		raise ValueError('''Model name not supported'''					)
	lowerCAmelCase__      :		Optional[int]    							=      original_model.decoder.tokenizer(__UpperCAmelCase		,	add_special_tokens=__UpperCAmelCase		,	return_tensors='''pt'''					)[
	    """input_ids"""
	]

	lowerCAmelCase__      :		Optional[int]    							=      original_model.encoder.model.patch_embed(__UpperCAmelCase					)
	lowerCAmelCase__      :		List[Any]    							=      model.encoder.embeddings(__UpperCAmelCase					)
	assert torch.allclose(__UpperCAmelCase		,	__UpperCAmelCase		,	atol=1e-3					)

	# verify encoder hidden states
	lowerCAmelCase__      :		Any    							=      original_model.encoder(__UpperCAmelCase					)
	lowerCAmelCase__      :		Optional[Any]    							=      model.encoder(__UpperCAmelCase					).last_hidden_state
	assert torch.allclose(__UpperCAmelCase		,	__UpperCAmelCase		,	atol=1e-2					)

	# verify decoder hidden states
	lowerCAmelCase__      :		str    							=      original_model(__UpperCAmelCase		,	__UpperCAmelCase		,	__UpperCAmelCase					).logits
	lowerCAmelCase__      :		List[str]    							=      model(__UpperCAmelCase		,	decoder_input_ids=__UpperCAmelCase					).logits
	assert torch.allclose(__UpperCAmelCase		,	__UpperCAmelCase		,	atol=1e-3					)
	print('''Looks ok!'''					)

	if pytorch_dump_folder_path is not None:
		print(f'Saving model and processor to {pytorch_dump_folder_path}'					)
		model.save_pretrained(__UpperCAmelCase					)
		processor.save_pretrained(__UpperCAmelCase					)

	if push_to_hub:
		model.push_to_hub('''nielsr/''' + model_name.split('''/'''					)[-1]		,	commit_message='''Update model'''					)
		processor.push_to_hub('''nielsr/''' + model_name.split('''/'''					)[-1]		,	commit_message='''Update model'''					)


if __name__ == "__main__":
							__UpperCamelCase    : str				 =						argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '''--model_name''',
							    default='''naver-clova-ix/donut-base-finetuned-docvqa''',
							    required=False,
							    type=str,
							    help='''Name of the original model you\'d like to convert.''',
							)
							parser.add_argument(
							    '''--pytorch_dump_folder_path''',
							    default=None,
							    required=False,
							    type=str,
							    help='''Path to the output PyTorch model directory.''',
							)
							parser.add_argument(
							    '''--push_to_hub''',
							    action='''store_true''',
							    help='''Whether or not to push the converted model and processor to the 🤗 hub.''',
							)

							__UpperCamelCase    : List[str]				 =						parser.parse_args()
							convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












"""simple docstring"""





import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

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


if is_torch_available():
							import torch

							from transformers.models.esm.modeling_esmfold import EsmForProteinFolding



class       __A		:

					def __init__(			self  :      str     ,		__snake_case  :      int     ,		__snake_case  :      Dict=1_3     ,		__snake_case  :      Dict=7     ,		__snake_case  :      str=False     ,		__snake_case  :      Dict=True     ,		__snake_case  :      Tuple=False     ,		__snake_case  :      List[Any]=False     ,		__snake_case  :      Dict=1_9     ,		__snake_case  :      Dict=3_2     ,		__snake_case  :      Union[str, Any]=5     ,		__snake_case  :      List[Any]=4     ,		__snake_case  :      List[Any]=3_7     ,		__snake_case  :      Union[str, Any]="gelu"     ,		__snake_case  :      List[str]=0.1     ,		__snake_case  :      Union[str, Any]=0.1     ,		__snake_case  :      List[str]=5_1_2     ,		__snake_case  :      List[str]=1_6     ,		__snake_case  :      Tuple=2     ,		__snake_case  :      Optional[Any]=0.02     ,		__snake_case  :      str=3     ,		__snake_case  :      List[str]=4     ,		__snake_case  :      Optional[int]=None     ,		)			->							Optional[int]:
											__magic_name__:	Any							=       parent
											__magic_name__:	Optional[Any]							=       batch_size
											__magic_name__:	Optional[Any]							=       seq_length
											__magic_name__:	Optional[Any]							=       is_training
											__magic_name__:	str							=       use_input_mask
											__magic_name__:	Optional[Any]							=       use_token_type_ids
											__magic_name__:	List[str]							=       use_labels
											__magic_name__:	List[str]							=       vocab_size
											__magic_name__:	Union[str, Any]							=       hidden_size
											__magic_name__:	List[Any]							=       num_hidden_layers
											__magic_name__:	Dict							=       num_attention_heads
											__magic_name__:	int							=       intermediate_size
											__magic_name__:	List[Any]							=       hidden_act
											__magic_name__:	int							=       hidden_dropout_prob
											__magic_name__:	List[str]							=       attention_probs_dropout_prob
											__magic_name__:	Tuple							=       max_position_embeddings
											__magic_name__:	Any							=       type_vocab_size
											__magic_name__:	Union[str, Any]							=       type_sequence_label_size
											__magic_name__:	List[str]							=       initializer_range
											__magic_name__:	Optional[Any]							=       num_labels
											__magic_name__:	List[str]							=       num_choices
											__magic_name__:	Tuple							=       scope

					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Any:
											__magic_name__:	Union[str, Any]							=       ids_tensor([self.batch_size, self.seq_length]     ,		self.vocab_size					)

											__magic_name__:	Dict							=       None
											if self.use_input_mask:
																	__magic_name__:	Optional[int]							=       random_attention_mask([self.batch_size, self.seq_length]					)

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

											__magic_name__:	Dict							=       self.get_config()

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

					def   lowerCamelCase__   (			self  :      int					)			->							Any:
											__magic_name__:	List[Any]							=       EsmConfig(
											    vocab_size=3_3     ,		hidden_size=self.hidden_size     ,		pad_token_id=1     ,		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     ,		is_folding_model=__snake_case     ,		esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}     ,		)
											return config

					def   lowerCamelCase__   (			self  :      Optional[int]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      Any     ,		__snake_case  :      List[Any]     ,		__snake_case  :      Any     ,		__snake_case  :      Union[str, Any]					)			->							Optional[Any]:
											__magic_name__:	str							=       EsmForProteinFolding(config=__snake_case					).float()
											model.to(__snake_case					)
											model.eval()
											__magic_name__:	str							=       model(__snake_case     ,		attention_mask=__snake_case					)
											__magic_name__:	Any							=       model(__snake_case					)
											__magic_name__:	Dict							=       model(__snake_case					)

											self.parent.assertEqual(result.positions.shape     ,		(8, self.batch_size, self.seq_length, 1_4, 3)					)
											self.parent.assertEqual(result.angles.shape     ,		(8, self.batch_size, self.seq_length, 7, 2)					)






					def   lowerCamelCase__   (			self  :      Any					)			->							int:
											__magic_name__:	str							=       self.prepare_config_and_inputs()
											(
											    (
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			
											):	Union[str, Any]							=       config_and_inputs
											__magic_name__:	Any							=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
											return config, inputs_dict



@require_torch
class       __A		(   SCREAMING_SNAKE_CASE_     ,SCREAMING_SNAKE_CASE_     ,unittest.TestCase			):
					UpperCAmelCase__								=				False

					UpperCAmelCase__								=				(EsmForProteinFolding,) if is_torch_available() else ()
					UpperCAmelCase__								=				()
					UpperCAmelCase__								=				{} if is_torch_available() else {}
					UpperCAmelCase__								=				False

					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Optional[Any]:
											__magic_name__:	Optional[int]							=       EsmFoldModelTester(self					)
											__magic_name__:	Optional[Any]							=       ConfigTester(self     ,		config_class=__snake_case     ,		hidden_size=3_7					)

					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Any:
											self.config_tester.run_common_tests()

					def   lowerCamelCase__   (			self  :      List[str]					)			->							str:
											__magic_name__:	Tuple							=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__snake_case					)

					@unittest.skip("""Does not support attention outputs"""					)
					def   lowerCamelCase__   (			self  :      List[str]					)			->							str:
											pass

					@unittest.skip
					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Tuple:
											pass

					@unittest.skip("""Esm does not support embedding resizing"""					)
					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Optional[int]:
											pass

					@unittest.skip("""Esm does not support embedding resizing"""					)
					def   lowerCamelCase__   (			self  :      str					)			->							str:
											pass

					@unittest.skip("""ESMFold does not support passing input embeds!"""					)
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							List[Any]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Any					)			->							List[str]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							List[Any]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      int					)			->							Tuple:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Optional[Any]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Optional[Any]					)			->							Optional[Any]:
											pass

					@unittest.skip("""ESMFold does not output hidden states in the normal way."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							List[Any]:
											pass

					@unittest.skip("""ESMfold does not output hidden states in the normal way."""					)
					def   lowerCamelCase__   (			self  :      Any					)			->							int:
											pass

					@unittest.skip("""ESMFold only has one output format."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							str:
											pass

					@unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality"""					)
					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Dict:
											pass

					@unittest.skip("""ESMFold does not support input chunking."""					)
					def   lowerCamelCase__   (			self  :      int					)			->							Any:
											pass

					@unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."""					)
					def   lowerCamelCase__   (			self  :      int					)			->							Union[str, Any]:
											pass

					@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
					def   lowerCamelCase__   (			self  :      Any					)			->							Optional[Any]:
											pass

					@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							int:
											pass

					@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							List[str]:
											pass

					@unittest.skip("""ESMFold doesn't support data parallel."""					)
					def   lowerCamelCase__   (			self  :      List[str]					)			->							str:
											pass






					@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							Dict:
											pass



@require_torch
class       __A		(   SCREAMING_SNAKE_CASE_			):

					@slow
					def   lowerCamelCase__   (			self  :      Optional[Any]					)			->							List[Any]:
											__magic_name__:	Any							=       EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1"""					).float()
											model.eval()
											__magic_name__:	Optional[int]							=       torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]					)
											__magic_name__:	Dict							=       model(__snake_case					)["""positions"""]
											__magic_name__:	Tuple							=       torch.tensor([2.5828, 0.7993, -10.9334]     ,		dtype=torch.floataa					)
											self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]     ,		__snake_case     ,		atol=1E-4					)					)



