


"""simple docstring"""



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
__snake_case       =       """src/transformers"""


# This is to make sure the transformers module imported is the one in the repo.
__snake_case       =       direct_transformers_import(TRANSFORMERS_PATH)


# Regexes that match TF/Flax/PT model names.
__snake_case       =       re.compile(R"""TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)""")
__snake_case       =       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.
__snake_case       =       re.compile(R"""(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)""")


# Fill this with tuples (pipeline_tag, model_mapping, auto_model)
__snake_case       =       [
    ("""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     __lowerCAmelCase		(		lowercase    :    Union[str, Any]					)      ->		int:


    """simple docstring"""





    snake_case       : Tuple    =       re.finditer(".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)"		,				lowercase					)
    return [m.group(0					) for m in matches]







def     __lowerCAmelCase		(		)      ->		Dict:


    """simple docstring"""





    snake_case       : Any    =       transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
    snake_case       : int    =       {
        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.
    snake_case       : str    =       collections.defaultdict(lowercase					)
    snake_case       : Optional[Any]    =       collections.defaultdict(lowercase					)
    snake_case       : Tuple    =       collections.defaultdict(lowercase					)

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

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

    snake_case       : List[Any]    =       set(list(pt_models.keys()					) + list(tf_models.keys()					) + list(flax_models.keys()					)					)
    snake_case       : Dict    =       list(lowercase					)
    all_models.sort()

    snake_case       : Tuple    =       {"model_type": all_models}
    snake_case       : str    =       [pt_models[t] for t in all_models]
    snake_case       : Any    =       [tf_models[t] for t in all_models]
    snake_case       : Dict    =       [flax_models[t] for t in all_models]

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

    snake_case       : Any    =       [processors[t] for t in all_models]

    return pd.DataFrame(lowercase					)







def     __lowerCAmelCase		(		lowercase    :    List[Any]					)      ->		Tuple:


    """simple docstring"""





    snake_case       : str    =       [
        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:
        snake_case       : Dict    =       [model_mapping, F'TF_{model_mapping}', F'FLAX_{model_mapping}']
        snake_case       : List[str]    =       [auto_class, F'TF_{auto_class}', F'Flax_{auto_class}']
        # Loop through all three frameworks
        for module, cls, mapping in zip(lowercase		,				lowercase		,				lowercase					):
            # The type of pipeline may not exist in this framework
            if not hasattr(lowercase		,				lowercase					):
                continue
            # First extract all model_names
            snake_case       : int    =       []
            for name in getattr(lowercase		,				lowercase					).values():
                if isinstance(lowercase		,				lowercase					):
                    model_names.append(lowercase					)
                else:
                    model_names.extend(list(lowercase					)					)

            # 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     __lowerCAmelCase		(		lowercase    :    Dict		,				lowercase    :    List[Any]					)      ->		Any:


    """simple docstring"""





    snake_case       : Union[str, Any]    =       get_frameworks_table()
    snake_case       : str    =       Dataset.from_pandas(lowercase					)

    snake_case       : Any    =       hf_hub_download(
        "huggingface/transformers-metadata"		,				"pipeline_tags.json"		,				repo_type="dataset"		,				token=lowercase					)
    snake_case       : Dict    =       Dataset.from_json(lowercase					)
    snake_case       : Tuple    =       {
        tags_dataset[i]["model_class"]: (tags_dataset[i]["pipeline_tag"], tags_dataset[i]["auto_class"])
        for i in range(len(lowercase					)					)
    }
    snake_case       : List[Any]    =       update_pipeline_and_auto_class_table(lowercase					)

    # Sort the model classes to avoid some nondeterministic updates to create false update commits.
    snake_case       : Dict    =       sorted(table.keys()					)
    snake_case       : int    =       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],
        }					)
    snake_case       : Optional[Any]    =       Dataset.from_pandas(lowercase					)

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

        if commit_sha is not None:
            snake_case       : Optional[int]    =       (
                F'Update with commit {commit_sha}\n\nSee: '
                F'https://github.com/huggingface/transformers/commit/{commit_sha}'
            )
        else:
            snake_case       : List[Any]    =       "Update"

        upload_folder(
            repo_id="huggingface/transformers-metadata"		,				folder_path=lowercase		,				repo_type="dataset"		,				token=lowercase		,				commit_message=lowercase		,				)







def     __lowerCAmelCase		(		)      ->		List[str]:


    """simple docstring"""





    snake_case       : Optional[Any]    =       {tag: cls for tag, _, cls in PIPELINE_TAGS_AND_AUTO_MODELS}
    snake_case       : Optional[Any]    =       transformers_module.pipelines.SUPPORTED_TASKS
    snake_case       : Tuple    =       []
    for key in pipeline_tasks:
        if key not in in_table:
            snake_case       : List[Any]    =       pipeline_tasks[key]["pt"]
            if isinstance(lowercase		,				(list, tuple)					):
                snake_case       : Tuple    =       model[0]
            snake_case       : Optional[int]    =       model.__name__
            if model not in in_table.values():
                missing.append(lowercase					)

    if len(lowercase					) > 0:
        snake_case       : Tuple    =       ", ".join(lowercase					)
        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__":
  __snake_case       =       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.""")
  __snake_case       =       parser.parse_args()

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









"""simple docstring"""



import argparse
import json

import requests
import timm
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import AutoImageProcessor, SwinConfig, SwinForImageClassification







def     __lowerCAmelCase		(		lowercase    :    Optional[int]					)      ->		Tuple:


    """simple docstring"""





    snake_case       : Optional[Any]    =       SwinConfig()
    snake_case       : Union[str, Any]    =       swin_name.split("_"					)

    snake_case       : str    =       name_split[1]
    snake_case       : Tuple    =       int(name_split[4]					)
    snake_case       : str    =       int(name_split[3][-1]					)

    if model_size == "tiny":
        snake_case       : str    =       96
        snake_case       : List[Any]    =       (2, 2, 6, 2)
        snake_case       : Union[str, Any]    =       (3, 6, 12, 24)
    elif model_size == "small":
        snake_case       : int    =       96
        snake_case       : Optional[Any]    =       (2, 2, 18, 2)
        snake_case       : Any    =       (3, 6, 12, 24)
    elif model_size == "base":
        snake_case       : List[str]    =       128
        snake_case       : int    =       (2, 2, 18, 2)
        snake_case       : Any    =       (4, 8, 16, 32)
    else:
        snake_case       : List[str]    =       192
        snake_case       : str    =       (2, 2, 18, 2)
        snake_case       : Any    =       (6, 12, 24, 48)

    if "in22k" in swin_name:
        snake_case       : List[Any]    =       2_1841
    else:
        snake_case       : str    =       1000
        snake_case       : Any    =       "huggingface/label-files"
        snake_case       : str    =       "imagenet-1k-id2label.json"
        snake_case       : Optional[int]    =       json.load(open(hf_hub_download(lowercase		,				lowercase		,				repo_type="dataset"					)		,				"r"					)					)
        snake_case       : List[Any]    =       {int(lowercase					): v for k, v in idalabel.items()}
        snake_case       : Optional[int]    =       idalabel
        snake_case       : Optional[int]    =       {v: k for k, v in idalabel.items()}

    snake_case       : Tuple    =       img_size
    snake_case       : Any    =       num_classes
    snake_case       : str    =       embed_dim
    snake_case       : Any    =       depths
    snake_case       : Dict    =       num_heads
    snake_case       : List[Any]    =       window_size

    return config







def     __lowerCAmelCase		(		lowercase    :    Optional[Any]					)      ->		Any:


    """simple docstring"""





    if "patch_embed.proj" in name:
        snake_case       : int    =       name.replace("patch_embed.proj"		,				"embeddings.patch_embeddings.projection"					)
    if "patch_embed.norm" in name:
        snake_case       : Union[str, Any]    =       name.replace("patch_embed.norm"		,				"embeddings.norm"					)
    if "layers" in name:
        snake_case       : Optional[Any]    =       "encoder." + name
    if "attn.proj" in name:
        snake_case       : List[Any]    =       name.replace("attn.proj"		,				"attention.output.dense"					)
    if "attn" in name:
        snake_case       : List[str]    =       name.replace("attn"		,				"attention.self"					)
    if "norm1" in name:
        snake_case       : List[str]    =       name.replace("norm1"		,				"layernorm_before"					)
    if "norm2" in name:
        snake_case       : Optional[int]    =       name.replace("norm2"		,				"layernorm_after"					)
    if "mlp.fc1" in name:
        snake_case       : int    =       name.replace("mlp.fc1"		,				"intermediate.dense"					)
    if "mlp.fc2" in name:
        snake_case       : Tuple    =       name.replace("mlp.fc2"		,				"output.dense"					)

    if name == "norm.weight":
        snake_case       : Optional[int]    =       "layernorm.weight"
    if name == "norm.bias":
        snake_case       : Tuple    =       "layernorm.bias"

    if "head" in name:
        snake_case       : Any    =       name.replace("head"		,				"classifier"					)
    else:
        snake_case       : Tuple    =       "swin." + name

    return name







def     __lowerCAmelCase		(		lowercase    :    List[Any]		,				lowercase    :    List[Any]					)      ->		Any:


    """simple docstring"""





    for key in orig_state_dict.copy().keys():
        snake_case       : int    =       orig_state_dict.pop(lowercase					)

        if "mask" in key:
            continue
        elif "qkv" in key:
            snake_case       : List[Any]    =       key.split("."					)
            snake_case       : Dict    =       int(key_split[1]					)
            snake_case       : Dict    =       int(key_split[3]					)
            snake_case       : str    =       model.swin.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size

            if "weight" in key:
                snake_case       : Dict    =       val[:dim, :]
                snake_case       : Tuple    =       val[
                    dim : dim * 2, :
                ]
                snake_case       : List[str]    =       val[-dim:, :]
            else:
                snake_case       : Dict    =       val[
                    :dim
                ]
                snake_case       : List[str]    =       val[
                    dim : dim * 2
                ]
                snake_case       : Union[str, Any]    =       val[
                    -dim:
                ]
        else:
            snake_case       : Any    =       val

    return orig_state_dict







def     __lowerCAmelCase		(		lowercase    :    Tuple		,				lowercase    :    Any					)      ->		Dict:


    """simple docstring"""





    snake_case       : str    =       timm.create_model(lowercase		,				pretrained=lowercase					)
    timm_model.eval()

    snake_case       : int    =       get_swin_config(lowercase					)
    snake_case       : Optional[int]    =       SwinForImageClassification(lowercase					)
    model.eval()

    snake_case       : int    =       convert_state_dict(timm_model.state_dict()		,				lowercase					)
    model.load_state_dict(lowercase					)

    snake_case       : Dict    =       "http://images.cocodataset.org/val2017/000000039769.jpg"

    snake_case       : Tuple    =       AutoImageProcessor.from_pretrained("microsoft/{}".format(swin_name.replace("_"		,				"-"					)					)					)
    snake_case       : Optional[int]    =       Image.open(requests.get(lowercase		,				stream=lowercase					).raw					)
    snake_case       : Dict    =       image_processor(images=lowercase		,				return_tensors="pt"					)

    snake_case       : str    =       timm_model(inputs["pixel_values"]					)
    snake_case       : Tuple    =       model(**lowercase					).logits

    assert torch.allclose(lowercase		,				lowercase		,				atol=1e-3					)

    print(F'Saving model {swin_name} to {pytorch_dump_folder_path}'					)
    model.save_pretrained(lowercase					)

    print(F'Saving image processor to {pytorch_dump_folder_path}'					)
    image_processor.save_pretrained(lowercase					)


if __name__ == "__main__":
  __snake_case       =       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--swin_name""",
      default="""swin_tiny_patch4_window7_224""",
      type=str,
      help="""Name of the Swin timm model you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )

  __snake_case       =       parser.parse_args()
  convert_swin_checkpoint(args.swin_name, args.pytorch_dump_folder_path)





