fix a few things

Author: xadupre

_doc/examples/plot_export_tiny_phi2.py

@@ -50,7 +50,7 @@
     data["n_weights"],
 )
 
-print(f"model {size / 2**20:1.3f} Mb with {n_weights // 1000} mille parameters.")
+print(f"model {size / 2**20:1.1f} Mb with {n_weights // 1000} thousands of parameters.")
 # %%
 # The original model has 2.7 billion parameters. It was divided by more than 10.
 # However, it can still be used with
@@ -209,6 +209,7 @@
 # Every mixture of models goes through a control flow (a test).
 # It also happens when a cache is truncated. The code of the model
 # needs to be changed. See example :ref:`l-plot-export-cond`.
+# Loops are not supported yet.
 #
 # Issue with dynamic shapes
 # +++++++++++++++++++++++++

onnx_diagnostic/tasks/image_classification.py

@@ -7,6 +7,13 @@
 
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
+    if (
+        hasattr(config, "model_type")
+        and config.model_type == "timm_wrapper"
+        and not hasattr(config, "num_hidden_layers")
+    ):
+        # We cannot reduce.
+        return {}
     check_hasattr(config, ("num_hidden_layers", "hidden_sizes"))
     kwargs = dict(
         num_hidden_layers=(
@@ -82,6 +89,20 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
     If the configuration is None, the function selects typical dimensions.
     """
     if config is not None:
+        if (
+            hasattr(config, "model_type")
+            and config.model_type == "timm_wrapper"
+            and not hasattr(config, "num_hidden_layers")
+        ):
+            input_size = config.pretrained_cfg["input_size"]
+            kwargs = dict(
+                batch_size=2,
+                input_width=input_size[-2],
+                input_height=input_size[-1],
+                input_channels=input_size[-3],
+            )
+            return kwargs, get_inputs
+
         check_hasattr(config, ("image_size", "architectures"), "num_channels")
     if config is not None:
         if hasattr(config, "image_size"):

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -88,6 +88,15 @@ def get_model_info(model_id) -> Any:
     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) -> str:
     """
@@ -126,7 +135,7 @@ def task_from_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: balls back to pretrained config
+    :param fall_back_to_pretrained: falls back to pretrained config
     :return: task
     """
     if not pretrained:
@@ -139,6 +148,9 @@ def task_from_id(
     try:
         return config.pipeline_tag
     except AttributeError:
+        guess = _guess_task_from_config(config)
+        if guess is not None:
+            return guess
         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}"

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -78,6 +78,7 @@
     MobileBertModel,feature-extraction
     MobileNetV1Model,image-feature-extraction
     MobileNetV2Model,image-feature-extraction
+    mobilenetv3_small_100,image-classification
     MobileViTForImageClassification,image-classification
     ModernBertForMaskedLM,fill-mask
     Phi4MMForCausalLM,MoE
@@ -202,7 +203,7 @@ def load_models_testing() -> List[str]:
 @functools.cache
 def load_architecture_task() -> Dict[str, str]:
     """
-    Returns a dictionary mapping architecture to task.
+    Returns a dictionary mapping architectures to tasks.
 
     import pprint
     from onnx_diagnostic.torch_models.hghub.hub_data import load_architecture_task

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -4,7 +4,7 @@
 import transformers
 from ...helpers.config_helper import update_config
 from ...tasks import reduce_model_config, random_input_kwargs
-from .hub_api import task_from_arch, get_pretrained_config
+from .hub_api import task_from_arch, task_from_id, get_pretrained_config
 
 
 def get_untrained_model_with_inputs(
@@ -64,17 +64,21 @@ def get_untrained_model_with_inputs(
         config = get_pretrained_config(
             model_id, use_preinstalled=use_preinstalled, **(model_kwargs or {})
         )
+    if hasattr(config, "architecture") and config.architecture:
+        archs = [config.architecture]
     archs = config.architectures  # type: ignore
-    assert archs is not None and len(archs) == 1, (
+    task = None
+    if archs is None:
+        task = task_from_id(model_id)
+    assert task is not None or (archs is not None and len(archs) == 1), (
         f"Unable to determine the architecture for model {model_id!r}, "
         f"architectures={archs!r}, conf={config}"
     )
-    arch = archs[0]
-    if verbose:
-        print(f"[get_untrained_model_with_inputs] architecture={arch!r}")
     if verbose:
+        print(f"[get_untrained_model_with_inputs] architectures={archs!r}")
         print(f"[get_untrained_model_with_inputs] cls={config.__class__.__name__!r}")
-    task = task_from_arch(arch)
+    if task is None:
+        task = task_from_arch(archs[0])
     if verbose:
         print(f"[get_untrained_model_with_inputs] task={task!r}")
 
@@ -106,7 +110,15 @@ def get_untrained_model_with_inputs(
     if inputs_kwargs:
         kwargs.update(inputs_kwargs)
 
-    model = getattr(transformers, arch)(config)
+    if archs is not None:
+        model = getattr(transformers, archs[0])(config)
+    else:
+        assert same_as_pretrained, (
+            f"Model {model_id!r} cannot be built, the model cannot be built. "
+            f"It must be downloaded. Use same_as_pretrained=True."
+        )
+        model = None
+
     # This line is important. Some models may produce different
     # outputs even with the same inputs in training mode.
     model.eval()