fix import issues

Author: xadupre

onnx_diagnostic/_command_lines_parser.py

@@ -1265,9 +1265,10 @@ def _size(name):
     print(f"-- load ep {args.ep!r}")
     begin = time.perf_counter()
     # We need to load the plugs.
-    from .torch_export_patches.patches.patch_transformers import PLUGS_Qwen25
+    from .torch_export_patches.patches.patch_transformers import get_transformers_plugs
 
-    assert len(PLUGS_Qwen25) == 1, "Missing PLUGS for Qwen2.5"
+    plugs = get_transformers_plugs()
+    assert plugs, "Missing PLUGS for Qwen2.5"
     ep = torch.export.load(args.ep)
     print(f"-- done in {time.perf_counter() - begin:1.1f}s")
 

onnx_diagnostic/torch_export_patches/patches/patch_transformers.py

@@ -1,4 +1,5 @@
 # transformers
+from typing import List
 from .patch_helper import _has_transformers
 
 from ._patch_transformers_attention import (
@@ -86,3 +87,11 @@
 
 
 from ._patch_transformers_sam_mask_decoder import patched_SamMaskDecoder
+
+
+def get_transformers_plugs() -> List["EagerDirectReplacementWithOnnx"]:  # noqa: F821
+    """Returns the necessary plugs to rewrite models."""
+    plugs = []
+    if patch_qwen2_5:
+        plugs.extend(PLUGS_Qwen25)
+    return plugs