Fix boolean parsing for capture-onnx-graph options

olive/cli/capture_onnx.py

@@ -37,6 +37,19 @@ def parse_dim_dict(s):
         raise argparse.ArgumentTypeError("Format must be key=value,... with positive integers as values") from exc
 
 
+def parse_bool(value):
+    if isinstance(value, bool):
+        return value
+
+    normalized = value.lower()
+    if normalized in {"true", "1", "yes", "on"}:
+        return True
+    if normalized in {"false", "0", "no", "off"}:
+        return False
+
+    raise argparse.ArgumentTypeError(f"invalid boolean value: {value!r}")
+
+
 class CaptureOnnxGraphCommand(BaseOliveCLICommand):
     @staticmethod
     def register_subcommand(parser: ArgumentParser):
@@ -147,21 +160,21 @@ def register_subcommand(parser: ArgumentParser):
         )
         mb_group.add_argument(
             "--exclude_embeds",
-            type=bool,
+            type=parse_bool,
             default=False,
             required=False,
             help="Remove embedding layer from your ONNX model.",
         )
         mb_group.add_argument(
             "--exclude_lm_head",
-            type=bool,
+            type=parse_bool,
             default=False,
             required=False,
             help="Remove language modeling head from your ONNX model.",
         )
         mb_group.add_argument(
             "--enable_cuda_graph",
-            type=bool,
+            type=parse_bool,
             default=None,  # Explicitly setting to None to differentiate between user intent and default.
             required=False,
             help=(

test/cli/test_capture_onnx_args.py

@@ -0,0 +1,49 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import argparse
+
+import pytest
+
+from olive.cli.capture_onnx import CaptureOnnxGraphCommand
+
+
+def _parse_capture_args(*args):
+    parser = argparse.ArgumentParser()
+    commands = parser.add_subparsers()
+    CaptureOnnxGraphCommand.register_subcommand(commands)
+    return parser.parse_args(["capture-onnx-graph", *args])
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        ("true", True),
+        ("1", True),
+        ("yes", True),
+        ("on", True),
+        ("false", False),
+        ("0", False),
+        ("no", False),
+        ("off", False),
+    ],
+)
+def test_capture_onnx_boolean_arguments(value, expected):
+    args = _parse_capture_args(
+        "--exclude_embeds",
+        value,
+        "--exclude_lm_head",
+        value,
+        "--enable_cuda_graph",
+        value,
+    )
+
+    assert args.exclude_embeds is expected
+    assert args.exclude_lm_head is expected
+    assert args.enable_cuda_graph is expected
+
+
+def test_capture_onnx_boolean_arguments_reject_invalid_value():
+    with pytest.raises(SystemExit):
+        _parse_capture_args("--exclude_embeds", "not-a-bool")