allow to convert a model with other task inputs

Author: xadupre

onnx_diagnostic/tasks/text_generation.py

@@ -19,6 +19,9 @@
 def reduce_model_config(config: Any) -> Dict[str, Any]:
     """Reduces a model size."""
     # FalconMambaConfig: use_mambapy
+    if hasattr(config, "text_config"):
+        # The model is probably of mixture of models used only for text.
+        config = config.text_config
     check_hasattr(
         config,
         ("head_dim", ("hidden_size", "num_attention_heads"), "use_mambapy"),
@@ -308,6 +311,9 @@ def random_input_kwargs(config: Any) -> Tuple[Dict[str, Any], Callable]:
 
     If the configuration is None, the function selects typical dimensions.
     """
+    if hasattr(config, "text_config"):
+        # The model is probably of mixture of models used only for text.
+        config = config.text_config
     if config is not None:
         check_hasattr(
             config,

onnx_diagnostic/torch_models/hghub/model_inputs.py

@@ -120,14 +120,16 @@ def get_untrained_model_with_inputs(
             **(model_kwargs or {}),
         )
 
-    model, task, mkwargs, diff_config = None, None, {}, None
+    model, task_, mkwargs, diff_config = None, None, {}, None
     if use_pretrained and same_as_pretrained:
         if model_id in HANDLED_MODELS:
-            model, task, config = load_specific_model(model_id, verbose=verbose)
+            model, task_, config = load_specific_model(model_id, verbose=verbose)
 
+    if task is None:
+        task = task_
     if model is None:
         arch = architecture_from_config(config)
-        if arch is None:
+        if task is None and arch is None:
             task = task_from_id(model_id, subfolder=subfolder)
         assert task is not None or arch is not None, (
             f"Unable to determine the architecture for model {model_id!r}, "

onnx_diagnostic/torch_models/validate.py

@@ -117,11 +117,21 @@ def _make_folder_name(
     drop_inputs: Optional[List[str]] = None,
     same_as_pretrained: bool = False,
     use_pretrained: bool = False,
+    task: Optional[str] = None,
 ) -> str:
     "Creates a filename unique based on the given options."
     els = [model_id.replace("/", "_")]
     if subfolder:
         els.append(subfolder.replace("/", "_"))
+    if not task:
+        els.append(task)
+    if drop_inputs:
+        ii = "-".join(f"{s[0]}{s[-1]}" for s in drop_inputs)
+        els.append(f"I-{ii.upper()}")
+    if use_pretrained:
+        els.append("TRAINED")
+    elif same_as_pretrained:
+        els.append("SAMESIZE")
     if exporter:
         els.append(exporter)
     if optimization:
@@ -142,14 +152,7 @@ def _make_folder_name(
         els.append(sdev)
     if opset is not None:
         els.append(f"op{opset}")
-    if drop_inputs:
-        ii = "-".join(f"{s[0]}{s[-1]}" for s in drop_inputs)
-        els.append(f"I-{ii.upper()}")
-    if use_pretrained:
-        els.append("TRAINED")
-    elif same_as_pretrained:
-        els.append("SAMESIZE")
-    return "-".join(els)
+    return "/".join(els)
 
 
 def version_summary() -> Dict[str, Union[int, float, str]]:
@@ -476,6 +479,7 @@ def validate_model(
             drop_inputs=drop_inputs,
             use_pretrained=use_pretrained,
             same_as_pretrained=same_as_pretrained,
+            task=task,
         )
         dump_folder = os.path.join(dump_folder, folder_name)
         if not os.path.exists(dump_folder):
@@ -490,6 +494,8 @@ def validate_model(
             print(f"[validate_model] validate model id {model_id!r}, subfolder={subfolder!r}")
         else:
             print(f"[validate_model] validate model id {model_id!r}")
+        if task:
+            print(f"[validate_model] with task {task!r}")
         print(f"[validate_model] patch={patch!r}")
         if model_options:
             print(f"[validate_model] model_options={model_options!r}")
@@ -765,6 +771,7 @@ def validate_model(
             ep = data["exported_program"]
             if verbose:
                 print(f"[validate_model] -- dumps exported program in {dump_folder!r}...")
+            folder_name = folder_name.replace("/", "-")
             with open(os.path.join(dump_folder, f"{folder_name}.ep"), "w") as f:
                 f.write(str(ep))
             torch.export.save(ep, os.path.join(dump_folder, f"{folder_name}.pt2"))
@@ -773,6 +780,7 @@ def validate_model(
             if verbose:
                 print("[validate_model] done (dump ep)")
         if "onnx_program" in data:
+            folder_name = folder_name.replace("/", "-")
             epo = data["onnx_program"]
             if verbose:
                 print(f"[validate_model] dumps onnx program in {dump_folder!r}...")