




import argparse
import collections
import os
import re
import tempfile

import pandas as pd
from datasets import Dataset
from huggingface_hub import hf_hub_download, upload_folder

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/update_metadata.py
__magic_name__  :  str					=       """src/transformers"""


# This is to make sure the transformers module imported is the one in the repo.
__magic_name__  :  Optional[int]					=       direct_transformers_import(TRANSFORMERS_PATH)


# Regexes that match TF/Flax/PT model names.
__magic_name__  :  Union[str, Any]					=       re.compile(R"""TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)""")
__magic_name__  :  Optional[Any]					=       re.compile(R"""Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)""")
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
__magic_name__  :  List[str]					=       re.compile(R"""(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)""")


# Fill this with tuples (pipeline_tag, model_mapping, auto_model)
__magic_name__  :  Dict					=       [
    ("""pretraining""", """MODEL_FOR_PRETRAINING_MAPPING_NAMES""", """AutoModelForPreTraining"""),
    ("""feature-extraction""", """MODEL_MAPPING_NAMES""", """AutoModel"""),
    ("""audio-classification""", """MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES""", """AutoModelForAudioClassification"""),
    ("""text-generation""", """MODEL_FOR_CAUSAL_LM_MAPPING_NAMES""", """AutoModelForCausalLM"""),
    ("""automatic-speech-recognition""", """MODEL_FOR_CTC_MAPPING_NAMES""", """AutoModelForCTC"""),
    ("""image-classification""", """MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES""", """AutoModelForImageClassification"""),
    ("""image-segmentation""", """MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES""", """AutoModelForImageSegmentation"""),
    ("""fill-mask""", """MODEL_FOR_MASKED_LM_MAPPING_NAMES""", """AutoModelForMaskedLM"""),
    ("""object-detection""", """MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES""", """AutoModelForObjectDetection"""),
    (
        """zero-shot-object-detection""",
        """MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES""",
        """AutoModelForZeroShotObjectDetection""",
    ),
    ("""question-answering""", """MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES""", """AutoModelForQuestionAnswering"""),
    ("""text2text-generation""", """MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES""", """AutoModelForSeq2SeqLM"""),
    ("""text-classification""", """MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES""", """AutoModelForSequenceClassification"""),
    ("""automatic-speech-recognition""", """MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES""", """AutoModelForSpeechSeq2Seq"""),
    (
        """table-question-answering""",
        """MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES""",
        """AutoModelForTableQuestionAnswering""",
    ),
    ("""token-classification""", """MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES""", """AutoModelForTokenClassification"""),
    ("""multiple-choice""", """MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES""", """AutoModelForMultipleChoice"""),
    (
        """next-sentence-prediction""",
        """MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES""",
        """AutoModelForNextSentencePrediction""",
    ),
    (
        """audio-frame-classification""",
        """MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES""",
        """AutoModelForAudioFrameClassification""",
    ),
    ("""audio-xvector""", """MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES""", """AutoModelForAudioXVector"""),
    (
        """document-question-answering""",
        """MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES""",
        """AutoModelForDocumentQuestionAnswering""",
    ),
    (
        """visual-question-answering""",
        """MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES""",
        """AutoModelForVisualQuestionAnswering""",
    ),
    ("""image-to-text""", """MODEL_FOR_FOR_VISION_2_SEQ_MAPPING_NAMES""", """AutoModelForVision2Seq"""),
    (
        """zero-shot-image-classification""",
        """MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES""",
        """AutoModelForZeroShotImageClassification""",
    ),
    ("""depth-estimation""", """MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES""", """AutoModelForDepthEstimation"""),
    ("""video-classification""", """MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES""", """AutoModelForVideoClassification"""),
    ("""mask-generation""", """MODEL_FOR_MASK_GENERATION_MAPPING_NAMES""", """AutoModelForMaskGeneration"""),
]
def 		a_							( __lowerCAmelCase				):
  lowerCAmelCase__										=      re.finditer('''.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'''   , __lowerCAmelCase				)
  return [m.group(0				) for m in matches]
def 		a_							( ):
  lowerCAmelCase__										=      transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
  lowerCAmelCase__										=      {
      config.replace('''Config'''   , ''''''				): model_type for model_type, config in config_maping_names.items()
  }

  # Dictionaries flagging if each model prefix has a backend in PT/TF/Flax.
  lowerCAmelCase__										=      collections.defaultdict(__lowerCAmelCase				)
  lowerCAmelCase__										=      collections.defaultdict(__lowerCAmelCase				)
  lowerCAmelCase__										=      collections.defaultdict(__lowerCAmelCase				)

  # Let's lookup through all transformers object (once) and find if models are supported by a given backend.
  for attr_name in dir(__lowerCAmelCase				):
    lowerCAmelCase__										=      None
    if _re_tf_models.match(__lowerCAmelCase				) is not None:
      lowerCAmelCase__										=      tf_models
      lowerCAmelCase__										=      _re_tf_models.match(__lowerCAmelCase				).groups()[0]
    elif _re_flax_models.match(__lowerCAmelCase				) is not None:
      lowerCAmelCase__										=      flax_models
      lowerCAmelCase__										=      _re_flax_models.match(__lowerCAmelCase				).groups()[0]
    elif _re_pt_models.match(__lowerCAmelCase				) is not None:
      lowerCAmelCase__										=      pt_models
      lowerCAmelCase__										=      _re_pt_models.match(__lowerCAmelCase				).groups()[0]

    if lookup_dict is not None:
      while len(__lowerCAmelCase				) > 0:
        if attr_name in model_prefix_to_model_type:
          lowerCAmelCase__										=      True
          break
        # Try again after removing the last word in the name
        lowerCAmelCase__										=      ''''''.join(camel_case_split(__lowerCAmelCase				)[:-1]				)

  lowerCAmelCase__										=      set(list(pt_models.keys()				) + list(tf_models.keys()				) + list(flax_models.keys()				)				)
  lowerCAmelCase__										=      list(__lowerCAmelCase				)
  all_models.sort()

  lowerCAmelCase__										=      {'''model_type''': all_models}
  lowerCAmelCase__										=      [pt_models[t] for t in all_models]
  lowerCAmelCase__										=      [tf_models[t] for t in all_models]
  lowerCAmelCase__										=      [flax_models[t] for t in all_models]

  # Now let's use the auto-mapping names to make sure
  lowerCAmelCase__										=      {}
  for t in all_models:
    if t in transformers_module.models.auto.processing_auto.PROCESSOR_MAPPING_NAMES:
      lowerCAmelCase__										=      '''AutoProcessor'''
    elif t in transformers_module.models.auto.tokenization_auto.TOKENIZER_MAPPING_NAMES:
      lowerCAmelCase__										=      '''AutoTokenizer'''
    elif t in transformers_module.models.auto.feature_extraction_auto.FEATURE_EXTRACTOR_MAPPING_NAMES:
      lowerCAmelCase__										=      '''AutoFeatureExtractor'''
    else:
      # Default to AutoTokenizer if a model has nothing, for backward compatibility.
      lowerCAmelCase__										=      '''AutoTokenizer'''

  lowerCAmelCase__										=      [processors[t] for t in all_models]

  return pd.DataFrame(__lowerCAmelCase				)
def 		a_							( __lowerCAmelCase				):
  lowerCAmelCase__										=      [
      transformers_module.models.auto.modeling_auto,
      transformers_module.models.auto.modeling_tf_auto,
      transformers_module.models.auto.modeling_flax_auto,
  ]
  for pipeline_tag, model_mapping, auto_class in PIPELINE_TAGS_AND_AUTO_MODELS:
    lowerCAmelCase__										=      [model_mapping, F"""TF_{model_mapping}""", F"""FLAX_{model_mapping}"""]
    lowerCAmelCase__										=      [auto_class, F"""TF_{auto_class}""", F"""Flax_{auto_class}"""]
    # Loop through all three frameworks
    for module, cls, mapping in zip(__lowerCAmelCase   , __lowerCAmelCase   , __lowerCAmelCase				):
      # The type of pipeline may not exist in this framework
      if not hasattr(__lowerCAmelCase   , __lowerCAmelCase				):
        continue
      # First extract all model_names
      lowerCAmelCase__										=      []
      for name in getattr(__lowerCAmelCase   , __lowerCAmelCase				).values():
        if isinstance(__lowerCAmelCase   , __lowerCAmelCase				):
          model_names.append(__lowerCAmelCase				)
        else:
          model_names.extend(list(__lowerCAmelCase				)				)

            # Add pipeline tag and auto model class for those models
      table.update({model_name: (pipeline_tag, cls) for model_name in model_names}				)

  return table
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase				):
  lowerCAmelCase__										=      get_frameworks_table()
  lowerCAmelCase__										=      Dataset.from_pandas(__lowerCAmelCase				)

  lowerCAmelCase__										=      hf_hub_download(
      '''huggingface/transformers-metadata'''   , '''pipeline_tags.json'''   , repo_type='''dataset'''   , token=__lowerCAmelCase				)
  lowerCAmelCase__										=      Dataset.from_json(__lowerCAmelCase				)
  lowerCAmelCase__										=      {
      tags_dataset[i]['''model_class''']: (tags_dataset[i]['''pipeline_tag'''], tags_dataset[i]['''auto_class'''])
      for i in range(len(__lowerCAmelCase				)				)
  }
  lowerCAmelCase__										=      update_pipeline_and_auto_class_table(__lowerCAmelCase				)

  # Sort the model classes to avoid some nondeterministic updates to create false update commits.
  lowerCAmelCase__										=      sorted(table.keys()				)
  lowerCAmelCase__										=      pd.DataFrame(
      {
          '''model_class''': model_classes,
          '''pipeline_tag''': [table[m][0] for m in model_classes],
          '''auto_class''': [table[m][1] for m in model_classes],
      }				)
  lowerCAmelCase__										=      Dataset.from_pandas(__lowerCAmelCase				)

  with tempfile.TemporaryDirectory() as tmp_dir:
    frameworks_dataset.to_json(os.path.join(__lowerCAmelCase   , '''frameworks.json'''				)				)
    tags_dataset.to_json(os.path.join(__lowerCAmelCase   , '''pipeline_tags.json'''				)				)

    if commit_sha is not None:
      lowerCAmelCase__										=      (
          F"""Update with commit {commit_sha}\n\nSee: """
          F"""https://github.com/huggingface/transformers/commit/{commit_sha}"""
      )
    else:
      lowerCAmelCase__										=      '''Update'''

    upload_folder(
        repo_id='''huggingface/transformers-metadata'''   , folder_path=__lowerCAmelCase   , repo_type='''dataset'''   , token=__lowerCAmelCase   , commit_message=__lowerCAmelCase   , )
def 		a_							( ):
  lowerCAmelCase__										=      {tag: cls for tag, _, cls in PIPELINE_TAGS_AND_AUTO_MODELS}
  lowerCAmelCase__										=      transformers_module.pipelines.SUPPORTED_TASKS
  lowerCAmelCase__										=      []
  for key in pipeline_tasks:
    if key not in in_table:
      lowerCAmelCase__										=      pipeline_tasks[key]['''pt''']
      if isinstance(__lowerCAmelCase   , (list, tuple)				):
        lowerCAmelCase__										=      model[0]
      lowerCAmelCase__										=      model.__name__
      if model not in in_table.values():
        missing.append(__lowerCAmelCase				)

  if len(__lowerCAmelCase				) > 0:
    lowerCAmelCase__										=      ''', '''.join(__lowerCAmelCase				)
    raise ValueError(
        '''The following pipeline tags are not present in the `PIPELINE_TAGS_AND_AUTO_MODELS` constant inside '''
        F"""`utils/update_metadata.py`: {msg}. Please add them!"""				)


if __name__ == "__main__":
      __magic_name__  :  List[str]					=       argparse.ArgumentParser()
      parser.add_argument("""--token""", type=str, help="""The token to use to push to the transformers-metadata dataset.""")
      parser.add_argument("""--commit_sha""", type=str, help="""The sha of the commit going with this update.""")
      parser.add_argument("""--check-only""", action="""store_true""", help="""Activate to just check all pipelines are present.""")
      __magic_name__  :  Optional[int]					=       parser.parse_args()

      if args.check_only:
            check_pipeline_tags()
      else:
            update_metadata(args.token, args.commit_sha)







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_							( __lowerCAmelCase				):
  lowerCAmelCase__										=      model.config

  lowerCAmelCase__										=      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   , )
  lowerCAmelCase__										=      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_							( __lowerCAmelCase				):
  if "encoder.model" in name:
    lowerCAmelCase__										=      name.replace('''encoder.model'''   , '''encoder'''				)
  if "decoder.model" in name:
    lowerCAmelCase__										=      name.replace('''decoder.model'''   , '''decoder'''				)
  if "patch_embed.proj" in name:
    lowerCAmelCase__										=      name.replace('''patch_embed.proj'''   , '''embeddings.patch_embeddings.projection'''				)
  if "patch_embed.norm" in name:
    lowerCAmelCase__										=      name.replace('''patch_embed.norm'''   , '''embeddings.norm'''				)
  if name.startswith('''encoder'''				):
    if "layers" in name:
      lowerCAmelCase__										=      '''encoder.''' + name
    if "attn.proj" in name:
      lowerCAmelCase__										=      name.replace('''attn.proj'''   , '''attention.output.dense'''				)
    if "attn" in name and "mask" not in name:
      lowerCAmelCase__										=      name.replace('''attn'''   , '''attention.self'''				)
    if "norm1" in name:
      lowerCAmelCase__										=      name.replace('''norm1'''   , '''layernorm_before'''				)
    if "norm2" in name:
      lowerCAmelCase__										=      name.replace('''norm2'''   , '''layernorm_after'''				)
    if "mlp.fc1" in name:
      lowerCAmelCase__										=      name.replace('''mlp.fc1'''   , '''intermediate.dense'''				)
    if "mlp.fc2" in name:
      lowerCAmelCase__										=      name.replace('''mlp.fc2'''   , '''output.dense'''				)

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

  return name
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase				):
  for key in orig_state_dict.copy().keys():
    lowerCAmelCase__										=      orig_state_dict.pop(__lowerCAmelCase				)

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

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

  return orig_state_dict
def 		a_							( __lowerCAmelCase   , __lowerCAmelCase=None   , __lowerCAmelCase=False				):
  # load original model
  lowerCAmelCase__										=      DonutModel.from_pretrained(__lowerCAmelCase				).eval()

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

  lowerCAmelCase__										=      original_model.state_dict()
  lowerCAmelCase__										=      convert_state_dict(__lowerCAmelCase   , __lowerCAmelCase				)
  model.load_state_dict(__lowerCAmelCase				)

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

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

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

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

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

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

      __magic_name__  :  Any					=       parser.parse_args()
      convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)

