onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -0,0 +1,422 @@
+import copy
+import functools
+import json
+import os
+import pprint
+import sys
+from typing import Any, Dict, List, Optional, Union
+import transformers
+from huggingface_hub import HfApi, model_info, hf_hub_download, list_repo_files
+from ...helpers.config_helper import update_config
+from . import hub_data_cached_configs
+from .hub_data import __date__, __data_tasks__, load_architecture_task, __data_arch_values__
+
+
+@functools.cache
+def get_architecture_default_values(architecture: str):
+    """
+    The configuration may miss information to build the dummy inputs.
+    This information returns the missing pieces.
+    """
+    assert architecture in __data_arch_values__, (
+        f"No known default values for {architecture!r}, "
+        f"expecting one architecture in {', '.join(sorted(__data_arch_values__))}"
+    )
+    return __data_arch_values__[architecture]
+
+
+@functools.cache
+def _retrieve_cached_configurations() -> Dict[str, transformers.PretrainedConfig]:
+    res = {}
+    for k, v in hub_data_cached_configs.__dict__.items():
+        if k.startswith("_ccached_"):
+            doc = v.__doc__
+            res[doc] = v
+    return res
+
+
+def get_cached_configuration(
+    name: str, exc: bool = False, **kwargs
+) -> Optional[transformers.PretrainedConfig]:
+    """
+    Returns cached configuration to avoid having to many accesses to internet.
+    It returns None if not Cache. The list of cached models follows.
+    If *exc* is True or if environment variable ``NOHTTP`` is defined,
+    the function raises an exception if *name* is not found.
+
+    .. runpython::
+
+        import pprint
+        from onnx_diagnostic.torch_models.hghub.hub_api import _retrieve_cached_configurations
+
+        configs = _retrieve_cached_configurations()
+        pprint.pprint(sorted(configs))
+    """
+    cached = _retrieve_cached_configurations()
+    assert cached, "no cached configuration, which is weird"
+    if name in cached:
+        conf = cached[name]()
+        if kwargs:
+            conf = copy.deepcopy(conf)
+            update_config(conf, kwargs)
+        return conf
+    assert not exc and not os.environ.get("NOHTTP", ""), (
+        f"Unable to find {name!r} (exc={exc}, "
+        f"NOHTTP={os.environ.get('NOHTTP', '')!r}) "
+        f"in {pprint.pformat(sorted(cached))}"
+    )
+    return None
+
+
+def get_pretrained_config(
+    model_id: str,
+    trust_remote_code: bool = True,
+    use_preinstalled: bool = True,
+    subfolder: Optional[str] = None,
+    use_only_preinstalled: bool = False,
+    **kwargs,
+) -> Any:
+    """
+    Returns the config for a model_id.
+
+    :param model_id: model id
+    :param trust_remote_code: trust_remote_code,
+        see :meth:`transformers.AutoConfig.from_pretrained`
+    :param use_preinstalled: if use_preinstalled, uses this version to avoid
+        accessing the network, if available, it is returned by
+        :func:`get_cached_configuration`, the cached list is mostly for
+        unit tests
+    :param subfolder: subfolder for the given model id
+    :param use_only_preinstalled: if True, raises an exception if not preinstalled
+    :param kwargs: additional kwargs
+    :return: a configuration
+    """
+    if use_preinstalled:
+        conf = get_cached_configuration(
+            model_id, exc=use_only_preinstalled, subfolder=subfolder, **kwargs
+        )
+        if conf is not None:
+            return conf
+    assert not use_only_preinstalled, (
+        f"Inconsistencies: use_only_preinstalled={use_only_preinstalled}, "
+        f"use_preinstalled={use_preinstalled!r}"
+    )
+    if subfolder:
+        try:
+            return transformers.AutoConfig.from_pretrained(
+                model_id, trust_remote_code=trust_remote_code, subfolder=subfolder, **kwargs
+            )
+        except ValueError:
+            # Then we try to download it.
+            config = hf_hub_download(
+                model_id, filename="config.json", subfolder=subfolder, **kwargs
+            )
+            try:
+                return transformers.AutoConfig.from_pretrained(
+                    config, trust_remote_code=trust_remote_code, **kwargs
+                )
+            except ValueError:
+                # Diffusers uses a dictionayr.
+                with open(config, "r") as f:
+                    return json.load(f)
+    return transformers.AutoConfig.from_pretrained(
+        model_id, trust_remote_code=trust_remote_code, **kwargs
+    )
+
+
+def get_model_info(model_id) -> Any:
+    """Returns the model info for a model_id."""
+    return model_info(model_id)
+
+
+def _guess_task_from_config(config: Any) -> Optional[str]:
+    """Tries to infer a task from the configuration."""
+    if hasattr(config, "bbox_loss_coefficient") and hasattr(config, "giou_loss_coefficient"):
+        return "object-detection"
+    if hasattr(config, "architecture") and config.architecture:
+        return task_from_arch(config.architecture)
+    return None
+
+
+@functools.cache
+def task_from_arch(
+    arch: str,
+    default_value: Optional[str] = None,
+    model_id: Optional[str] = None,
+    subfolder: Optional[str] = None,
+) -> str:
+    """
+    This function relies on stored information. That information needs to be refresh.
+
+    :param arch: architecture name
+    :param default_value: default value in case the task cannot be determined
+    :param model_id: unused unless the architecture does not help.
+    :param subfolder: subfolder
+    :return: task
+
+    .. runpython::
+
+        from onnx_diagnostic.torch_models.hghub.hub_data import __date__
+        print("last refresh", __date__)
+
+    List of supported architectures, see
+    :func:`load_architecture_task
+    <onnx_diagnostic.torch_models.hghub.hub_data.load_architecture_task>`.
+    """
+    data = load_architecture_task()
+    if arch not in data and model_id:
+        # Let's try with the model id.
+        return task_from_id(model_id, subfolder=subfolder)
+    if default_value is not None:
+        return data.get(arch, default_value)
+    assert arch in data, (
+        f"Architecture {arch!r} is unknown, last refresh in {__date__}. "
+        f"``onnx_diagnostic.torch_models.hghub.hub_data.__data_arch__`` "
+        f"needs to be updated (model_id={(model_id or '?')!r})."
+    )
+    return data[arch]
+
+
+def _trygetattr(config, attname):
+    try:
+        return getattr(config, attname)
+    except AttributeError:
+        return None
+
+
+def architecture_from_config(config) -> Optional[str]:
+    """Guesses the architecture (class) of the model described by this config."""
+    if isinstance(config, dict):
+        if "_class_name" in config:
+            return config["_class_name"]
+        if "architecture" in config:
+            return config["architecture"]
+        if config.get("architectures", []):
+            return config["architectures"][0]
+    if hasattr(config, "_class_name"):
+        return config._class_name
+    if hasattr(config, "architecture"):
+        return config.architecture
+    if hasattr(config, "architectures") and config.architectures:
+        return config.architectures[0]
+    if hasattr(config, "__dict__"):
+        if "_class_name" in config.__dict__:
+            return config.__dict__["_class_name"]
+        if "architecture" in config.__dict__:
+            return config.__dict__["architecture"]
+        if config.__dict__.get("architectures", []):
+            return config.__dict__["architectures"][0]
+    return None
+
+
+def find_package_source(config) -> Optional[str]:
+    """Guesses the package the class models from."""
+    if isinstance(config, dict):
+        if "_diffusers_version" in config:
+            return "diffusers"
+    if hasattr(config, "_diffusers_version"):
+        return "diffusers"
+    if hasattr(config, "__dict__"):
+        if "_diffusers_version" in config.__dict__:
+            return "diffusers"
+    return "transformers"
+
+
+def task_from_id(
+    model_id: str,
+    default_value: Optional[str] = None,
+    pretrained: bool = False,
+    fall_back_to_pretrained: bool = True,
+    subfolder: Optional[str] = None,
+) -> str:
+    """
+    Returns the task attached to a model id.
+
+    :param model_id: model id
+    :param default_value: if specified, the function returns this value
+        if the task cannot be determined
+    :param pretrained: uses the config
+    :param fall_back_to_pretrained: falls back to pretrained config
+    :param subfolder: subfolder
+    :return: task
+    """
+    if not pretrained:
+        try:
+            transformers.pipelines.get_task(model_id)
+        except RuntimeError:
+            if not fall_back_to_pretrained:
+                raise
+    config = get_pretrained_config(model_id, subfolder=subfolder)
+    tag = _trygetattr(config, "pipeline_tag")
+    if tag is not None:
+        return tag
+
+    guess = _guess_task_from_config(config)
+    if guess is not None:
+        return guess
+    data = load_architecture_task()
+    if subfolder:
+        full_id = f"{model_id}//{subfolder}"
+        if full_id in data:
+            return data[full_id]
+    if model_id in data:
+        return data[model_id]
+    arch = architecture_from_config(config)
+    if arch is None:
+        if model_id.startswith("google/bert_"):
+            return "fill-mask"
+    assert arch is not None, (
+        f"Cannot return the task of {model_id!r}, pipeline_tag is not setup, "
+        f"config={config}. The task can be added in "
+        f"``onnx_diagnostic.torch_models.hghub.hub_data.__data_arch__``."
+    )
+    return task_from_arch(arch, default_value=default_value)
+
+
+def task_from_tags(tags: Union[str, List[str]]) -> str:
+    """
+    Guesses the task from the list of tags.
+    If given by a string, ``|`` should be the separator.
+    """
+    if isinstance(tags, str):
+        tags = tags.split("|")
+    stags = set(tags)
+    for task in __data_tasks__:
+        if task in stags:
+            return task
+    raise ValueError(f"Unable to guess the task from tags={tags!r}")
+
+
+def enumerate_model_list(
+    n: int = 50,
+    pipeline_tag: Optional[str] = None,
+    search: Optional[str] = None,
+    dump: Optional[str] = None,
+    filter: Optional[Union[str, List[str]]] = None,
+    verbose: int = 0,
+):
+    """
+    Enumerates models coming from :epkg:`huggingface_hub`.
+
+    :param n: number of models to retrieve (-1 for all)
+    :param pipeline_tag: see :meth:`huggingface_hub.HfApi.list_models`
+    :param search: see :meth:`huggingface_hub.HfApi.list_models`
+    :param filter: see :meth:`huggingface_hub.HfApi.list_models`
+    :param dump: dumps the result in this csv file
+    :param verbose: show progress
+    """
+    api = HfApi()
+    models = api.list_models(
+        pipeline_tag=pipeline_tag,
+        search=search,
+        full=True,
+        filter=filter,
+        limit=n if n > 0 else None,
+    )
+    seen = 0
+    found = 0
+
+    if dump:
+        with open(dump, "w") as f:
+            f.write(
+                ",".join(
+                    [
+                        "id",
+                        "model_name",
+                        "author",
+                        "created_at",
+                        "last_modified",
+                        "downloads",
+                        "downloads_all_time",
+                        "likes",
+                        "trending_score",
+                        "private",
+                        "gated",
+                        "tags",
+                    ]
+                )
+            )
+            f.write("\n")
+
+    for m in models:
+        seen += 1  # noqa: SIM113
+        if verbose and seen % 1000 == 0:
+            print(f"[enumerate_model_list] {seen} models, found {found}")
+        if verbose > 1:
+            print(
+                f"[enumerate_model_list]     id={m.id!r}, "
+                f"library={m.library_name!r}, task={m.task!r}"
+            )
+        with open(dump, "a") as f:  # type: ignore
+            f.write(
+                ",".join(
+                    map(
+                        str,
+                        [
+                            m.id,
+                            getattr(m, "model_name", "") or "",
+                            m.author or "",
+                            str(m.created_at or "").split(" ")[0],
+                            str(m.last_modified or "").split(" ")[0],
+                            m.downloads or "",
+                            m.downloads_all_time or "",
+                            m.likes or "",
+                            m.trending_score or "",
+                            m.private or "",
+                            m.gated or "",
+                            (
+                                ("|".join(m.tags)).replace(",", "_").replace(" ", "_")
+                                if m.tags
+                                else ""
+                            ),
+                        ],
+                    )
+                )
+            )
+            f.write("\n")
+        yield m
+        found += 1  # noqa: SIM113
+        if n >= 0:
+            n -= 1
+            if n == 0:
+                break
+
+
+def download_code_modelid(
+    model_id: str, verbose: int = 0, add_path_to_sys_path: bool = True
+) -> List[str]:
+    """
+    Downloads the code for a given model id.
+
+    :param model_id: model id
+    :param verbose: verbosity
+    :param add_path_to_sys_path: add folder where the files are downloaded to sys.path
+    :return: list of downloaded files
+    """
+    if verbose:
+        print(f"[download_code_modelid] retrieve file list for {model_id!r}")
+    files = list_repo_files(model_id)
+    pyfiles = [name for name in files if os.path.splitext(name)[-1] == ".py"]
+    if verbose:
+        print(f"[download_code_modelid] python files {pyfiles}")
+    absfiles = []
+    paths = set()
+    for i, name in enumerate(pyfiles):
+        if verbose:
+            print(f"[download_code_modelid] download file {i+1}/{len(pyfiles)}: {name!r}")
+        r = hf_hub_download(repo_id=model_id, filename=name)
+        p = os.path.split(r)[0]
+        paths.add(p)
+        absfiles.append(r)
+    if add_path_to_sys_path:
+        for p in paths:
+            init = os.path.join(p, "__init__.py")
+            if not os.path.exists(init):
+                with open(init, "w"):
+                    pass
+            if p in sys.path:
+                continue
+            if verbose:
+                print(f"[download_code_modelid] add {p!r} to 'sys.path'")
+            sys.path.insert(0, p)
+    return absfiles

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -0,0 +1,234 @@
+import io
+import functools
+import textwrap
+from typing import Dict, List
+
+__date__ = "2025-06-21"
+
+__data_arch_values__ = {"ResNetForImageClassification": dict(image_size=224)}
+
+__data_arch__ = textwrap.dedent(
+    """
+    architecture,task
+    ASTModel,feature-extraction
+    AutoencoderKL,image-to-image
+    AlbertModel,feature-extraction
+    BeitForImageClassification,image-classification
+    BartForConditionalGeneration,summarization
+    BartModel,feature-extraction
+    BertForMaskedLM,fill-mask
+    BertForSequenceClassification,text-classification
+    BertModel,sentence-similarity
+    BigBirdModel,feature-extraction
+    BlenderbotModel,feature-extraction
+    BloomModel,feature-extraction
+    CLIPModel,zero-shot-image-classification
+    CLIPTextModel,feature-extraction
+    CLIPVisionModel,feature-extraction
+    CamembertModel,feature-extraction
+    CodeGenModel,feature-extraction
+    ConvBertModel,feature-extraction
+    ConvNextForImageClassification,image-classification
+    ConvNextV2Model,image-feature-extraction
+    CosmosTransformer3DModel,image-to-video
+    CvtModel,feature-extraction
+    DPTModel,image-feature-extraction
+    Data2VecAudioModel,feature-extraction
+    Data2VecTextModel,feature-extraction
+    Data2VecVisionModel,image-feature-extraction
+    DebertaModel,feature-extraction
+    DebertaV2Model,feature-extraction
+    DecisionTransformerModel,reinforcement-learning
+    DeepseekV3ForCausalLM,text-generation
+    DeiTModel,image-feature-extraction
+    DetrModel,image-feature-extraction
+    Dinov2Model,image-feature-extraction
+    DistilBertForSequenceClassification,text-classification
+    DistilBertModel,feature-extraction
+    DonutSwinModel,feature-extraction
+    ElectraModel,feature-extraction
+    EsmModel,feature-extraction
+    FalconMambaForCausalLM,text-generation
+    GLPNModel,image-feature-extraction
+    GPT2LMHeadModel,text-generation
+    GPTBigCodeModel,feature-extraction
+    GPTJModel,feature-extraction
+    GPTNeoModel,feature-extraction
+    GPTNeoXForCausalLM,text-generation
+    GptOssForCausalLM,text-generation
+    GemmaForCausalLM,text-generation
+    Gemma2ForCausalLM,text-generation
+    Gemma3ForConditionalGeneration,image-text-to-text
+    Gemma3ForCausalLM,text-generation
+    Glm4vMoeForConditionalGeneration,image-text-to-text
+    GraniteForCausalLM,text-generation
+    GroupViTModel,feature-extraction
+    HieraForImageClassification,image-classification
+    HubertModel,feature-extraction
+    IBertModel,feature-extraction
+    IdeficsForVisionText2Text,image-text-to-text
+    ImageGPTModel,image-feature-extraction
+    LayoutLMModel,feature-extraction
+    LayoutLMv3Model,feature-extraction
+    LevitModel,image-feature-extraction
+    LiltModel,feature-extraction
+    LlamaForCausalLM,text-generation
+    LongT5Model,feature-extraction
+    LongformerModel,feature-extraction
+    MCTCTModel,feature-extraction
+    MPNetForMaskedLM,sentence-similarity
+    MPNetModel,feature-extraction
+    MT5Model,feature-extraction
+    MarianMTModel,text2text-generation
+    MarkupLMModel,feature-extraction
+    MaskFormerForInstanceSegmentation,image-segmentation
+    MegatronBertModel,feature-extraction
+    MgpstrForSceneTextRecognition,feature-extraction
+    MistralForCausalLM,text-generation
+    MobileBertModel,feature-extraction
+    MobileNetV1Model,image-feature-extraction
+    MobileNetV2Model,image-feature-extraction
+    mobilenetv3_small_100,image-classification
+    MobileViTForImageClassification,image-classification
+    ModernBertForMaskedLM,fill-mask
+    Phi4MMForCausalLM,MoE
+    MoonshineForConditionalGeneration,automatic-speech-recognition
+    MptForCausalLM,text-generation
+    MusicgenForConditionalGeneration,text-to-audio
+    NystromformerModel,feature-extraction
+    OPTModel,feature-extraction
+    Olmo2ForCausalLM,text-generation
+    OlmoForCausalLM,text-generation
+    OwlViTModel,feature-extraction
+    Owlv2Model,feature-extraction
+    PatchTSMixerForPrediction,no-pipeline-tag
+    PatchTSTForPrediction,no-pipeline-tag
+    PegasusModel,feature-extraction
+    Phi3ForCausalLM,text-generation
+    PhiForCausalLM,text-generation
+    PhiMoEForCausalLM,text-generation
+    Pix2StructForConditionalGeneration,image-to-text
+    PLBartForConditionalGeneration,text2text-generation
+    PoolFormerModel,image-feature-extraction
+    PvtForImageClassification,image-classification
+    Qwen2ForCausalLM,text-generation
+    Qwen2_5_VLForConditionalGeneration,image-text-to-text
+    Qwen3MoeForCausalLM,text-generation
+    RTDetrForObjectDetection,object-detection
+    RegNetModel,image-feature-extraction
+    RemBertModel,feature-extraction
+    ResNetForImageClassification,image-classification
+    RoFormerModel,feature-extraction
+    RobertaForMaskedLM,sentence-similarity
+    RobertaModel,feature-extraction
+    RtDetrV2ForObjectDetection,object-detection
+    SEWDModel,feature-extraction
+    SEWModel,feature-extraction
+    SamModel,mask-generation
+    SegformerModel,image-feature-extraction
+    SiglipModel,zero-shot-image-classification
+    SiglipVisionModel,image-feature-extraction
+    Speech2TextModel,feature-extraction
+    SpeechT5ForTextToSpeech,text-to-audio
+    SplinterModel,feature-extraction
+    SqueezeBertModel,feature-extraction
+    Swin2SRModel,image-feature-extraction
+    SwinModel,image-feature-extraction
+    Swinv2Model,image-feature-extraction
+    T5ForConditionalGeneration,text2text-generation
+    TableTransformerModel,image-feature-extraction
+    TableTransformerForObjectDetection,object-detection
+    UNet2DConditionModel,text-to-image
+    UniSpeechForSequenceClassification,audio-classification
+    ViTForImageClassification,image-classification
+    ViTMAEModel,image-feature-extraction
+    ViTMSNForImageClassification,image-classification
+    VisionEncoderDecoderModel,document-question-answering
+    VitPoseForPoseEstimation,keypoint-detection
+    VitsModel,text-to-audio
+    Wav2Vec2ConformerForCTC,automatic-speech-recognition
+    Wav2Vec2Model,feature-extraction
+    WhisperForConditionalGeneration,automatic-speech-recognition
+    XLMModel,feature-extraction
+    XLMRobertaForCausalLM,text-generation
+    XLMRobertaForMaskedLM,fill-mask
+    XLMRobertaModel,sentence-similarity
+    Wav2Vec2ForCTC,automatic-speech-recognition
+    YolosForObjectDetection,object-detection
+    YolosModel,image-feature-extraction
+    Alibaba-NLP/gte-large-en-v1.5,sentence-similarity
+    emilyalsentzer/Bio_ClinicalBERT,fill-mask
+    nvidia/Cosmos-Predict2-2B-Video2World//transformer,image-to-video"""
+)
+
+__data_tasks__ = [
+    "audio-classification",
+    "automatic-speech-recognition",
+    "document-question-answering",
+    "feature-extraction",
+    "fill-mask",
+    "image-classification",
+    "image-feature-extraction",
+    "image-segmentation",
+    "image-text-to-text",
+    "image-to-text",
+    "keypoint-detection",
+    "mask-generation",
+    "no-pipeline-tag",
+    "object-detection",
+    "reinforcement-learning",
+    "sentence-similarity",
+    "summarization",
+    "text-classification",
+    "text-generation",
+    "text-to-image",
+    "text-to-audio",
+    "text2text-generation",
+    "zero-shot-image-classification",
+]
+
+__models_testing__ = """
+hf-internal-testing/tiny-random-BeitForImageClassification
+hf-internal-testing/tiny-random-convnext
+fxmarty/tiny-random-GemmaForCausalLM
+hf-internal-testing/tiny-random-GPTNeoXForCausalLM
+hf-internal-testing/tiny-random-GraniteForCausalLM
+hf-internal-testing/tiny-random-HieraForImageClassification
+fxmarty/tiny-llama-fast-tokenizer
+sshleifer/tiny-marian-en-de
+hf-internal-testing/tiny-random-MaskFormerForInstanceSegmentation
+echarlaix/tiny-random-mistral
+hf-internal-testing/tiny-random-mobilevit
+hf-internal-testing/tiny-random-MoonshineForConditionalGeneration
+hf-internal-testing/tiny-random-OlmoForCausalLM
+hf-internal-testing/tiny-random-Olmo2ForCausalLM
+echarlaix/tiny-random-PhiForCausalLM
+Xenova/tiny-random-Phi3ForCausalLM
+fxmarty/pix2struct-tiny-random
+fxmarty/tiny-dummy-qwen2
+hf-internal-testing/tiny-random-ViTMSNForImageClassification
+hf-internal-testing/tiny-random-YolosModel
+hf-internal-testing/tiny-xlm-roberta
+HuggingFaceM4/tiny-random-idefics
+"""
+
+
+@functools.cache
+def load_models_testing() -> List[str]:
+    """Returns model ids for testing."""
+    return [_.strip() for _ in __models_testing__.split("\n") if _.strip()]
+
+
+@functools.cache
+def load_architecture_task() -> Dict[str, str]:
+    """
+    Returns a dictionary mapping architectures to tasks.
+
+    import pprint
+    from onnx_diagnostic.torch_models.hghub.hub_data import load_architecture_task
+    pprint.pprint(load_architecture_task())
+    """
+    import pandas
+
+    df = pandas.read_csv(io.StringIO(__data_arch__))
+    return dict(zip(list(df["architecture"]), list(df["task"])))