fix dummy inputs

Author: xadupre

_unittests/ut_torch_models/test_hghub_api.py

@@ -98,7 +98,7 @@ def test_task_from_tags(self):
         ]
         for tags, etask in _tags:
             with self.subTest(tags=tags, task=etask):
-                task = task_from_tags(tags, True)
+                task = task_from_tags(tags)
                 self.assertEqual(etask, task)
 
 

_unittests/ut_torch_models/test_hghub_model.py

@@ -20,13 +20,27 @@ def test_config_class_from_architecture(self):
         self.assertEqual(config, transformers.LlamaConfig)
 
     @hide_stdout()
-    def test_get_untrained_model_with_inputs(self):
+    def test_get_untrained_model_with_inputs_tiny_llm(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)
+        self.assertEqual((data["size"], data["n_weights"]), (1858125824, 464531456))
+
+    @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((data["size"], data["n_weights"]), (126190824, 31547706))
+
+    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((data["size"], data["n_weights"]), (4291141632, 1072785408))
 
 
 if __name__ == "__main__":

onnx_diagnostic/torch_models/hghub/hub_api.py

@@ -5,7 +5,7 @@
 from .hub_data import __date__, __data_tasks__, load_architecture_task
 
 
-def get_pretrained_config(model_id) -> str:
+def get_pretrained_config(model_id: str) -> str:
     """Returns the config for a model_id."""
     return transformers.AutoConfig.from_pretrained(model_id)
 
@@ -61,7 +61,7 @@ def task_from_id(model_id: str, pretrained: bool = False) -> str:
 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 given by a string, ``|`` should be the separator.
     """
     if isinstance(tags, str):
         tags = tags.split("|")

onnx_diagnostic/torch_models/hghub/hub_data.py

@@ -5,6 +5,7 @@
 
 __data_tasks__ = [
     "automatic-speech-recognition",
+    "image-text-to-text",
     "image-to-text",
     "text-generation",
     "object-detection",

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -10,20 +10,31 @@
 
 
 @functools.cache
-def config_class_from_architecture(arch: str) -> type:
+def config_class_from_architecture(arch: str, exc: bool = False) -> type:
     """
     Retrieves the configuration class for a given architecture.
+
+    :param arch: architecture (clas name)
+    :param exc: raise an exception if not found
+    :return: type
     """
     cls = getattr(transformers, arch)
     mod_name = cls.__module__
     mod = importlib.import_module(mod_name)
     source = inspect.getsource(mod)
     reg = re.compile("config: ([A-Za-z0-9]+)")
     fall = reg.findall(source)
+    if len(fall) == 0:
+        assert not exc, (
+            f"Unable to guess Configuration class name for arch={arch!r}, "
+            f"module={mod_name!r}, no candidate, source is\n{source}"
+        )
+        return None
     unique = set(fall)
     assert len(unique) == 1, (
         f"Unable to guess Configuration class name for arch={arch!r}, "
-        f"module={mod_name!r}, source is\n{source}"
+        f"module={mod_name!r}, found={unique} (#{len(unique)}), "
+        f"source is\n{source}"
     )
     cls_name = unique.pop()
     return getattr(transformers, cls_name)
@@ -81,7 +92,14 @@ def get_untrained_model_with_inputs(
     arch = archs[0]
     if verbose:
         print(f"[get_untrained_model_with_inputs] architecture={arch!r}")
-    cls = config_class_from_architecture(arch)
+    cls = config_class_from_architecture(arch, exc=False)
+    if cls is None:
+        if verbose:
+            print(
+                "[get_untrained_model_with_inputs] no found config name in the code, loads it"
+            )
+        config = get_pretrained_config(model_id)
+        cls = config.__class__
     if verbose:
         print(f"[get_untrained_model_with_inputs] cls={cls.__name__!r}")
 
@@ -107,12 +125,16 @@ def get_untrained_model_with_inputs(
             batch_size=2,
             sequence_length=30,
             sequence_length2=3,
-            num_hidden_layers=config.num_hidden_layers,
-            num_key_value_heads=config.num_key_value_heads,
             head_dim=getattr(
                 config, "head_dim", config.hidden_size // config.num_attention_heads
             ),
             max_token_id=config.vocab_size - 1,
+            num_hidden_layers=config.num_hidden_layers,
+            num_key_value_heads=(
+                config.num_key_value_heads
+                if hasattr(config, "num_key_value_heads")
+                else config.num_attention_heads
+            ),
         )
         if inputs_kwargs:
             kwargs.update(inputs_kwargs)