From 5fe46d2dcae17d66f5b89769a453f2820bc5fca8 Mon Sep 17 00:00:00 2001 From: Raja Mukerji Date: Mon, 7 Sep 2026 09:43:39 -0700 Subject: [PATCH] feat(cli): expose --layer-selection so layer strategy is reachable without the Python API --- obliteratus/cli.py | 17 ++++ obliteratus/remote.py | 2 + tests/test_cli_layer_selection.py | 125 ++++++++++++++++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 tests/test_cli_layer_selection.py diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 15feb5a..4b28faa 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -325,6 +325,15 @@ def main(argv: list[str] | None = None): ) p.add_argument("--regularization", type=float, default=None, help="Override: fraction to preserve (0.0-1.0)") p.add_argument("--refinement-passes", type=int, default=None, help="Override: number of iterative passes") + p.add_argument( + "--layer-selection", type=str, default=None, + choices=["knee_cosmic", "knee", "all", "all_except_first", "middle60", "top_k"], + help=( + "Override which layers the method selects. Defaults to the " + "method's own setting (knee_cosmic for most). 'all' matches the " + "Heretic method and lets per-layer weights do the narrowing." + ), + ) p.add_argument( "--min-layer-fraction", type=float, default=None, help="Optional layer floor as fraction of depth; e.g. 0.75 keeps only the final quarter.", @@ -475,6 +484,10 @@ def main(argv: list[str] | None = None): si_parser.add_argument("--n-directions", type=int, default=None) si_parser.add_argument("--regularization", type=float, default=None) si_parser.add_argument("--refinement-passes", type=int, default=None) + si_parser.add_argument( + "--layer-selection", type=str, default=None, + choices=["knee_cosmic", "knee", "all", "all_except_first", "middle60", "top_k"], + ) si_parser.add_argument("--min-layer-fraction", type=float, default=None) si_parser.add_argument("--max-layer-fraction", type=float, default=None) si_parser.add_argument("--harmless-pc-count", type=int, default=None) @@ -1001,6 +1014,7 @@ def _cmd_self_improve(args): direction_method=args.direction_method, regularization=regularization, refinement_passes=refinement_passes, + layer_selection=getattr(args, "layer_selection", None), min_layer_fraction=args.min_layer_fraction, max_layer_fraction=args.max_layer_fraction, harmless_pc_count=args.harmless_pc_count, @@ -1499,6 +1513,7 @@ def _cmd_abliterate(args): direction_method=getattr(args, "direction_method", None), regularization=args.regularization, refinement_passes=args.refinement_passes, + layer_selection=getattr(args, "layer_selection", None), min_layer_fraction=getattr(args, "min_layer_fraction", None), max_layer_fraction=getattr(args, "max_layer_fraction", None), harmless_pc_count=getattr(args, "harmless_pc_count", None), @@ -1825,6 +1840,8 @@ def _cmd_remote_abliterate(args): kwargs["regularization"] = args.regularization if args.refinement_passes is not None: kwargs["refinement_passes"] = args.refinement_passes + if getattr(args, "layer_selection", None) is not None: + kwargs["layer_selection"] = args.layer_selection if getattr(args, "min_layer_fraction", None) is not None: kwargs["min_layer_fraction"] = args.min_layer_fraction if getattr(args, "max_layer_fraction", None) is not None: diff --git a/obliteratus/remote.py b/obliteratus/remote.py index 96a9f7e..bcdd325 100644 --- a/obliteratus/remote.py +++ b/obliteratus/remote.py @@ -312,6 +312,7 @@ class RemoteRunner: refinement_passes: int | None = None, large_model: bool = False, verify_sample_size: int | None = None, + layer_selection: str | None = None, min_layer_fraction: float | None = None, max_layer_fraction: float | None = None, harmless_pc_count: int | None = None, @@ -351,6 +352,7 @@ class RemoteRunner: if verify_sample_size is not None: parts.extend(["--verify-sample-size", str(verify_sample_size)]) optional_values = ( + ("--layer-selection", layer_selection), ("--min-layer-fraction", min_layer_fraction), ("--max-layer-fraction", max_layer_fraction), ("--harmless-pc-count", harmless_pc_count), diff --git a/tests/test_cli_layer_selection.py b/tests/test_cli_layer_selection.py new file mode 100644 index 0000000..360142e --- /dev/null +++ b/tests/test_cli_layer_selection.py @@ -0,0 +1,125 @@ +"""CPU-safe contract tests for the ``--layer-selection`` CLI override. + +``AbliterationPipeline`` has accepted ``layer_selection`` since layer selection +was made configurable, and ``_distill`` dispatches on six distinct values, but no +command exposed it — reaching anything other than a method's built-in default +required constructing the pipeline in Python. These tests pin the flag to the +values the implementation actually dispatches on, so the two cannot drift. +""" + +from __future__ import annotations + +import re +import shlex + +import pytest + +from obliteratus import cli +from obliteratus.remote import RemoteConfig, RemoteRunner + +# The values `AbliterationPipeline._distill` branches on. "knee_cosmic" is the +# implicit default (the `else` arm) and is spelled out here so selecting it +# explicitly is possible rather than only reachable by omission. +EXPECTED_CHOICES = {"knee_cosmic", "knee", "all", "all_except_first", "middle60", "top_k"} + + +@pytest.mark.parametrize("command", ["obliterate", "abliterate"]) +def test_layer_selection_rejects_an_unknown_strategy(command, monkeypatch): + """An unrecognised value must fail at parse time, not fall through silently. + + `_distill` treats every unknown value as the default `knee_cosmic` arm, so + without constrained choices a typo would run a different strategy than the + one asked for and report success. + """ + monkeypatch.setattr(cli, "_cmd_abliterate", lambda _args: None) + + with pytest.raises(SystemExit) as excinfo: + cli.main([command, "org/model", "--layer-selection", "not-a-strategy"]) + + assert excinfo.value.code == 2 + + +@pytest.mark.parametrize("command", ["obliterate", "abliterate"]) +@pytest.mark.parametrize("strategy", sorted(EXPECTED_CHOICES)) +def test_layer_selection_parses_and_reaches_the_command(command, strategy, monkeypatch): + captured = {} + monkeypatch.setattr(cli, "_cmd_abliterate", lambda args: captured.update(vars(args))) + + cli.main([command, "org/model", "--layer-selection", strategy]) + + assert captured["layer_selection"] == strategy + + +def test_layer_selection_defaults_to_none_so_the_method_keeps_its_own_setting(monkeypatch): + """Omitting the flag must not override the method's configured strategy. + + The pipeline resolves `layer_selection or method_cfg[...]`, so passing + anything other than None here would silently flatten every method's default. + """ + captured = {} + monkeypatch.setattr(cli, "_cmd_abliterate", lambda args: captured.update(vars(args))) + + cli.main(["obliterate", "org/model"]) + + assert captured["layer_selection"] is None + + +def test_self_improve_accepts_the_same_strategies(monkeypatch, tmp_path): + captured = {} + monkeypatch.setattr(cli, "_cmd_self_improve", lambda args: captured.update(vars(args))) + + cli.main([ + "self-improve", "org/model", + "--audit", str(tmp_path / "audit.json"), + "--output-dir", str(tmp_path / "out"), + "--layer-selection", "all", + ]) + + assert captured["layer_selection"] == "all" + + +def test_remote_command_forwards_layer_selection(): + runner = RemoteRunner(RemoteConfig(host="example.invalid", user="operator")) + + command = runner.build_obliterate_command( + "org/model", + method="optimized", + layer_selection="all", + ) + + tokens = shlex.split(command) + assert "--layer-selection" in tokens + assert tokens[tokens.index("--layer-selection") + 1] == "all" + + +def test_remote_command_omits_the_flag_when_unset(): + runner = RemoteRunner(RemoteConfig(host="example.invalid", user="operator")) + + command = runner.build_obliterate_command("org/model", method="optimized") + + assert "--layer-selection" not in shlex.split(command) + + +def test_every_offered_strategy_is_one_distill_actually_dispatches_on(): + """Guard against the flag and the implementation drifting apart. + + A choice the CLI offers but `_distill` does not branch on would fall into the + default arm and silently run `knee_cosmic` instead — the failure this flag + exists to make impossible. + """ + import inspect + + from obliteratus.abliterate import AbliterationPipeline + + source = inspect.getsource(AbliterationPipeline._distill) + + # knee_cosmic is the implicit `else` arm and so has no equality branch. + for strategy in sorted(EXPECTED_CHOICES - {"knee_cosmic"}): + assert f'selection_method == "{strategy}"' in source, ( + f"CLI offers {strategy!r} but _distill has no branch for it" + ) + + branches = set(re.findall(r'selection_method == "([a-z_0-9]+)"', source)) + assert branches | {"knee_cosmic"} == EXPECTED_CHOICES, ( + f"_distill dispatches on {branches | {'knee_cosmic'}}, CLI offers {EXPECTED_CHOICES}" + )