return configuration as well

Author: xadupre

_unittests/ut_torch_models/test_hghub_api.py

@@ -3,6 +3,7 @@
 from onnx_diagnostic.ext_test_case import (
     ExtTestCase,
     hide_stdout,
+    long_test,
     never_test,
     requires_torch,
     requires_transformers,
@@ -15,7 +16,10 @@
     task_from_arch,
     task_from_tags,
 )
-from onnx_diagnostic.torch_models.hghub.hub_data import load_architecture_task
+from onnx_diagnostic.torch_models.hghub.hub_data import (
+    load_architecture_task,
+    load_models_testing,
+)
 
 
 class TestHuggingFaceHubApi(ExtTestCase):
@@ -111,6 +115,17 @@ def test_task_from_tags(self):
                 task = task_from_tags(tags)
                 self.assertEqual(etask, task)
 
+    def test_model_testings(self):
+        models = load_models_testing()
+        self.assertNotEmpty(models)
+
+    @long_test()
+    def test_model_testings_and_architctures(self):
+        models = load_models_testing()
+        for mid in models:
+            task = task_from_id(mid)
+            self.assertNotEmpty(task)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

_unittests/ut_torch_models/test_hghub_model.py

@@ -23,6 +23,10 @@ def test_config_class_from_architecture(self):
     def test_get_untrained_model_with_inputs_tiny_llm(self):
         mid = "arnir0/Tiny-LLM"
         data = get_untrained_model_with_inputs(mid, verbose=1)
+        self.assertEqual(
+            set(data),
+            {"model", "inputs", "dynamic_shapes", "configuration", "size", "n_weights"},
+        )
         model, inputs = data["model"], data["inputs"]
         model(**inputs)
         self.assertEqual((1858125824, 464531456), (data["size"], data["n_weights"]))

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -37,25 +37,32 @@ def task_from_arch(arch: str) -> str:
     return data[arch]
 
 
-def task_from_id(model_id: str, pretrained: bool = False) -> str:
+def task_from_id(
+    model_id: str, pretrained: bool = False, fall_back_to_pretrained: bool = True
+) -> str:
     """
     Returns the task attached to a model id.
 
     :param model_id: model id
     :param pretrained: uses the config
+    :param fall_back_to_pretrained: balls back to pretrained config
     :return: task
     """
-    if pretrained:
-        config = get_pretrained_config(model_id)
+    if not pretrained:
         try:
-            return config.pipeline_tag
-        except AttributeError:
-            assert config.architectures is not None and len(config.architectures) == 1, (
-                f"Cannot return the task of {model_id!r}, pipeline_tag is not setup, "
-                f"architectures={config.architectures} in config={config}"
-            )
-            return task_from_arch(config.architectures[0])
-    return transformers.pipelines.get_task(model_id)
+            transformers.pipelines.get_task(model_id)
+        except RuntimeError:
+            if not fall_back_to_pretrained:
+                raise
+    config = get_pretrained_config(model_id)
+    try:
+        return config.pipeline_tag
+    except AttributeError:
+        assert config.architectures is not None and len(config.architectures) == 1, (
+            f"Cannot return the task of {model_id!r}, pipeline_tag is not setup, "
+            f"architectures={config.architectures} in config={config}"
+        )
+        return task_from_arch(config.architectures[0])
 
 
 def task_from_tags(tags: Union[str, List[str]]) -> str:

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -1,30 +1,9 @@
 import io
 import functools
+from typing import Dict, List
 
 __date__ = "2025-03-26"
 
-__data_tasks__ = [
-    "automatic-speech-recognition",
-    "image-text-to-text",
-    "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
@@ -143,9 +122,61 @@
 YolosModel,image-feature-extraction
 """
 
+__data_tasks__ = [
+    "automatic-speech-recognition",
+    "image-text-to-text",
+    "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",
+]
+
+__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
+"""
+
+
+@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():
+def load_architecture_task() -> Dict[str, str]:
     """
     Returns a dictionary mapping architecture to task.
 

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -62,7 +62,7 @@ def get_untrained_model_with_inputs(
     :param model_kwargs: to change the model generation
     :param verbose: display found information
     :param dynamic_rope: use dynamic rope (see :class:`transformers.LlamaConfig`)
-    :return: dictionary
+    :return: dictionary with a model, inputs, dynamic shapes, and the configuration
 
     Example:
 
@@ -79,6 +79,7 @@ def get_untrained_model_with_inputs(
         print("-- number of parameters:", data["n_weights"])
         print("-- inputs:", string_type(data["inputs"], with_shape=True))
         print("-- dynamic shapes:", pprint.pformat(data["dynamic_shapes"]))
+        print("-- configuration:", pprint.pformat(data["configuration"]))
     """
     if verbose:
         print(f"[get_untrained_model_with_inputs] model_id={model_id!r}")
@@ -139,7 +140,7 @@ def get_untrained_model_with_inputs(
         if inputs_kwargs:
             kwargs.update(inputs_kwargs)
 
-        return get_inputs_for_text_generation(model, **kwargs)
+        return get_inputs_for_text_generation(model, config, **kwargs)
     raise NotImplementedError(f"Input generation for task {task!r} not implemented yet.")
 
 
@@ -155,6 +156,7 @@ def compute_model_size(model: torch.nn.Module) -> Tuple[int, int]:
 
 def get_inputs_for_text_generation(
     model: torch.nn.Module,
+    config: Optional[Any],
     max_token_id: int,
     num_key_value_heads: int,
     num_hidden_layers: int,
@@ -167,6 +169,7 @@ def get_inputs_for_text_generation(
 ):
     """
     :param model: model to get the missing information
+    :param config: configuration used to generate the model
     :param head_dim: last dimension of the cache
     :param batch_size: batch size
     :param sequence_length: sequence length
@@ -216,5 +219,10 @@ def get_inputs_for_text_generation(
     )
     sizes = compute_model_size(model)
     return dict(
-        model=model, inputs=inputs, dynamic_shapes=shapes, size=sizes[0], n_weights=sizes[1]
+        model=model,
+        inputs=inputs,
+        dynamic_shapes=shapes,
+        size=sizes[0],
+        n_weights=sizes[1],
+        configuration=config,
     )