diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index bad1873..1bfce95 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -973,6 +973,12 @@ class AbliterationPipeline: # refusal rate measurement. Default 30 gives ~3.3% resolution; # increase for tighter confidence intervals (reviewer feedback). self.verify_sample_size = verify_sample_size if verify_sample_size is not None else 30 + if refusal_max_tokens is not None and ( + isinstance(refusal_max_tokens, bool) + or not isinstance(refusal_max_tokens, int) + or refusal_max_tokens <= 0 + ): + raise ValueError("refusal_max_tokens must be a positive integer") self.refusal_max_tokens = refusal_max_tokens if refusal_max_tokens is not None else 128 # Large model mode: conservative defaults for 120B+ models. diff --git a/obliteratus/cli.py b/obliteratus/cli.py index c762d61..8b33af6 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -24,6 +24,17 @@ _BANNER = r""" """ +def _positive_int(value: str) -> int: + """Parse a strictly positive integer for public CLI boundaries.""" + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError("must be a positive integer") from exc + if parsed <= 0: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed + + def _add_gpu_args(parser): """Add --gpus flag for multi-GPU control.""" gpu_group = parser.add_argument_group("GPU selection") @@ -241,7 +252,7 @@ def main(argv: list[str] | None = None): "Increase for tighter confidence intervals (e.g. 100 for ~1%% resolution).", ) p.add_argument( - "--refusal-max-tokens", type=int, default=None, + "--refusal-max-tokens", type=_positive_int, default=None, help="Max new tokens to generate per response in the refusal test (default: 128).", ) p.add_argument( diff --git a/tests/test_abliterate.py b/tests/test_abliterate.py index db85022..f1c2ddc 100644 --- a/tests/test_abliterate.py +++ b/tests/test_abliterate.py @@ -180,8 +180,22 @@ class TestPipelineInit: assert pipeline.dtype == "float16" assert pipeline.output_dir == Path("abliterated") assert pipeline.trust_remote_code is False + assert pipeline.refusal_max_tokens == 128 assert pipeline.handle is None + @pytest.mark.parametrize("invalid", [0, -1, 1.5, True]) + def test_refusal_max_tokens_must_be_positive_integer(self, invalid): + with pytest.raises( + ValueError, match="refusal_max_tokens must be a positive integer", + ): + AbliterationPipeline(model_name="test-model", refusal_max_tokens=invalid) + + def test_refusal_max_tokens_accepts_positive_override(self): + pipeline = AbliterationPipeline( + model_name="test-model", refusal_max_tokens=512, + ) + assert pipeline.refusal_max_tokens == 512 + def test_default_method_is_advanced(self): pipeline = AbliterationPipeline(model_name="test-model") assert pipeline.method == "advanced" diff --git a/tests/test_cli_boundaries.py b/tests/test_cli_boundaries.py index d3e8987..406c1f2 100644 --- a/tests/test_cli_boundaries.py +++ b/tests/test_cli_boundaries.py @@ -58,6 +58,25 @@ def test_main_routes_remote_commands(monkeypatch, argv, target): command.assert_called_once() +def test_refusal_max_tokens_cli_default_and_positive_override(monkeypatch): + command = Mock() + monkeypatch.setattr(cli, "_cmd_abliterate", command) + + cli.main(["abliterate", "local/model"]) + assert command.call_args.args[0].refusal_max_tokens is None + + command.reset_mock() + cli.main(["abliterate", "local/model", "--refusal-max-tokens", "512"]) + assert command.call_args.args[0].refusal_max_tokens == 512 + + +@pytest.mark.parametrize("invalid", ["0", "-1", "not-an-integer"]) +def test_refusal_max_tokens_cli_rejects_invalid_values(invalid): + with pytest.raises(SystemExit) as exc: + cli.main(["abliterate", "local/model", "--refusal-max-tokens", invalid]) + assert exc.value.code == 2 + + def test_version_is_stable_and_does_not_dispatch(capsys): from obliteratus import __version__ @@ -311,7 +330,8 @@ def _remote_args(**overrides): "max_layer_fraction": 0.9, "harmless_pc_count": 1, "shield_concept_count": 2, "shield_ridge": 0.1, "shield_residualize": True, "shield_layer_penalty": 0.2, "projection_target": "all", "projection_row_fraction": 0.5, "large_model": True, - "verify_sample_size": 5, "config": "config.yml", "preset": "quick", "methods": ["advanced"], + "verify_sample_size": 5, "refusal_max_tokens": 512, + "config": "config.yml", "preset": "quick", "methods": ["advanced"], "hub_org": "org", "hub_repo": None, "dataset": "builtin", } values.update(overrides) @@ -332,6 +352,13 @@ def test_remote_runner_factory_and_commands(monkeypatch): runner.run_obliterate.return_value = "results" cli._cmd_remote_abliterate(args) assert runner.run_obliterate.call_args.kwargs["projection_row_fraction"] == 0.5 + assert runner.run_obliterate.call_args.kwargs["refusal_max_tokens"] == 512 + + args.refusal_max_tokens = None + runner.run_obliterate.reset_mock() + cli._cmd_remote_abliterate(args) + assert "refusal_max_tokens" not in runner.run_obliterate.call_args.kwargs + runner.run_config.return_value = "results" cli._cmd_remote_run(args) runner.run_tourney.return_value = "results" @@ -399,11 +426,12 @@ def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp min_layer_fraction=0.1, max_layer_fraction=0.9, harmless_pc_count=1, shield_concept_count=2, shield_ridge=0.1, shield_residualize=True, shield_layer_penalty=0.2, projection_target="all", projection_row_fraction=0.5, - quantization=None, large_model=False, verify_sample_size=3, + quantization=None, large_model=False, verify_sample_size=3, refusal_max_tokens=512, residue_file=["audit.json"], dataset="builtin", residue_weight=2, residue_max=4, contribute=True, contribute_notes="fixture", ) cli._cmd_abliterate(args) + assert factory.call_args.kwargs["refusal_max_tokens"] == 512 assert (result_path / "hard_negative_residue.json").is_file() telemetry.assert_called_once_with(pipeline)