



import importlib.metadata
import warnings
from copy import deepcopy

from packaging import version

from ..utils import logging
from .import_utils import is_accelerate_available, is_bitsandbytes_available


if is_bitsandbytes_available():
  import bitsandbytes as bnb
  import torch
  import torch.nn as nn

  from ..pytorch_utils import ConvaD

if is_accelerate_available():
  from accelerate import init_empty_weights
  from accelerate.utils import find_tied_parameters

lowercase_             =     logging.get_logger(__name__)
def   a   (    A__			:     Union[str, Any]   ,     A__			:     Tuple   ,     A__			:     Union[str, Any]   ,     A__			:     List[str]=None   ,     A__			:     List[Any]=None  )					->   List[Any]:





      """simple docstring"""



      if "." in tensor_name:
            _lowercase			 =tensor_name.split('.'  )
            for split in splits[:-1]:
                  _lowercase			 =getattr(A__   ,     A__  )
                  if new_module is None:
                        raise ValueError(F'''{module} has no attribute {split}.'''  )
                  _lowercase			 =new_module
            _lowercase			 =splits[-1]

      if tensor_name not in module._parameters and tensor_name not in module._buffers:
            raise ValueError(F'''{module} does not have a parameter or a buffer named {tensor_name}.'''  )
      _lowercase			 =tensor_name in module._buffers
      _lowercase			 =getattr(A__   ,     A__  )

      if old_value.device == torch.device('meta'  ) and device not in ["meta", torch.device('meta'  )] and value is None:
            raise ValueError(F'''{tensor_name} is on the meta device, we need a `value` to put in on {device}.'''  )

      _lowercase			 =False
      _lowercase			 =False
      if is_buffer or not is_bitsandbytes_available():
            _lowercase			 =False
            _lowercase			 =False
      else:
            _lowercase			 =hasattr(bnb.nn   ,     'Params4bit'  ) and isinstance(module._parameters[tensor_name]   ,     bnb.nn.Paramsabit  )
            _lowercase			 =isinstance(module._parameters[tensor_name]   ,     bnb.nn.IntaParams  )

      if is_abit or is_abit:
            _lowercase			 =module._parameters[tensor_name]
            if param.device.type != "cuda":
                  if value is None:
                        _lowercase			 =old_value.to(A__  )
                  elif isinstance(A__   ,     torch.Tensor  ):
                        _lowercase			 =value.to('cpu'  )
                        if value.dtype == torch.inta:
                              _lowercase			 =version.parse(importlib.metadata.version('bitsandbytes'  )  ) > version.parse(
                                  '0.37.2'  )
                              if not is_abit_serializable:
                                    raise ValueError(
                                        'Detected int8 weights but the version of bitsandbytes is not compatible with int8 serialization. '
                                        'Make sure to download the latest `bitsandbytes` version. `pip install --upgrade bitsandbytes`.'  )
                  else:
                        _lowercase			 =torch.tensor(A__   ,     device='cpu'  )

                  # Support models using `Conv1D` in place of `nn.Linear` (e.g. gpt2) by transposing the weight matrix prior to quantization.
                  # Since weights are saved in the correct "orientation", we skip transposing when loading.
                  if issubclass(module.source_cls   ,     A__  ) and fpaa_statistics is None:
                        _lowercase			 =new_value.T

                  _lowercase			 =old_value.__dict__
                  if is_abit:
                        _lowercase			 =bnb.nn.IntaParams(A__   ,     requires_grad=A__   ,     **A__  ).to(A__  )
                  elif is_abit:
                        _lowercase			 =bnb.nn.Paramsabit(A__   ,     requires_grad=A__   ,     **A__  ).to(A__  )

                  _lowercase			 =new_value
                  if fpaa_statistics is not None:
                        setattr(module.weight   ,     'SCB'   ,     fpaa_statistics.to(A__  )  )

      else:
            if value is None:
                  _lowercase			 =old_value.to(A__  )
            elif isinstance(A__   ,     torch.Tensor  ):
                  _lowercase			 =value.to(A__  )
            else:
                  _lowercase			 =torch.tensor(A__   ,     device=A__  )

            if is_buffer:
                  _lowercase			 =new_value
            else:
                  _lowercase			 =nn.Parameter(A__   ,     requires_grad=old_value.requires_grad  )
                  _lowercase			 =new_value
def   a   (    A__			:     List[Any]   ,     A__			:     Tuple=None   ,     A__			:     Optional[int]=None   ,     A__			:     Tuple=None   ,     A__			:     int=False  )					->   Any:





      """simple docstring"""



      for name, module in model.named_children():
            if current_key_name is None:
                  _lowercase			 =[]
            current_key_name.append(A__  )

            if (isinstance(A__   ,     nn.Linear  ) or isinstance(A__   ,     A__  )) and name not in modules_to_not_convert:
                  # Check if the current key is not in the `modules_to_not_convert`
                  if not any(key in '.'.join(A__  ) for key in modules_to_not_convert  ):
                        with init_empty_weights():
                              if isinstance(A__   ,     A__  ):
                                    _lowercase , _lowercase			 =module.weight.shape
                              else:
                                    _lowercase			 =module.in_features
                                    _lowercase			 =module.out_features

                              if quantization_config.quantization_method() == "llm_int8":
                                    _lowercase			 =bnb.nn.LinearabitLt(
                                        A__   ,     A__   ,     module.bias is not None   ,     has_fpaa_weights=quantization_config.llm_inta_has_fpaa_weight   ,     threshold=quantization_config.llm_inta_threshold   ,     )
                                    _lowercase			 =True
                              else:
                                    if (
                                        quantization_config.llm_inta_skip_modules is not None
                                        and name in quantization_config.llm_inta_skip_modules
                                    ):
                                          pass
                                    else:
                                          _lowercase			 =bnb.nn.Linearabit(
                                              A__   ,     A__   ,     module.bias is not None   ,     quantization_config.bnb_abit_compute_dtype   ,     compress_statistics=quantization_config.bnb_abit_use_double_quant   ,     quant_type=quantization_config.bnb_abit_quant_type   ,     )
                                          _lowercase			 =True
                    # Store the module class in case we need to transpose the weight later
                              _lowercase			 =type(A__  )
                              # Force requires grad to False to avoid unexpected errors
                              model._modules[name].requires_grad_(A__  )
            if len(list(module.children()  )  ) > 0:
                  _lowercase , _lowercase			 =_replace_with_bnb_linear(
                      A__   ,     A__   ,     A__   ,     A__   ,     has_been_replaced=A__   ,     )
            # Remove the last key for recursion
            current_key_name.pop(-1  )
      return model, has_been_replaced
def   a   (    A__			:     Any   ,     A__			:     Union[str, Any]=None   ,     A__			:     Optional[Any]=None   ,     A__			:     int=None  )					->   Tuple:





      """simple docstring"""



      _lowercase			 =['lm_head'] if modules_to_not_convert is None else modules_to_not_convert
      _lowercase , _lowercase			 =_replace_with_bnb_linear(
          A__   ,     A__   ,     A__   ,     A__  )

      if not has_been_replaced:
            logger.warning(
                'You are loading your model in 8bit or 4bit but no linear modules were found in your model.'
                ' Please double check your model architecture, or submit an issue on github if you think this is'
                ' a bug.'  )

      return model
def   a   (    *A__			:     Optional[Any]   ,     **A__			:     List[Any]  )					->   Optional[Any]:





      """simple docstring"""



      warnings.warn(
          '`replace_8bit_linear` will be deprecated in a future version, please use `replace_with_bnb_linear` instead'   ,     A__   ,     )
      return replace_with_bnb_linear(*A__   ,     **A__  )
def   a   (    *A__			:     Dict   ,     **A__			:     Optional[Any]  )					->   str:





      """simple docstring"""



      warnings.warn(
          '`set_module_8bit_tensor_to_device` will be deprecated in a future version, please use `set_module_quantized_tensor_to_device` instead'   ,     A__   ,     )
      return set_module_quantized_tensor_to_device(*A__   ,     **A__  )







def   a   (    A__			:     List[str]  )					->   Any:





      """simple docstring"""



      _lowercase			 =deepcopy(A__  )  # this has 0 cost since it is done inside `init_empty_weights` context manager`
      tied_model.tie_weights()

      _lowercase			 =find_tied_parameters(A__  )
      # For compatibility with Accelerate < 0.18
      if isinstance(A__   ,     A__  ):
            _lowercase			 =sum(list(tied_params.values()  )   ,     []  ) + list(tied_params.keys()  )
      else:
            _lowercase			 =sum(A__   ,     []  )
      _lowercase			 =len(A__  ) > 0

      # Check if it is a base model
      _lowercase			 =not hasattr(A__   ,     model.base_model_prefix  )

      # Ignore this for base models (BertModel, GPT2Model, etc.)
      if (not has_tied_params) and is_base_model:
            return []

      # otherwise they have an attached head
      _lowercase			 =list(model.named_children()  )
      _lowercase			 =[list_modules[-1][0]]

      # add last module together with tied weights
      _lowercase			 =set(A__  ) - set(A__  )
      _lowercase			 =list(set(A__  )  ) + list(A__  )

      # remove ".weight" from the keys
      _lowercase			 =['.weight', '.bias']
      _lowercase			 =[]
      for name in list_untouched:
            for name_to_remove in names_to_remove:
                  if name_to_remove in name:
                        _lowercase			 =name.replace(A__   ,     ''  )
            filtered_module_names.append(A__  )

      return filtered_module_names





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   (    A__			:     Tuple  )					->   Any:





      """simple docstring"""



      _lowercase			 =model.config

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

      return encoder_config, decoder_config
def   a   (    A__			:     Union[str, Any]  )					->   str:





      """simple docstring"""



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

            if name == "encoder.norm.weight":
                  _lowercase			 ='encoder.layernorm.weight'
            if name == "encoder.norm.bias":
                  _lowercase			 ='encoder.layernorm.bias'

      return name
def   a   (    A__			:     Optional[Any]   ,     A__			:     List[Any]  )					->   List[str]:





      """simple docstring"""



      for key in orig_state_dict.copy().keys():
            _lowercase			 =orig_state_dict.pop(A__  )

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

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

      return orig_state_dict
def   a   (    A__			:     str   ,     A__			:     List[str]=None   ,     A__			:     List[Any]=False  )					->   List[str]:





      """simple docstring"""



      _lowercase			 =DonutModel.from_pretrained(A__  ).eval()

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

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

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

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

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

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

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

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

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

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


