


"""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__      (    SCREAMING_SNAKE_CASE   :   Tuple    ):


  '''simple docstring'''







  lowerCAmelCase					:					Tuple								=							model.config

  lowerCAmelCase					:					Optional[Any]								=							DonutSwinConfig(
      image_size=original_config.input_size    ,				patch_size=4    ,				depths=original_config.encoder_layer    ,				num_heads=[4, 8, 1_6, 3_2]    ,				window_size=original_config.window_size    ,				embed_dim=1_2_8    ,				)
  lowerCAmelCase					:					Optional[int]								=							MBartConfig(
      is_decoder=lowerCamelCase__    ,				is_encoder_decoder=lowerCamelCase__    ,				add_cross_attention=lowerCamelCase__    ,				decoder_layers=original_config.decoder_layer    ,				max_position_embeddings=original_config.max_position_embeddings    ,				vocab_size=len(
          model.decoder.tokenizer    )    ,				scale_embedding=lowerCamelCase__    ,				add_final_layer_norm=lowerCamelCase__    ,				)

  return encoder_config, decoder_config
def  a__      (    SCREAMING_SNAKE_CASE   :   str    ):


  '''simple docstring'''







  if "encoder.model" in name:
    lowerCAmelCase					:					List[Any]								=							name.replace("encoder.model"    ,				"encoder"    )
  if "decoder.model" in name:
    lowerCAmelCase					:					Optional[Any]								=							name.replace("decoder.model"    ,				"decoder"    )
  if "patch_embed.proj" in name:
    lowerCAmelCase					:					Union[str, Any]								=							name.replace("patch_embed.proj"    ,				"embeddings.patch_embeddings.projection"    )
  if "patch_embed.norm" in name:
    lowerCAmelCase					:					Any								=							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					:					List[str]								=							name.replace("attn.proj"    ,				"attention.output.dense"    )
    if "attn" in name and "mask" not in name:
      lowerCAmelCase					:					int								=							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					:					Tuple								=							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					:					Union[str, Any]								=							"encoder.layernorm.weight"
    if name == "encoder.norm.bias":
      lowerCAmelCase					:					Dict								=							"encoder.layernorm.bias"

  return name
def  a__      (    SCREAMING_SNAKE_CASE   :   List[str]    ,				SCREAMING_SNAKE_CASE   :   Optional[Any]    ):


  '''simple docstring'''







  for key in orig_state_dict.copy().keys():
    lowerCAmelCase					:					Optional[int]								=							orig_state_dict.pop(lowerCamelCase__    )

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

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

  return orig_state_dict
def  a__      (    SCREAMING_SNAKE_CASE   :   Union[str, Any]    ,				SCREAMING_SNAKE_CASE   :   int=None    ,				SCREAMING_SNAKE_CASE   :   Dict=False    ):


  '''simple docstring'''







  lowerCAmelCase					:					List[Any]								=							DonutModel.from_pretrained(lowerCamelCase__    ).eval()

  # load HuggingFace model
  lowerCAmelCase						,  lowerCAmelCase					:					Union[str, Any]								=							get_configs(lowerCamelCase__    )
  lowerCAmelCase					:					List[Any]								=							DonutSwinModel(lowerCamelCase__    )
  lowerCAmelCase					:					Optional[Any]								=							MBartForCausalLM(lowerCamelCase__    )
  lowerCAmelCase					:					Optional[Any]								=							VisionEncoderDecoderModel(encoder=lowerCamelCase__    ,				decoder=lowerCamelCase__    )
  model.eval()

  lowerCAmelCase					:					List[str]								=							original_model.state_dict()
  lowerCAmelCase					:					Optional[int]								=							convert_state_dict(lowerCamelCase__    ,				lowerCamelCase__    )
  model.load_state_dict(lowerCamelCase__    )

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

  lowerCAmelCase					:					Any								=							XLMRobertaTokenizerFast.from_pretrained(lowerCamelCase__    ,				from_slow=lowerCamelCase__    )
  lowerCAmelCase					:					Dict								=							DonutImageProcessor(
      do_align_long_axis=original_model.config.align_long_axis    ,				size=original_model.config.input_size[::-1]    )
  lowerCAmelCase					:					Optional[Any]								=							DonutProcessor(lowerCamelCase__    ,				lowerCamelCase__    )
  lowerCAmelCase					:					int								=							processor(lowerCamelCase__    ,				return_tensors="pt"    ).pixel_values

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

  lowerCAmelCase					:					int								=							original_model.encoder.model.patch_embed(lowerCamelCase__    )
  lowerCAmelCase						,  lowerCAmelCase					:					str								=							model.encoder.embeddings(lowerCamelCase__    )
  assert torch.allclose(lowerCamelCase__    ,				lowerCamelCase__    ,				atol=1E-3    )

  # verify encoder hidden states
  lowerCAmelCase					:					Any								=							original_model.encoder(lowerCamelCase__    )
  lowerCAmelCase					:					Any								=							model.encoder(lowerCamelCase__    ).last_hidden_state
  assert torch.allclose(lowerCamelCase__    ,				lowerCamelCase__    ,				atol=1E-2    )

  # verify decoder hidden states
  lowerCAmelCase					:					Optional[Any]								=							original_model(lowerCamelCase__    ,				lowerCamelCase__    ,				lowerCamelCase__    ).logits
  lowerCAmelCase					:					Dict								=							model(lowerCamelCase__    ,				decoder_input_ids=lowerCamelCase__    ).logits
  assert torch.allclose(lowerCamelCase__    ,				lowerCamelCase__    ,				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(lowerCamelCase__    )
    processor.save_pretrained(lowerCamelCase__    )

  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)









"""simple docstring"""







import multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
     import sqlitea

     import sqlalchemy


class 				SCREAMING_SNAKE_CASE__    (		lowercase    ):
       """simple docstring"""


       def __init__(						self		,					snake_case__		,					snake_case__		,					snake_case__ = None		,					snake_case__ = None		,					snake_case__ = False		,					**snake_case__		,					):






         """simple docstring"""





         super().__init__(features=snake_case__		,					cache_dir=snake_case__		,					keep_in_memory=snake_case__		,					**snake_case__						)
         lowerCAmelCase					:					Dict								=							Sql(
             cache_dir=snake_case__		,					features=snake_case__		,					sql=snake_case__		,					con=snake_case__		,					**snake_case__		,					)







       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							None
         lowerCAmelCase					:					Optional[int]								=							None
         lowerCAmelCase					:					str								=							None
         lowerCAmelCase					:					int								=							None

         self.builder.download_and_prepare(
             download_config=snake_case__		,					download_mode=snake_case__		,					verification_mode=snake_case__		,					base_path=snake_case__		,					)

         # Build dataset for splits
         lowerCAmelCase					:					str								=							self.builder.as_dataset(
             split="train"		,					verification_mode=snake_case__		,					in_memory=self.keep_in_memory						)
         return dataset


class 				SCREAMING_SNAKE_CASE__    :
       """simple docstring"""


       def __init__(						self		,					snake_case__		,					snake_case__		,					snake_case__		,					snake_case__ = None		,					snake_case__ = None		,					**snake_case__		,					):






         """simple docstring"""





         if num_proc is not None and num_proc <= 0:
           raise ValueError(f"""num_proc {num_proc} must be an integer > 0."""						)

         lowerCAmelCase					:					Tuple								=							dataset
         lowerCAmelCase					:					Tuple								=							name
         lowerCAmelCase					:					List[str]								=							con
         lowerCAmelCase					:					List[Any]								=							batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
         lowerCAmelCase					:					Optional[int]								=							num_proc
         lowerCAmelCase					:					Optional[int]								=							to_sql_kwargs

       def 	lowercase__							(						self						):






         """simple docstring"""





         lowerCAmelCase					:					List[str]								=							self.to_sql_kwargs.pop("sql"		,					snake_case__						)
         lowerCAmelCase					:					List[Any]								=							self.to_sql_kwargs.pop("con"		,					snake_case__						)
         lowerCAmelCase					:					Dict								=							self.to_sql_kwargs.pop("index"		,					snake_case__						)

         lowerCAmelCase					:					Dict								=							self._write(index=snake_case__		,					**self.to_sql_kwargs						)
         return written

       def 	lowercase__							(						self		,					snake_case__						):






         """simple docstring"""





         lowerCAmelCase						,  lowerCAmelCase						,  lowerCAmelCase					:					Optional[int]								=							args
         lowerCAmelCase					:					Optional[int]								=							{**to_sql_kwargs, "if_exists": "append"} if offset > 0 else to_sql_kwargs
         lowerCAmelCase					:					List[Any]								=							query_table(
             table=self.dataset.data		,					key=slice(snake_case__		,					offset + self.batch_size						)		,					indices=self.dataset._indices		,					)
         lowerCAmelCase					:					Tuple								=							batch.to_pandas()
         lowerCAmelCase					:					str								=							df.to_sql(self.name		,					self.con		,					index=snake_case__		,					**snake_case__						)
         return num_rows or len(snake_case__						)







       def 	lowercase__							(						self		,					snake_case__		,					**snake_case__						):






         """simple docstring"""





         lowerCAmelCase					:					Tuple								=							0

         if self.num_proc is None or self.num_proc == 1:
           for offset in logging.tqdm(
               range(0		,					len(self.dataset						)		,					self.batch_size						)		,					unit="ba"		,					disable=not logging.is_progress_bar_enabled()		,					desc="Creating SQL from Arrow format"		,					):
             written += self._batch_sql((offset, index, to_sql_kwargs)						)
         else:
           lowerCAmelCase						,  lowerCAmelCase					:					Dict								=							len(self.dataset						), self.batch_size
           with multiprocessing.Pool(self.num_proc						) as pool:
             for num_rows in logging.tqdm(
                 pool.imap(
                     self._batch_sql		,					[(offset, index, to_sql_kwargs) for offset in range(0		,					snake_case__		,					snake_case__						)]		,					)		,					total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size		,					unit="ba"		,					disable=not logging.is_progress_bar_enabled()		,					desc="Creating SQL from Arrow format"		,					):
               written += num_rows

         return written





