add tiny

Author: xadupre

_doc/api/torch_models/hghub/index.rst

@@ -7,6 +7,9 @@ onnx_diagnostic.torch_models.hghub
 
     hub_api
     hub_data
+    model_inputs
+
+.. autofunction:: onnx_diagnostic.torch_models.hghub.get_untrained_model_with_inputs
 
 .. automodule:: onnx_diagnostic.torch_models.hghub
     :members:

_doc/api/torch_models/hghub/model_inputs.rst

@@ -0,0 +1,8 @@
+
+onnx_diagnostic.torch_models.hghub.model_inputs
+===============================================
+
+.. automodule:: onnx_diagnostic.torch_models.hghub.model_inputs
+    :members:
+    :no-undoc-members:
+    :exclude-members: get_untrained_model_with_inputs

_doc/conf.py

@@ -232,5 +232,6 @@
         "arnir0/Tiny-LLM": "https://huggingface.co/arnir0/Tiny-LLM",
         "microsoft/phi-2": "https://huggingface.co/microsoft/phi-2",
         "microsoft/Phi-3.5-mini-instruct": "https://huggingface.co/microsoft/Phi-3.5-mini-instruct",
+        "microsoft/Phi-3.5-vision-instruct": "https://huggingface.co/microsoft/Phi-3.5-vision-instruct",
     }
 )

_unittests/ut_torch_models/test_hghub.py

@@ -0,0 +1,106 @@
+import unittest
+import pandas
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    hide_stdout,
+    never_test,
+    requires_torch,
+    requires_transformers,
+)
+from onnx_diagnostic.torch_models.hghub.hub_api import (
+    enumerate_model_list,
+    get_model_info,
+    get_pretrained_config,
+    task_from_id,
+    task_from_arch,
+    task_from_tags,
+)
+from onnx_diagnostic.torch_models.hghub.hub_data import load_architecture_task
+
+
+class TestHuggingFaceHubApi(ExtTestCase):
+
+    @requires_transformers("4.50")  # we limit to some versions of the CI
+    @requires_torch("2.7")
+    def test_enumerate_model_list(self):
+        models = list(
+            enumerate_model_list(
+                2,
+                verbose=1,
+                dump="test_enumerate_model_list.csv",
+                filter="text-generation",
+                library="transformers",
+            )
+        )
+        self.assertEqual(len(models), 2)
+        df = pandas.read_csv("test_enumerate_model_list.csv")
+        self.assertEqual(df.shape, (2, 12))
+        tasks = [task_from_id(c) for c in df.id]
+        self.assertEqual(["text-generation", "text-generation"], tasks)
+
+    @requires_transformers("4.50")
+    @requires_torch("2.7")
+    def test_task_from_id(self):
+        for name, etask in [
+            ("arnir0/Tiny-LLM", "text-generation"),
+            ("microsoft/phi-2", "text-generation"),
+            ("microsoft/Phi-3.5-mini-instruct", "text-generation"),
+            ("microsoft/Phi-3.5-vision-instruct", "text-generation"),
+        ]:
+            with self.subTest(name=name, task=etask):
+                task = task_from_id(name, True)
+                self.assertEqual(etask, task)
+
+    @requires_transformers("4.50")
+    @requires_torch("2.7")
+    @hide_stdout()
+    def test_get_pretrained_config(self):
+        conf = get_pretrained_config("microsoft/phi-2")
+        self.assertNotEmpty(conf)
+        print(conf)
+
+    @requires_transformers("4.50")
+    @requires_torch("2.7")
+    @hide_stdout()
+    def test_get_model_info(self):
+        info = get_model_info("microsoft/phi-2")
+        self.assertEqual(info.pipeline_tag, "text-generation")
+
+        info = get_model_info("microsoft/Phi-3.5-vision-instruct")
+        self.assertEqual(info.pipeline_tag, "image-text-to-text")
+
+        info = get_model_info("microsoft/Phi-4-multimodal-instruct")
+        self.assertEqual(info.pipeline_tag, "automatic-speech-recognition")
+
+    def test_task_from_arch(self):
+        task = task_from_arch("LlamaForCausalLM")
+        self.assertEqual("text-generation", task)
+
+    @never_test()
+    def test_hf_all_models(self):
+        list(enumerate_model_list(-1, verbose=1, dump="test_hf_all_models.csv"))
+
+    def test_load_architecture_task(self):
+        data = load_architecture_task()
+        print(set(data.values()))
+
+    def test_task_from_tags(self):
+        _tags = [
+            ("text-generation|nlp|code|en|text-generation-inference", "text-generation"),
+            (
+                "text-generation|nlp|code|vision|image-text-to-text|conversational",
+                "image-text-to-text",
+            ),
+            (
+                "text-generation|nlp|code|audio|automatic-speech-recognition|speech-summarization|speech-translation|visual-question-answering",
+                "automatic-speech-recognition",
+            ),
+        ]
+        for tags, etask in _tags:
+            with self.subTest(tags=tags, task=etask):
+                task = task_from_tags(tags, True)
+                self.assertEqual(etask, task)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_torch_models/test_hghub_api.py

@@ -0,0 +1,33 @@
+import unittest
+import transformers
+from onnx_diagnostic.ext_test_case import (
+    ExtTestCase,
+    hide_stdout,
+    requires_torch,
+    requires_transformers,
+)
+from onnx_diagnostic.torch_models.hghub.model_inputs import (
+    config_class_from_architecture,
+    get_untrained_model_with_inputs,
+)
+
+
+class TestHuggingFaceHubModel(ExtTestCase):
+    @requires_transformers("4.50")  # we limit to some versions of the CI
+    @requires_torch("2.7")
+    def test_config_class_from_architecture(self):
+        config = config_class_from_architecture("LlamaForCausalLM")
+        self.assertEqual(config, transformers.LlamaConfig)
+
+    @hide_stdout()
+    def test_get_untrained_model_with_inputs(self):
+        mid = "arnir0/Tiny-LLM"
+        data = get_untrained_model_with_inputs(mid, verbose=1)
+        model, inputs = data["model"], data["inputs"]
+        model(**inputs)
+        self.assertEqual(data["size"], 1858125824)
+        self.assertEqual(data["n_weights"], 464531456)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

_unittests/ut_torch_models/test_hghub_model.py

@@ -0,0 +1 @@
+from .model_inputs import get_untrained_model_with_inputs

onnx_diagnostic/torch_models/hghub/__init__.py

@@ -1,14 +1,21 @@
+import functools
 from typing import List, Optional, Union
 import transformers
-from huggingface_hub import HfApi
-from .hub_data import __date__, load_architecture_task
+from huggingface_hub import HfApi, model_info
+from .hub_data import __date__, __data_tasks__, load_architecture_task
 
 
 def get_pretrained_config(model_id) -> str:
     """Returns the config for a model_id."""
     return transformers.AutoConfig.from_pretrained(model_id)
 
 
+def get_model_info(model_id) -> str:
+    """Returns the model info for a model_id."""
+    return model_info(model_id)
+
+
+@functools.cache
 def task_from_arch(arch: str) -> str:
     """
     This function relies on stored information. That information needs to be refresh.
@@ -51,6 +58,20 @@ def task_from_id(model_id: str, pretrained: bool = False) -> str:
     return transformers.pipelines.get_task(model_id)
 
 
+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 separater.
+    """
+    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,
     task: Optional[str] = None,
@@ -92,6 +113,7 @@ def enumerate_model_list(
                 ",".join(
                     [
                         "id",
+                        "model_name",
                         "author",
                         "created_at",
                         "last_modified",
@@ -123,6 +145,7 @@ def enumerate_model_list(
                         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],

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -3,6 +3,27 @@
 
 __date__ = "2025-03-26"
 
+__data_tasks__ = [
+    "automatic-speech-recognition",
+    "image-to-text",
+    "text-generation",
+    "object-detection",
+    "document-question-answering",
+    "feature-extraction",
+    "text-to-audio",
+    "zero-shot-image-classification",
+    "image-segmentation",
+    "reinforcement-learning",
+    "no-pipeline-tag",
+    "image-classification",
+    "text2text-generation",
+    "mask-generation",
+    "keypoint-detection",
+    "audio-classification",
+    "image-feature-extraction",
+    "fill-mask",
+]
+
 __data_arch__ = """
 architecture,task
 ASTModel,feature-extraction