





import unittest

from transformers import MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING, AutoTokenizer, is_vision_available
from transformers.pipelines import pipeline
from transformers.pipelines.document_question_answering import apply_tesseract
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_detectrona,
    require_pytesseract,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
  from PIL import Image

  from transformers.image_utils import load_image
else:






  class 					lowercase__ :

     '''simple docstring'''


     @staticmethod
     def    UpperCamelCase__						(					*__magic_name__,	**__magic_name__							)				->   List[Any]:






         """simple docstring"""
         pass




  def 	lowerCAmelCase_					(						__UpperCAmelCase:    List[str]   )   ->			Any:
      return None


# This is a pinned image from a specific revision of a document question answering space, hosted by HuggingFace,
# so we can expect it to be available.
UpperCAmelCase_  							=				(
    'https://huggingface.co/spaces/impira/docquery/resolve/2f6c96314dc84dfda62d40de9da55f2f5165d403/invoice.png'
)






@is_pipeline_test
@require_torch
@require_vision
class 					lowercase__ (			unittest.TestCase       ):

   '''simple docstring'''


   a		:							List[str]     	=	MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING




   @require_pytesseract
   @require_vision
   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__,	__magic_name__							)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       Optional[Any]      =   pipeline(
           '''document-question-answering''',	model=__magic_name__,	tokenizer=__magic_name__,	image_processor=__magic_name__							)

       UpperCamelCase__							:       Tuple      =   INVOICE_URL
       UpperCamelCase__							:       Union[str, Any]      =   list(zip(*apply_tesseract(load_image(__magic_name__							),	__magic_name__,	''''''							)							)							)
       UpperCamelCase__							:       int      =   '''What is the placebo?'''
       UpperCamelCase__							:       Optional[int]      =   [
           {
               '''image''': load_image(__magic_name__							),
               '''question''': question,
           },
           {
               '''image''': image,
               '''question''': question,
           },
           {
               '''image''': image,
               '''question''': question,
               '''word_boxes''': word_boxes,
           },
       ]
       return dqa_pipeline, examples




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__							)				->   Tuple:






       """simple docstring"""
       UpperCamelCase__							:       Optional[int]      =   dqa_pipeline(__magic_name__,	top_k=2							)
       self.assertEqual(
           __magic_name__,	[
               [
                   {'''score''': ANY(__magic_name__							), '''answer''': ANY(__magic_name__							), '''start''': ANY(__magic_name__							), '''end''': ANY(__magic_name__							)},
                   {'''score''': ANY(__magic_name__							), '''answer''': ANY(__magic_name__							), '''start''': ANY(__magic_name__							), '''end''': ANY(__magic_name__							)},
               ]
           ]
           * 3,	)




   @require_torch
   @require_detectrona
   @require_pytesseract
   def    UpperCamelCase__						(					self							)				->   Optional[int]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   pipeline('''document-question-answering''',	model='''hf-internal-testing/tiny-random-layoutlmv2'''							)
       UpperCamelCase__							:       Tuple      =   INVOICE_URL
       UpperCamelCase__							:       Dict      =   '''How many cats are there?'''

       UpperCamelCase__							:       int      =   [
           {'''score''': 0.0001, '''answer''': '''oy 2312/2019''', '''start''': 38, '''end''': 39},
           {'''score''': 0.0001, '''answer''': '''oy 2312/2019 DUE''', '''start''': 38, '''end''': 40},
       ]
       UpperCamelCase__							:       List[Any]      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	top_k=2							)
       self.assertEqual(nested_simplify(__magic_name__,	decimals=4							),	__magic_name__							)

       UpperCamelCase__							:       Any      =   dqa_pipeline({'''image''': image, '''question''': question},	top_k=2							)
       self.assertEqual(nested_simplify(__magic_name__,	decimals=4							),	__magic_name__							)

       # This image does not detect ANY text in it, meaning layoutlmv2 should fail.
       # Empty answer probably
       UpperCamelCase__							:       Optional[Any]      =   '''./tests/fixtures/tests_samples/COCO/000000039769.png'''
       UpperCamelCase__							:       Any      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	top_k=2							)
       self.assertEqual(__magic_name__,	[]							)

       # We can optionnally pass directly the words and bounding boxes
       UpperCamelCase__							:       List[Any]      =   '''./tests/fixtures/tests_samples/COCO/000000039769.png'''
       UpperCamelCase__							:       Union[str, Any]      =   []
       UpperCamelCase__							:       Optional[Any]      =   []
       UpperCamelCase__							:       Optional[int]      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	words=__magic_name__,	boxes=__magic_name__,	top_k=2							)
       self.assertEqual(__magic_name__,	[]							)




   @slow
   @require_torch
   @require_detectrona
   @require_pytesseract
   def    UpperCamelCase__						(					self							)				->   Any:






       """simple docstring"""
       UpperCamelCase__							:       Optional[int]      =   pipeline(
           '''document-question-answering''',	model='''tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa''',	revision='''9977165''',	)
       UpperCamelCase__							:       Dict      =   INVOICE_URL
       UpperCamelCase__							:       str      =   '''What is the invoice number?'''

       UpperCamelCase__							:       Any      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.9944, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               {'''score''': 0.0009, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
           ],	)

       UpperCamelCase__							:       int      =   dqa_pipeline({'''image''': image, '''question''': question},	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.9944, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               {'''score''': 0.0009, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
           ],	)

       UpperCamelCase__							:       Tuple      =   dqa_pipeline(
           [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}],	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               [
                   {'''score''': 0.9944, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                   {'''score''': 0.0009, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               ],
           ]
           * 2,	)




   @slow
   @require_torch
   @require_detectrona
   @require_pytesseract
   def    UpperCamelCase__						(					self							)				->   Union[str, Any]:






       """simple docstring"""
       UpperCamelCase__							:       Union[str, Any]      =   pipeline(
           '''document-question-answering''',	model='''tiennvcs/layoutlmv2-base-uncased-finetuned-docvqa''',	revision='''9977165''',	max_seq_len=50,	)
       UpperCamelCase__							:       Optional[Any]      =   INVOICE_URL
       UpperCamelCase__							:       List[Any]      =   '''What is the invoice number?'''

       UpperCamelCase__							:       List[Any]      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.9974, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
               {'''score''': 0.9948, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
           ],	)

       UpperCamelCase__							:       Tuple      =   dqa_pipeline({'''image''': image, '''question''': question},	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.9974, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
               {'''score''': 0.9948, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
           ],	)

       UpperCamelCase__							:       Optional[int]      =   dqa_pipeline(
           [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}],	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               [
                   {'''score''': 0.9974, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
                   {'''score''': 0.9948, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               ]
           ]
           * 2,	)




   @slow
   @require_torch
   @require_pytesseract
   @require_vision
   def    UpperCamelCase__						(					self							)				->   Optional[int]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   AutoTokenizer.from_pretrained(
           '''impira/layoutlm-document-qa''',	revision='''3dc6de3''',	add_prefix_space=__magic_name__							)
       UpperCamelCase__							:       List[str]      =   pipeline(
           '''document-question-answering''',	model='''impira/layoutlm-document-qa''',	tokenizer=__magic_name__,	revision='''3dc6de3''',	)
       UpperCamelCase__							:       Union[str, Any]      =   INVOICE_URL
       UpperCamelCase__							:       str      =   '''What is the invoice number?'''

       UpperCamelCase__							:       Tuple      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.4251, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               {'''score''': 0.0819, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
           ],	)

       UpperCamelCase__							:       List[Any]      =   dqa_pipeline({'''image''': image, '''question''': question},	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.4251, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               {'''score''': 0.0819, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
           ],	)

       UpperCamelCase__							:       Tuple      =   dqa_pipeline(
           [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}],	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               [
                   {'''score''': 0.4251, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                   {'''score''': 0.0819, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
               ]
           ]
           * 2,	)

       UpperCamelCase__							:       str      =   list(zip(*apply_tesseract(load_image(__magic_name__							),	__magic_name__,	''''''							)							)							)

       # This model should also work if `image` is set to None
       UpperCamelCase__							:       Any      =   dqa_pipeline({'''image''': None, '''word_boxes''': word_boxes, '''question''': question},	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.4251, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               {'''score''': 0.0819, '''answer''': '''1110212019''', '''start''': 23, '''end''': 23},
           ],	)




   @slow
   @require_torch
   @require_pytesseract
   @require_vision
   def    UpperCamelCase__						(					self							)				->   Union[str, Any]:






       """simple docstring"""
       UpperCamelCase__							:       List[Any]      =   AutoTokenizer.from_pretrained(
           '''impira/layoutlm-document-qa''',	revision='''3dc6de3''',	add_prefix_space=__magic_name__							)
       UpperCamelCase__							:       int      =   pipeline(
           '''document-question-answering''',	model='''impira/layoutlm-document-qa''',	tokenizer=__magic_name__,	revision='''3dc6de3''',	max_seq_len=50,	)
       UpperCamelCase__							:       List[str]      =   INVOICE_URL
       UpperCamelCase__							:       Union[str, Any]      =   '''What is the invoice number?'''

       UpperCamelCase__							:       Dict      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.9999, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               {'''score''': 0.9998, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
           ],	)

       UpperCamelCase__							:       Tuple      =   dqa_pipeline(
           [{'''image''': image, '''question''': question}, {'''image''': image, '''question''': question}],	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               [
                   {'''score''': 0.9999, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
                   {'''score''': 0.9998, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               ]
           ]
           * 2,	)

       UpperCamelCase__							:       Tuple      =   list(zip(*apply_tesseract(load_image(__magic_name__							),	__magic_name__,	''''''							)							)							)

       # This model should also work if `image` is set to None
       UpperCamelCase__							:       Dict      =   dqa_pipeline({'''image''': None, '''word_boxes''': word_boxes, '''question''': question},	top_k=2							)
       self.assertEqual(
           nested_simplify(__magic_name__,	decimals=4							),	[
               {'''score''': 0.9999, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
               {'''score''': 0.9998, '''answer''': '''us-001''', '''start''': 16, '''end''': 16},
           ],	)




   @slow
   @require_torch
   def    UpperCamelCase__						(					self							)				->   Optional[int]:






       """simple docstring"""
       UpperCamelCase__							:       str      =   pipeline(
           '''document-question-answering''',	model='''naver-clova-ix/donut-base-finetuned-docvqa''',	tokenizer=AutoTokenizer.from_pretrained('''naver-clova-ix/donut-base-finetuned-docvqa'''							),	feature_extractor='''naver-clova-ix/donut-base-finetuned-docvqa''',	)

       UpperCamelCase__							:       List[str]      =   INVOICE_URL
       UpperCamelCase__							:       int      =   '''What is the invoice number?'''
       UpperCamelCase__							:       Tuple      =   dqa_pipeline(image=__magic_name__,	question=__magic_name__,	top_k=2							)
       self.assertEqual(nested_simplify(__magic_name__,	decimals=4							),	[{'''answer''': '''us-001'''}]							)




   @require_tf
   @unittest.skip('''Document question answering not implemented in TF'''							)
   def    UpperCamelCase__						(					self							)				->   int:






       """simple docstring"""
       pass








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 	lowerCAmelCase_					(						__UpperCAmelCase:    Union[str, Any]   )   ->			Union[str, Any]:
    UpperCamelCase__							:       List[str]      =   model.config

    UpperCamelCase__							:       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=128	,     )
    UpperCamelCase__							:       Union[str, Any]      =   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 	lowerCAmelCase_					(						__UpperCAmelCase:    str   )   ->			Tuple:
    if "encoder.model" in name:
        UpperCamelCase__							:       str      =   name.replace('''encoder.model'''	,     '''encoder'''   )
    if "decoder.model" in name:
        UpperCamelCase__							:       Optional[int]      =   name.replace('''decoder.model'''	,     '''decoder'''   )
    if "patch_embed.proj" in name:
        UpperCamelCase__							:       Dict      =   name.replace('''patch_embed.proj'''	,     '''embeddings.patch_embeddings.projection'''   )
    if "patch_embed.norm" in name:
        UpperCamelCase__							:       Optional[int]      =   name.replace('''patch_embed.norm'''	,     '''embeddings.norm'''   )
    if name.startswith('''encoder'''   ):
        if "layers" in name:
            UpperCamelCase__							:       Optional[Any]      =   '''encoder.''' + name
        if "attn.proj" in name:
            UpperCamelCase__							:       Optional[int]      =   name.replace('''attn.proj'''	,     '''attention.output.dense'''   )
        if "attn" in name and "mask" not in name:
            UpperCamelCase__							:       str      =   name.replace('''attn'''	,     '''attention.self'''   )
        if "norm1" in name:
            UpperCamelCase__							:       List[str]      =   name.replace('''norm1'''	,     '''layernorm_before'''   )
        if "norm2" in name:
            UpperCamelCase__							:       Union[str, Any]      =   name.replace('''norm2'''	,     '''layernorm_after'''   )
        if "mlp.fc1" in name:
            UpperCamelCase__							:       Any      =   name.replace('''mlp.fc1'''	,     '''intermediate.dense'''   )
        if "mlp.fc2" in name:
            UpperCamelCase__							:       List[str]      =   name.replace('''mlp.fc2'''	,     '''output.dense'''   )

        if name == "encoder.norm.weight":
            UpperCamelCase__							:       Optional[Any]      =   '''encoder.layernorm.weight'''
        if name == "encoder.norm.bias":
            UpperCamelCase__							:       Any      =   '''encoder.layernorm.bias'''

    return name




def 	lowerCAmelCase_					(						__UpperCAmelCase:    str	,     __UpperCAmelCase:    str   )   ->			Dict:
    for key in orig_state_dict.copy().keys():
        UpperCamelCase__							:       Dict      =   orig_state_dict.pop(__UpperCAmelCase   )

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

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

    return orig_state_dict




def 	lowerCAmelCase_					(						__UpperCAmelCase:    Optional[Any]	,     __UpperCAmelCase:    List[Any]=None	,     __UpperCAmelCase:    Tuple=False   )   ->			Optional[int]:
    # load original model
    UpperCamelCase__							:       List[Any]      =   DonutModel.from_pretrained(__UpperCAmelCase   ).eval()

    # load HuggingFace model
    UpperCamelCase__     ,UpperCamelCase__							:       List[str]      =   get_configs(__UpperCAmelCase   )
    UpperCamelCase__							:       int      =   DonutSwinModel(__UpperCAmelCase   )
    UpperCamelCase__							:       Optional[Any]      =   MBartForCausalLM(__UpperCAmelCase   )
    UpperCamelCase__							:       Optional[int]      =   VisionEncoderDecoderModel(encoder=__UpperCAmelCase	,     decoder=__UpperCAmelCase   )
    model.eval()

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

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

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

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

    UpperCamelCase__							:       List[str]      =   original_model.encoder.model.patch_embed(__UpperCAmelCase   )
    UpperCamelCase__     ,UpperCamelCase__							:       List[str]      =   model.encoder.embeddings(__UpperCAmelCase   )
    assert torch.allclose(__UpperCAmelCase	,     __UpperCAmelCase	,     atol=1e-3   )

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

    # verify decoder hidden states
    UpperCamelCase__							:       Optional[Any]      =   original_model(__UpperCAmelCase	,     __UpperCAmelCase	,     __UpperCAmelCase   ).logits
    UpperCamelCase__							:       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_  							=				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_  							=				parser.parse_args()
  convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)

