refacto

xadupre · xadupre · commit a90e77157ed5 · 2025-03-27T23:36:59.000+01:00
diff --git a/_unittests/ut_torch_models/test_hghub_api.py b/_unittests/ut_torch_models/test_hghub_api.py
@@ -26,21 +26,22 @@ class TestHuggingFaceHubApi(ExtTestCase):
 
     @requires_transformers("4.50")  # we limit to some versions of the CI
     @requires_torch("2.7")
+    @hide_stdout()
     def test_enumerate_model_list(self):
         models = list(
             enumerate_model_list(
                 2,
                 verbose=1,
                 dump="test_enumerate_model_list.csv",
-                filter="text-generation",
+                filter="image-classification",
                 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)
+        tasks = [task_from_id(c, "missing") for c in df.id]
+        self.assertEqual(len(tasks), 2)
 
     @requires_transformers("4.50")
     @requires_torch("2.7")
diff --git a/_unittests/ut_torch_models/test_hghub_model.py b/_unittests/ut_torch_models/test_hghub_model.py
@@ -29,27 +29,36 @@ def test_get_untrained_model_with_inputs_tiny_llm(self):
         data = get_untrained_model_with_inputs(mid, verbose=1)
         self.assertEqual(
             set(data),
-            {"model", "inputs", "dynamic_shapes", "configuration", "size", "n_weights"},
+            {
+                "model",
+                "inputs",
+                "dynamic_shapes",
+                "configuration",
+                "size",
+                "n_weights",
+                "input_kwargs",
+                "model_kwargs",
+            },
         )
         model, inputs = data["model"], data["inputs"]
         model(**inputs)
-        self.assertEqual((1858125824, 464531456), (data["size"], data["n_weights"]))
+        self.assertEqual((51955968, 12988992), (data["size"], data["n_weights"]))
 
     @hide_stdout()
     def test_get_untrained_model_with_inputs_tiny_xlm_roberta(self):
         mid = "hf-internal-testing/tiny-xlm-roberta"  # XLMRobertaConfig
         data = get_untrained_model_with_inputs(mid, verbose=1)
         model, inputs = data["model"], data["inputs"]
         model(**inputs)
-        self.assertEqual((126190824, 31547706), (data["size"], data["n_weights"]))
+        self.assertEqual((8642088, 2160522), (data["size"], data["n_weights"]))
 
     @hide_stdout()
     def test_get_untrained_model_with_inputs_tiny_gpt_neo(self):
         mid = "hf-internal-testing/tiny-random-GPTNeoXForCausalLM"
         data = get_untrained_model_with_inputs(mid, verbose=1)
         model, inputs = data["model"], data["inputs"]
         model(**inputs)
-        self.assertEqual((4291141632, 1072785408), (data["size"], data["n_weights"]))
+        self.assertEqual((316712, 79178), (data["size"], data["n_weights"]))
 
     @hide_stdout()
     def test_get_untrained_model_with_inputs_phi_2(self):
@@ -60,7 +69,7 @@ def test_get_untrained_model_with_inputs_phi_2(self):
         # different expected value for different version of transformers
         self.assertIn(
             (data["size"], data["n_weights"]),
-            [(1040293888, 260073472), (1040498688, 260124672)],
+            [(453330944, 113332736)],
         )
 
     @hide_stdout()
@@ -70,7 +79,7 @@ def test_get_untrained_model_with_inputs_beit(self):
         model, inputs = data["model"], data["inputs"]
         model(**inputs)
         # different expected value for different version of transformers
-        self.assertIn((data["size"], data["n_weights"]), [(30732296, 7683074)])
+        self.assertIn((data["size"], data["n_weights"]), [(111448, 27862)])
 
     @hide_stdout()
     @long_test()
@@ -81,7 +90,7 @@ def _diff(c1, c2):
                 if isinstance(v, (str, dict, list, tuple, int, float)) and v != getattr(
                     c2, k, None
                 ):
-                    rows.append(f"{k} :: -- {v} ++ {getattr(c2, k, "MISS")}")
+                    rows.append(f"{k} :: -- {v} ++ {getattr(c2, k, 'MISS')}")
             return "\n".join(rows)
 
         # UNHIDE=1 LONGTEST=1 python _unittests/ut_torch_models/test_hghub_model.py -k L -f
diff --git a/onnx_diagnostic/torch_models/hghub/hub_api.py b/onnx_diagnostic/torch_models/hghub/hub_api.py
@@ -5,9 +5,11 @@
 from .hub_data import __date__, __data_tasks__, load_architecture_task
 
 
-def get_pretrained_config(model_id: str) -> str:
+def get_pretrained_config(model_id: str, trust_remote_code: bool = True) -> str:
     """Returns the config for a model_id."""
-    return transformers.AutoConfig.from_pretrained(model_id)
+    return transformers.AutoConfig.from_pretrained(
+        model_id, trust_remote_code=trust_remote_code
+    )
 
 
 def get_model_info(model_id) -> str:
@@ -16,11 +18,12 @@ def get_model_info(model_id) -> str:
 
 
 @functools.cache
-def task_from_arch(arch: str) -> str:
+def task_from_arch(arch: str, default_value: 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
     :return: task
 
     .. runpython::
@@ -33,17 +36,24 @@ def task_from_arch(arch: str) -> str:
     <onnx_diagnostic.torch_models.hghub.hub_data.load_architecture_task>`.
     """
     data = load_architecture_task()
+    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__}"
     return data[arch]
 
 
 def task_from_id(
-    model_id: str, pretrained: bool = False, fall_back_to_pretrained: bool = True
+    model_id: str,
+    default_value: Optional[str] = None,
+    pretrained: bool = False,
+    fall_back_to_pretrained: bool = True,
 ) -> 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: balls back to pretrained config
     :return: task
@@ -62,7 +72,7 @@ def task_from_id(
             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 task_from_arch(config.architectures[0], default_value=default_value)
 
 
 def task_from_tags(tags: Union[str, List[str]]) -> str:
diff --git a/onnx_diagnostic/torch_models/hghub/model_inputs.py b/onnx_diagnostic/torch_models/hghub/model_inputs.py
@@ -2,7 +2,7 @@
 import importlib
 import inspect
 import re
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 import transformers
 from ...cache_helpers import make_dynamic_cache
@@ -46,6 +46,104 @@ def _update_config(config: Any, kwargs: Dict[str, Any]):
             setattr(config, k, v)
 
 
+def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
+    """Reduces a model size."""
+    if task == "text-generation":
+        kwargs = dict(
+            head_dim=getattr(
+                config, "head_dim", config.hidden_size // config.num_attention_heads
+            ),
+            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_key_value_heads=(
+                config.num_key_value_heads
+                if hasattr(config, "num_key_value_heads")
+                else config.num_attention_heads
+            ),
+            intermediate_size=(
+                min(config.intermediate_size, 24576 // 4)
+                if config.intermediate_size % 4 == 0
+                else config.intermediate_size
+            ),
+            hidden_size=(
+                min(config.hidden_size, 3072 // 4)
+                if config.hidden_size % 4 == 0
+                else config.hidden_size
+            ),
+        )
+    elif task == "image-classification":
+        if isinstance(config.image_size, int):
+            kwargs = dict(
+                batch_size=2,
+                input_width=config.image_size,
+                input_height=config.image_size,
+                input_channels=config.num_channels,
+            )
+        else:
+            kwargs = dict(
+                batch_size=2,
+                input_width=config.image_size[0],
+                input_height=config.image_size[1],
+                input_channels=config.num_channels,
+            )
+    else:
+        raise NotImplementedError(f"Input generation for task {task!r} not implemented yet.")
+
+    for k, v in kwargs.items():
+        setattr(config, k, v)
+    return kwargs
+
+
+def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:
+    """Inputs kwargs"""
+    if task == "text-generation":
+        kwargs = dict(
+            batch_size=2,
+            sequence_length=30,
+            sequence_length2=3,
+            head_dim=getattr(
+                config, "head_dim", config.hidden_size // config.num_attention_heads
+            ),
+            dummy_max_token_id=config.vocab_size - 1,
+            num_hidden_layers=min(config.num_hidden_layers, 2),
+            num_key_value_heads=(
+                config.num_key_value_heads
+                if hasattr(config, "num_key_value_heads")
+                else config.num_attention_heads
+            ),
+            intermediate_size=(
+                min(config.intermediate_size, 24576 // 4)
+                if config.intermediate_size % 4 == 0
+                else config.intermediate_size
+            ),
+            hidden_size=(
+                min(config.hidden_size, 3072 // 4)
+                if config.hidden_size % 4 == 0
+                else config.hidden_size
+            ),
+        )
+        fct = get_inputs_for_text_generation
+    elif task == "image-classification":
+        if isinstance(config.image_size, int):
+            kwargs = dict(
+                batch_size=2,
+                input_width=config.image_size,
+                input_height=config.image_size,
+                input_channels=config.num_channels,
+            )
+        else:
+            kwargs = dict(
+                batch_size=2,
+                input_width=config.image_size[0],
+                input_height=config.image_size[1],
+                input_channels=config.num_channels,
+            )
+        fct = get_inputs_for_image_classification  # type: ignore
+    else:
+        raise NotImplementedError(f"Input generation for task {task!r} not implemented yet.")
+
+    return kwargs, fct
+
+
 def get_untrained_model_with_inputs(
     model_id: str,
     config: Optional[Any] = None,
@@ -114,63 +212,26 @@ def get_untrained_model_with_inputs(
         config.rope_scaling = (
             {"rope_type": "dynamic", "factor": 10.0} if dynamic_rope else None
         )
+
+    # updating the configuration
+    if not same_as_pretrained:
+        mkwargs = reduce_model_config(config, task)
+    else:
+        mkwargs = {}
     if model_kwargs:
         for k, v in model_kwargs.items():
             setattr(config, k, v)
-
-    if task == "text-generation":
-        kwargs = dict(
-            batch_size=2,
-            sequence_length=30,
-            sequence_length2=3,
-            head_dim=getattr(
-                config, "head_dim", config.hidden_size // config.num_attention_heads
-            ),
-            dummy_max_token_id=config.vocab_size - 1,
-            num_hidden_layers=min(config.num_hidden_layers, 2),
-            num_key_value_heads=(
-                config.num_key_value_heads
-                if hasattr(config, "num_key_value_heads")
-                else config.num_attention_heads
-            ),
-            intermediate_size=(
-                min(config.intermediate_size, 24576 // 4)
-                if config.intermediate_size % 4 == 0
-                else config.intermediate_size
-            ),
-            hidden_size=(
-                min(config.hidden_size, 3072 // 4)
-                if config.hidden_size % 4 == 0
-                else config.hidden_size
-            ),
-        )
-
-        fct = get_inputs_for_text_generation
-    elif task == "image-classification":
-        if isinstance(config.image_size, int):
-            kwargs = dict(
-                batch_size=2,
-                input_width=config.image_size,
-                input_height=config.image_size,
-                input_channels=config.num_channels,
-            )
-        else:
-            kwargs = dict(
-                batch_size=2,
-                input_width=config.image_size[0],
-                input_height=config.image_size[1],
-                input_channels=config.num_channels,
-            )
-        fct = get_inputs_for_image_classification
-    else:
-        raise NotImplementedError(f"Input generation for task {task!r} not implemented yet.")
-
+            mkwargs[k] = v
+    # input kwargs
+    kwargs, fct = random_input_kwargs(config, task)
     if inputs_kwargs:
         kwargs.update(inputs_kwargs)
-    true_kwargs = (inputs_kwargs or {}) if same_as_pretrained else kwargs
-    _update_config(config, true_kwargs)
+
     model = getattr(transformers, arch)(config)
-    return fct(model, config, **true_kwargs)
+    res = fct(model, config, **kwargs)
+    res["input_kwargs"] = kwargs
+    res["model_kwargs"] = mkwargs
+    return res
 
 
 def compute_model_size(model: torch.nn.Module) -> Tuple[int, int]:
