codingutf8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license utility that updates the metadata of the transformers library in the repository huggingfacetransformersmetadata usage for an update as used by the github action updatemetadata bash python utilsupdatemetadata py token token commitsha commitsha usage to check all pipelines are properly defined in the constant pipelinetagsandautomodels of this script so that new pipelines are properly added as metadata as used in make repoconsistency bash python utilsupdatemetadata py checkonly all paths are set with the intent you should run this script from the root of the repo with the command python utilsupdatemetadata py this is to make sure the transformers module imported is the one in the repo regexes that match tfflaxpt model names will match any tf or flax model too so need to be in an else branch afterthe two previous regexes fill this with tuples pipelinetag modelmapping automodel split a camelcased name into words args identifier str the camelcased name to parse returns liststr the list of words in the identifier as seprated by capital letters example py camelcasesplitcamelcasedclass camel cased class regex thanks to https stackoverflow comquestions29916065howtodocamelcasesplitinpython generates a dataframe containing the supported auto classes for each model type using the content of the auto modules dictionary model names to config dictionaries flagging if each model prefix has a backend in pttfflax let s lookup through all transformers object once and find if models are supported by a given backend try again after removing the last word in the name now let s find the right processing class for each model in order we check if there is a processor then a tokenizer then a featureextractor then an imageprocessor default to autotokenizer if a model has nothing for backward compatibility update the table maping models to pipelines and auto classes without removing old keys if they don t exist anymore args table dictstr tuplestr str the existing table mapping model names to a tuple containing the pipeline tag and the autoclass name with which they should be used returns dictstr tuplestr str the updated table in the same format loop through all three frameworks the type of pipeline may not exist in this framework first extract all modelnames add pipeline tag and auto model class for those models update the metadata for the transformers repo in huggingfacetransformersmetadata args token str a valid token giving write access to huggingfacetransformersmetadata commitsha str the commit sha on transformers corresponding to this update sort the model classes to avoid some nondeterministic updates to create false update commits check all pipeline tags are properly defined in the pipelinetagsandautomodels constant of this script coding utf 8 2021 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license utility that updates the metadata of the transformers library in the repository huggingface transformers metadata usage for an update as used by the github action update_metadata bash python utils update_metadata py token token commit_sha commit_sha usage to check all pipelines are properly defined in the constant pipeline_tags_and_auto_models of this script so that new pipelines are properly added as metadata as used in make repo consistency bash python utils update_metadata py check only 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 this is to make sure the transformers module imported is the one in the repo regexes that match tf flax pt model names will match any tf or flax model too so need to be in an else branch afterthe two previous regexes fill this with tuples pipeline_tag model_mapping auto_model split a camel cased name into words args identifier str the camel cased name to parse returns list str the list of words in the identifier as seprated by capital letters example py camel_case_split camelcasedclass camel cased class regex thanks to https stackoverflow com questions 29916065 how to do camelcase split in python generates a dataframe containing the supported auto classes for each model type using the content of the auto modules dictionary model names to config dictionaries flagging if each model prefix has a backend in pt tf flax let s lookup through all transformers object once and find if models are supported by a given backend try again after removing the last word in the name now let s find the right processing class for each model in order we check if there is a processor then a tokenizer then a featureextractor then an imageprocessor default to autotokenizer if a model has nothing for backward compatibility update the table maping models to pipelines and auto classes without removing old keys if they don t exist anymore args table dict str tuple str str the existing table mapping model names to a tuple containing the pipeline tag and the auto class name with which they should be used returns dict str tuple str str the updated table in the same format loop through all three frameworks the type of pipeline may not exist in this framework first extract all model_names add pipeline tag and auto model class for those models update the metadata for the transformers repo in huggingface transformers metadata args token str a valid token giving write access to huggingface transformers metadata commit_sha str the commit sha on transformers corresponding to this update sort the model classes to avoid some nondeterministic updates to create false update commits check all pipeline tags are properly defined in the pipeline_tags_and_auto_models constant of this script
import argparse
import collections
import os
import re
import tempfile
from typing import Dict, List, Tuple

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

from transformers.utils import direct_transformers_import




TRANSFORMERS_PATH = "src/transformers"



transformers_module = direct_transformers_import(TRANSFORMERS_PATH)



_re_tf_models = re.compile(r"TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
_re_flax_models = re.compile(r"Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")

_re_pt_models = re.compile(r"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")



PIPELINE_TAGS_AND_AUTO_MODELS = [
    ("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"),
    ("image-to-image", "MODEL_FOR_IMAGE_TO_IMAGE_MAPPING_NAMES", "AutoModelForImageToImage"),
    ("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"),
    ("text-to-audio", "MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES", "AutoModelForTextToSpectrogram"),
    ("text-to-audio", "MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES", "AutoModelForTextToWaveform"),
]


def camel_case_split(identifier: str) -> List[str]:
    
    
    matches = re.finditer(".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)", identifier)
    return [m.group(0) for m in matches]


def get_frameworks_table() -> pd.DataFrame:
    
    
    config_maping_names = transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
    model_prefix_to_model_type = {
        config.replace("Config", ""): model_type for model_type, config in config_maping_names.items()
    }

    
    pt_models = collections.defaultdict(bool)
    tf_models = collections.defaultdict(bool)
    flax_models = collections.defaultdict(bool)

    
    for attr_name in dir(transformers_module):
        lookup_dict = None
        if _re_tf_models.match(attr_name) is not None:
            lookup_dict = tf_models
            attr_name = _re_tf_models.match(attr_name).groups()[0]
        elif _re_flax_models.match(attr_name) is not None:
            lookup_dict = flax_models
            attr_name = _re_flax_models.match(attr_name).groups()[0]
        elif _re_pt_models.match(attr_name) is not None:
            lookup_dict = pt_models
            attr_name = _re_pt_models.match(attr_name).groups()[0]

        if lookup_dict is not None:
            while len(attr_name) > 0:
                if attr_name in model_prefix_to_model_type:
                    lookup_dict[model_prefix_to_model_type[attr_name]] = True
                    break
                
                attr_name = "".join(camel_case_split(attr_name)[:-1])

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

    data = {"model_type": all_models}
    data["pytorch"] = [pt_models[t] for t in all_models]
    data["tensorflow"] = [tf_models[t] for t in all_models]
    data["flax"] = [flax_models[t] for t in all_models]

    
    
    processors = {}
    for t in all_models:
        if t in transformers_module.models.auto.processing_auto.PROCESSOR_MAPPING_NAMES:
            processors[t] = "AutoProcessor"
        elif t in transformers_module.models.auto.tokenization_auto.TOKENIZER_MAPPING_NAMES:
            processors[t] = "AutoTokenizer"
        elif t in transformers_module.models.auto.image_processing_auto.IMAGE_PROCESSOR_MAPPING_NAMES:
            processors[t] = "AutoImageProcessor"
        elif t in transformers_module.models.auto.feature_extraction_auto.FEATURE_EXTRACTOR_MAPPING_NAMES:
            processors[t] = "AutoFeatureExtractor"
        else:
            
            processors[t] = "AutoTokenizer"

    data["processor"] = [processors[t] for t in all_models]

    return pd.DataFrame(data)


def update_pipeline_and_auto_class_table(table: Dict[str, Tuple[str, str]]) -> Dict[str, Tuple[str, str]]:
    
    auto_modules = [
        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:
        model_mappings = [model_mapping, f"TF_{model_mapping}", f"FLAX_{model_mapping}"]
        auto_classes = [auto_class, f"TF_{auto_class}", f"Flax_{auto_class}"]
        
        for module, cls, mapping in zip(auto_modules, auto_classes, model_mappings):
            
            if not hasattr(module, mapping):
                continue
            
            model_names = []
            for name in getattr(module, mapping).values():
                if isinstance(name, str):
                    model_names.append(name)
                else:
                    model_names.extend(list(name))

            
            table.update({model_name: (pipeline_tag, cls) for model_name in model_names})

    return table


def update_metadata(token: str, commit_sha: str):
    
    frameworks_table = get_frameworks_table()
    frameworks_dataset = Dataset.from_pandas(frameworks_table)

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

    
    model_classes = sorted(table.keys())
    tags_table = 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],
        }
    )
    tags_dataset = Dataset.from_pandas(tags_table)

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

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

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


def check_pipeline_tags():
    
    in_table = {tag: cls for tag, _, cls in PIPELINE_TAGS_AND_AUTO_MODELS}
    pipeline_tasks = transformers_module.pipelines.SUPPORTED_TASKS
    missing = []
    for key in pipeline_tasks:
        if key not in in_table:
            model = pipeline_tasks[key]["pt"]
            if isinstance(model, (list, tuple)):
                model = model[0]
            model = model.__name__
            if model not in in_table.values():
                missing.append(key)

    if len(missing) > 0:
        msg = ", ".join(missing)
        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__":
    parser = 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.")
    args = parser.parse_args()

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