




"""simple docstring"""



import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
      import os

      import jax.numpy as jnp
      from jax import jit

      from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
      from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

      SCREAMING_SNAKE_CASE    =	"0.12"  # assumed parallelism: 8


if is_torch_available():
      import torch







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_      ,			lowercase_=None						)						->		Optional[Any]:
 if rng is None:
  A__	=    random.Random()

 A__	=    1
 for dim in shape:
  total_dims *= dim

 A__	=    []
 for _ in range(__lowerCamelCase						):
  values.append(rng.randint(0      ,			vocab_size - 1						)						)

 A__	=    np.array(__lowerCamelCase      ,			dtype=jnp.intaa						).reshape(__lowerCamelCase						)

 return output







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_=None						)						->		str:
 A__	=    ids_tensor(__lowerCamelCase      ,			vocab_size=2      ,			rng=__lowerCamelCase						)
 # make sure that at least one token is attended to for each batch
 A__	=    1
 return attn_mask




@require_flax
class      UpperCAmelCase_  :
    lowercase__						 =				None
    lowercase__						 =				()




    def 							__magic_name__							(   self						:	Optional[int]				)				->							Tuple:




     '''simple docstring'''

     A__,				A__	=    self.model_tester.prepare_config_and_inputs_for_common()

     # cut to half length & take max batch_size 3
     A__	=    2
     A__	=    inputs["input_ids"].shape[-1] // 2
     A__	=    inputs["input_ids"][:max_batch_size, :sequence_length]

     A__	=    jnp.ones_like(_lowerCamelCase				)
     A__	=    attention_mask[:max_batch_size, :sequence_length]

     # generate max 5 tokens
     A__	=    input_ids.shape[-1] + 5
     if config.eos_token_id is not None and config.pad_token_id is None:
      # hack to allow generate for models such as GPT2 as is done in `generate()`
      A__	=    config.eos_token_id
     return config, input_ids, attention_mask, max_length




    @is_pt_flax_cross_test
    def 							__magic_name__							(   self						:	Optional[Any]				)				->							Optional[Any]:




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    False
     A__	=    max_length
     A__	=    0

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
      A__	=    getattr(_lowerCamelCase     ,			_lowerCamelCase				)
      A__	=    pt_model_class(_lowerCamelCase				).eval()
      A__	=    load_flax_weights_in_pytorch_model(_lowerCamelCase     ,			flax_model.params				)

      A__	=    flax_model.generate(_lowerCamelCase				).sequences
      A__	=    pt_model.generate(torch.tensor(_lowerCamelCase     ,			dtype=torch.long				)				)

      if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
       A__	=    flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

      self.assertListEqual(pt_generation_outputs.numpy().tolist()     ,			flax_generation_outputs.tolist()				)




    def 							__magic_name__							(   self						:	int				)				->							Any:




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    False
     A__	=    max_length

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




    def 							__magic_name__							(   self						:	Any				)				->							Tuple:




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    True
     A__	=    max_length

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




    def 							__magic_name__							(   self						:	int				)				->							Dict:




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    False
     A__	=    max_length
     A__	=    2

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




    def 							__magic_name__							(   self						:	Dict				)				->							Optional[int]:




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    False
     A__	=    max_length
     A__	=    2
     A__	=    2

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[0]     ,			input_ids.shape[0] * config.num_return_sequences				)




    def 							__magic_name__							(   self						:	Union[str, Any]				)				->							str:




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    True
     A__	=    max_length
     A__	=    0.8
     A__	=    10
     A__	=    0.3
     A__	=    1
     A__	=    8
     A__	=    9

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




    def 							__magic_name__							(   self						:	int				)				->							Optional[Any]:




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    max_length
     A__	=    1
     A__	=    8
     A__	=    9

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




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




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()
     A__	=    max_length
     A__	=    2
     A__	=    1
     A__	=    8
     A__	=    9

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




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




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()

     # pad attention mask on the left
     A__	=    attention_mask.at[(0, 0)].set(0				)

     A__	=    False
     A__	=    max_length

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase     ,			attention_mask=_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase     ,			attention_mask=_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




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




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()

     # pad attention mask on the left
     A__	=    attention_mask.at[(0, 0)].set(0				)

     A__	=    True
     A__	=    max_length

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase     ,			attention_mask=_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase     ,			attention_mask=_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)




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




     '''simple docstring'''

     A__,				A__,				A__,				A__	=    self._get_input_ids_and_config()

     # pad attention mask on the left
     A__	=    attention_mask.at[(0, 0)].set(0				)

     A__	=    2
     A__	=    max_length

     for model_class in self.all_generative_model_classes:
      A__	=    model_class(_lowerCamelCase				)

      A__	=    model.generate(_lowerCamelCase     ,			attention_mask=_lowerCamelCase				).sequences
      self.assertEqual(generation_outputs.shape[-1]     ,			_lowerCamelCase				)

      A__	=    jit(model.generate				)
      A__	=    jit_generate(_lowerCamelCase     ,			attention_mask=_lowerCamelCase				).sequences

      self.assertListEqual(generation_outputs.tolist()     ,			jit_generation_outputs.tolist()				)






@require_flax
class      UpperCAmelCase_  (						unittest.TestCase							):




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




     '''simple docstring'''

     A__	=    AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"				)
     A__	=    FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"				)

     A__	=    "Hello world"
     A__	=    tokenizer(_lowerCamelCase     ,			return_tensors="np"				).input_ids

     # typos are quickly detected (the correct argument is `do_sample`)
     with self.assertRaisesRegex(_lowerCamelCase     ,			"do_samples"				):
      model.generate(_lowerCamelCase     ,			do_samples=_lowerCamelCase				)

     # arbitrary arguments that will not be used anywhere are also not accepted
     with self.assertRaisesRegex(_lowerCamelCase     ,			"foo"				):
      A__	=    {"foo": "bar"}
      model.generate(_lowerCamelCase     ,			**_lowerCamelCase				)











"""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       (							lowercase_						)						->		List[str]:
 A__	=    model.config

 A__	=    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      ,			)
 A__	=    MBartConfig(
     is_decoder=lowercase_      ,			is_encoder_decoder=lowercase_      ,			add_cross_attention=lowercase_      ,			decoder_layers=original_config.decoder_layer      ,			max_position_embeddings=original_config.max_position_embeddings      ,			vocab_size=len(
         model.decoder.tokenizer						)      ,			scale_embedding=lowercase_      ,			add_final_layer_norm=lowercase_      ,			)

 return encoder_config, decoder_config







def   _SCREAMING_SNAKE_CASE       (							lowercase_						)						->		int:
 if "encoder.model" in name:
  A__	=    name.replace("encoder.model"      ,			"encoder"						)
 if "decoder.model" in name:
  A__	=    name.replace("decoder.model"      ,			"decoder"						)
 if "patch_embed.proj" in name:
  A__	=    name.replace("patch_embed.proj"      ,			"embeddings.patch_embeddings.projection"						)
 if "patch_embed.norm" in name:
  A__	=    name.replace("patch_embed.norm"      ,			"embeddings.norm"						)
 if name.startswith("encoder"						):
  if "layers" in name:
   A__	=    "encoder." + name
  if "attn.proj" in name:
   A__	=    name.replace("attn.proj"      ,			"attention.output.dense"						)
  if "attn" in name and "mask" not in name:
   A__	=    name.replace("attn"      ,			"attention.self"						)
  if "norm1" in name:
   A__	=    name.replace("norm1"      ,			"layernorm_before"						)
  if "norm2" in name:
   A__	=    name.replace("norm2"      ,			"layernorm_after"						)
  if "mlp.fc1" in name:
   A__	=    name.replace("mlp.fc1"      ,			"intermediate.dense"						)
  if "mlp.fc2" in name:
   A__	=    name.replace("mlp.fc2"      ,			"output.dense"						)

  if name == "encoder.norm.weight":
   A__	=    "encoder.layernorm.weight"
  if name == "encoder.norm.bias":
   A__	=    "encoder.layernorm.bias"

 return name







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_						)						->		Any:
 for key in orig_state_dict.copy().keys():
  A__	=    orig_state_dict.pop(lowercase_						)

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

   if "weight" in key:
    A__	=    val[:dim, :]
    A__	=    val[dim : dim * 2, :]
    A__	=    val[-dim:, :]
   else:
    A__	=    val[:dim]
    A__	=    val[dim : dim * 2]
    A__	=    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:
   A__	=    val

 return orig_state_dict







def   _SCREAMING_SNAKE_CASE       (							lowercase_      ,			lowercase_=None      ,			lowercase_=False						)						->		Dict:
 # load original model
 A__	=    DonutModel.from_pretrained(lowercase_						).eval()

 # load HuggingFace model
 A__,				A__	=    get_configs(lowercase_						)
 A__	=    DonutSwinModel(lowercase_						)
 A__	=    MBartForCausalLM(lowercase_						)
 A__	=    VisionEncoderDecoderModel(encoder=lowercase_      ,			decoder=lowercase_						)
 model.eval()

 A__	=    original_model.state_dict()
 A__	=    convert_state_dict(lowercase_      ,			lowercase_						)
 model.load_state_dict(lowercase_						)

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

 A__	=    XLMRobertaTokenizerFast.from_pretrained(lowercase_      ,			from_slow=lowercase_						)
 A__	=    DonutImageProcessor(
     do_align_long_axis=original_model.config.align_long_axis      ,			size=original_model.config.input_size[::-1]						)
 A__	=    DonutProcessor(lowercase_      ,			lowercase_						)
 A__	=    processor(lowercase_      ,			return_tensors="pt"						).pixel_values

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

 A__	=    original_model.encoder.model.patch_embed(lowercase_						)
 A__,				A__	=    model.encoder.embeddings(lowercase_						)
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-3						)

 # verify encoder hidden states
 A__	=    original_model.encoder(lowercase_						)
 A__	=    model.encoder(lowercase_						).last_hidden_state
 assert torch.allclose(lowercase_      ,			lowercase_      ,			atol=1E-2						)

 # verify decoder hidden states
 A__	=    original_model(lowercase_      ,			lowercase_      ,			lowercase_						).logits
 A__	=    model(lowercase_      ,			decoder_input_ids=lowercase_						).logits
 assert torch.allclose(lowercase_      ,			lowercase_      ,			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(lowercase_						)
  processor.save_pretrained(lowercase_						)

 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__":
      SCREAMING_SNAKE_CASE    =	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.",
      )

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





