




"""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 					a			(    __snake_case  :					str	):



				'''simple docstring'''





				UpperCAmelCase_       :List[Any]         =   model.config

				UpperCAmelCase_       :str         =   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=128,     )
				UpperCAmelCase_       :Optional[int]         =   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 					a			(    __snake_case  :					List[str]	):



				'''simple docstring'''





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

								if name == "encoder.norm.weight":
												UpperCAmelCase_       :int         =   '''encoder.layernorm.weight'''
								if name == "encoder.norm.bias":
												UpperCAmelCase_       :List[str]         =   '''encoder.layernorm.bias'''

				return name



def 					a			(    __snake_case  :					str,     __snake_case  :					Optional[int]	):



				'''simple docstring'''





				for key in orig_state_dict.copy().keys():
								UpperCAmelCase_       :Tuple         =   orig_state_dict.pop(UpperCamelCase__	)

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

												if "weight" in key:
																UpperCAmelCase_       :Union[str, Any]         =   val[:dim, :]
																UpperCAmelCase_       :Optional[Any]         =   val[dim : dim * 2, :]
																UpperCAmelCase_       :Any         =   val[-dim:, :]
												else:
																UpperCAmelCase_       :Optional[Any]         =   val[:dim]
																UpperCAmelCase_       :Union[str, Any]         =   val[dim : dim * 2]
																UpperCAmelCase_       :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:
												UpperCAmelCase_       :Union[str, Any]         =   val

				return orig_state_dict



def 					a			(    __snake_case  :					Union[str, Any],     __snake_case  :					int=None,     __snake_case  :					str=False	):



				'''simple docstring'''





				UpperCAmelCase_       :Tuple         =   DonutModel.from_pretrained(UpperCamelCase__	).eval()

				# load HuggingFace model
				UpperCAmelCase_      ,UpperCAmelCase_       :Union[str, Any]         =   get_configs(UpperCamelCase__	)
				UpperCAmelCase_       :Dict         =   DonutSwinModel(UpperCamelCase__	)
				UpperCAmelCase_       :Any         =   MBartForCausalLM(UpperCamelCase__	)
				UpperCAmelCase_       :Optional[Any]         =   VisionEncoderDecoderModel(encoder=UpperCamelCase__,     decoder=UpperCamelCase__	)
				model.eval()

				UpperCAmelCase_       :int         =   original_model.state_dict()
				UpperCAmelCase_       :List[Any]         =   convert_state_dict(UpperCamelCase__,     UpperCamelCase__	)
				model.load_state_dict(UpperCamelCase__	)

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

				UpperCAmelCase_       :Optional[int]         =   XLMRobertaTokenizerFast.from_pretrained(UpperCamelCase__,     from_slow=UpperCamelCase__	)
				UpperCAmelCase_       :Optional[Any]         =   DonutImageProcessor(
				    do_align_long_axis=original_model.config.align_long_axis,     size=original_model.config.input_size[::-1]	)
				UpperCAmelCase_       :List[Any]         =   DonutProcessor(UpperCamelCase__,     UpperCamelCase__	)
				UpperCAmelCase_       :int         =   processor(UpperCamelCase__,     return_tensors='''pt'''	).pixel_values

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

				UpperCAmelCase_       :int         =   original_model.encoder.model.patch_embed(UpperCamelCase__	)
				UpperCAmelCase_      ,UpperCAmelCase_       :Any         =   model.encoder.embeddings(UpperCamelCase__	)
				assert torch.allclose(UpperCamelCase__,     UpperCamelCase__,     atol=1E-3	)

				# verify encoder hidden states
				UpperCAmelCase_       :Any         =   original_model.encoder(UpperCamelCase__	)
				UpperCAmelCase_       :Optional[int]         =   model.encoder(UpperCamelCase__	).last_hidden_state
				assert torch.allclose(UpperCamelCase__,     UpperCamelCase__,     atol=1E-2	)

				# verify decoder hidden states
				UpperCAmelCase_       :Tuple         =   original_model(UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__	).logits
				UpperCAmelCase_       :Union[str, Any]         =   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__":
						__lowerCamelCase										=    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.",
						)

						__lowerCamelCase										=    parser.parse_args()
						convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









import warnings
from typing import List

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import is_flax_available, is_tf_available, is_torch_available




class 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       ["image_processor", "tokenizer"]
							lowerCamelCase_							     =       "OwlViTImageProcessor"
							lowerCamelCase_							     =       ("CLIPTokenizer", "CLIPTokenizerFast")


							def __init__(				self     :Optional[Any]				,					__A     :int=None				,					__A     :Optional[int]=None				,					**__A     :str		) -> Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   None
									if "feature_extractor" in kwargs:
											warnings.warn(
											    """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
											    """ instead."""				,					__A				,					)
											SCREAMING_SNAKE_CASE__								=   kwargs.pop("""feature_extractor"""		)

									SCREAMING_SNAKE_CASE__								=   image_processor if image_processor is not None else feature_extractor
									if image_processor is None:
											raise ValueError("""You need to specify an `image_processor`."""		)
									if tokenizer is None:
											raise ValueError("""You need to specify a `tokenizer`."""		)

									super().__init__(__A				,					__A		)


							def __call__(				self     :str				,					__A     :Dict=None				,					__A     :List[str]=None				,					__A     :str=None				,					__A     :Optional[int]="max_length"				,					__A     :Tuple="np"				,					**__A     :int		) -> Tuple:



									"""simple docstring"""



									if text is None and query_images is None and images is None:
											raise ValueError(
											    """You have to specify at least one text or query image or image. All three cannot be none."""		)

									if text is not None:
											if isinstance(__A				,					__A		) or (isinstance(__A				,					__A		) and not isinstance(text[0]				,					__A		)):
													SCREAMING_SNAKE_CASE__								=   [self.tokenizer(__A				,					padding=__A				,					return_tensors=__A				,					**__A		)]

											elif isinstance(__A				,					__A		) and isinstance(text[0]				,					__A		):
													SCREAMING_SNAKE_CASE__								=   []

													# Maximum number of queries across batch
													SCREAMING_SNAKE_CASE__								=   max([len(__A		) for t in text]		)

													# Pad all batch samples to max number of text queries
													for t in text:
															if len(__A		) != max_num_queries:
																	SCREAMING_SNAKE_CASE__								=   t + [""" """] * (max_num_queries - len(__A		))

															SCREAMING_SNAKE_CASE__								=   self.tokenizer(__A				,					padding=__A				,					return_tensors=__A				,					**__A		)
															encodings.append(__A		)
											else:
													raise TypeError("""Input text should be a string, a list of strings or a nested list of strings"""		)

											if return_tensors == "np":
													SCREAMING_SNAKE_CASE__								=   np.concatenate([encoding["""input_ids"""] for encoding in encodings]				,					axis=0		)
													SCREAMING_SNAKE_CASE__								=   np.concatenate([encoding["""attention_mask"""] for encoding in encodings]				,					axis=0		)

											elif return_tensors == "jax" and is_flax_available():
													import jax.numpy as jnp

													SCREAMING_SNAKE_CASE__								=   jnp.concatenate([encoding["""input_ids"""] for encoding in encodings]				,					axis=0		)
													SCREAMING_SNAKE_CASE__								=   jnp.concatenate([encoding["""attention_mask"""] for encoding in encodings]				,					axis=0		)

											elif return_tensors == "pt" and is_torch_available():
													import torch

													SCREAMING_SNAKE_CASE__								=   torch.cat([encoding["""input_ids"""] for encoding in encodings]				,					dim=0		)
													SCREAMING_SNAKE_CASE__								=   torch.cat([encoding["""attention_mask"""] for encoding in encodings]				,					dim=0		)

											elif return_tensors == "tf" and is_tf_available():
													import tensorflow as tf

													SCREAMING_SNAKE_CASE__								=   tf.stack([encoding["""input_ids"""] for encoding in encodings]				,					axis=0		)
													SCREAMING_SNAKE_CASE__								=   tf.stack([encoding["""attention_mask"""] for encoding in encodings]				,					axis=0		)

											else:
													raise ValueError("""Target return tensor type could not be returned"""		)

											SCREAMING_SNAKE_CASE__								=   BatchEncoding()
											SCREAMING_SNAKE_CASE__								=   input_ids
											SCREAMING_SNAKE_CASE__								=   attention_mask

									if query_images is not None:
											SCREAMING_SNAKE_CASE__								=   BatchEncoding()
											SCREAMING_SNAKE_CASE__								=   self.image_processor(
											    __A				,					return_tensors=__A				,					**__A		).pixel_values
											SCREAMING_SNAKE_CASE__								=   query_pixel_values

									if images is not None:
											SCREAMING_SNAKE_CASE__								=   self.image_processor(__A				,					return_tensors=__A				,					**__A		)

									if text is not None and images is not None:
											SCREAMING_SNAKE_CASE__								=   image_features.pixel_values
											return encoding
									elif query_images is not None and images is not None:
											SCREAMING_SNAKE_CASE__								=   image_features.pixel_values
											return encoding
									elif text is not None or query_images is not None:
											return encoding
									else:
											return BatchEncoding(data=dict(**__A		)				,					tensor_type=__A		)


							def   _snake_case	(				self     :List[Any]				,					*__A     :Dict				,					**__A     :Dict		) -> Optional[int]:



									"""simple docstring"""



									return self.image_processor.post_process(*__A				,					**__A		)


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



									"""simple docstring"""



									return self.image_processor.post_process_object_detection(*__A				,					**__A		)


							def   _snake_case	(				self     :str				,					*__A     :List[str]				,					**__A     :Union[str, Any]		) -> Any:



									"""simple docstring"""



									return self.image_processor.post_process_image_guided_detection(*__A				,					**__A		)


							def   _snake_case	(				self     :Dict				,					*__A     :List[str]				,					**__A     :List[str]		) -> int:



									"""simple docstring"""



									return self.tokenizer.batch_decode(*__A				,					**__A		)


							def   _snake_case	(				self     :Dict				,					*__A     :Dict				,					**__A     :List[str]		) -> str:



									"""simple docstring"""



									return self.tokenizer.decode(*__A				,					**__A		)


							@property
							def   _snake_case	(				self     :List[Any]		) -> Optional[int]:



									"""simple docstring"""



									warnings.warn(
									    """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""				,					__A				,					)
									return self.image_processor_class


							@property
							def   _snake_case	(				self     :Any		) -> Optional[Any]:



									"""simple docstring"""



									warnings.warn(
									    """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""				,					__A				,					)
									return self.image_processor