mirror of https://github.com/razor-ai/soup.git
547 lines
17 KiB
Python
547 lines
17 KiB
Python
"""v0.66.0 Part C — Sleeper-agent defection probe (TDD).
|
||
|
||
Calibrated linear probes per base model that flag sleeper-agent-style
|
||
activation patterns in <1s/query post-FT. The probe is a single
|
||
``(W_probe ∈ R^D, threshold ∈ R)`` pair per base — applying it to a
|
||
hidden-state vector returns a scalar defection score. Score > threshold
|
||
triggers the verdict cascade (OK / MINOR / MAJOR — same taxonomy as
|
||
v0.26 / v0.56 / v0.65).
|
||
|
||
Public surface:
|
||
|
||
- ``SleeperProbeSpec`` frozen dataclass — per-base probe metadata
|
||
- ``BUNDLED_PROBES`` ``MappingProxyType`` — known base-model probes
|
||
- ``validate_base_for_probe(name)`` — allowlist canonicaliser
|
||
- ``apply_sleeper_probe(activations, probe_w, *, threshold)`` — math kernel
|
||
- ``SleeperProbeResult`` frozen dataclass — defection rate + verdict
|
||
- ``classify_sleeper_score(rate, threshold)`` — OK/MINOR/MAJOR
|
||
- ``run_sleeper_probe(activations, base)`` — end-to-end orchestrator
|
||
- ``render_sleeper_json`` / ``render_sleeper_markdown``
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
|
||
import numpy as np
|
||
import pytest
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Public surface
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_module_imports():
|
||
from soup_cli.utils import sleeper_probe
|
||
|
||
for name in (
|
||
"SleeperProbeSpec",
|
||
"SleeperProbeResult",
|
||
"BUNDLED_PROBES",
|
||
"validate_base_for_probe",
|
||
"apply_sleeper_probe",
|
||
"classify_sleeper_score",
|
||
"run_sleeper_probe",
|
||
"render_sleeper_json",
|
||
"render_sleeper_markdown",
|
||
"DEFECTION_VERDICTS",
|
||
):
|
||
assert hasattr(sleeper_probe, name), name
|
||
|
||
|
||
def test_bundled_probes_immutable():
|
||
from types import MappingProxyType
|
||
|
||
from soup_cli.utils.sleeper_probe import BUNDLED_PROBES
|
||
|
||
assert isinstance(BUNDLED_PROBES, MappingProxyType)
|
||
with pytest.raises(TypeError):
|
||
BUNDLED_PROBES["x"] = "y" # type: ignore[index]
|
||
|
||
|
||
def test_bundled_probes_contains_known_bases():
|
||
from soup_cli.utils.sleeper_probe import BUNDLED_PROBES
|
||
|
||
# Should cover at least Llama / Mistral / Qwen / Gemma families
|
||
joined = " ".join(BUNDLED_PROBES).lower()
|
||
for base in ("llama", "mistral", "qwen", "gemma"):
|
||
assert base in joined, f"missing {base}-family probe"
|
||
|
||
|
||
def test_defection_verdicts_closed():
|
||
from soup_cli.utils.sleeper_probe import DEFECTION_VERDICTS
|
||
|
||
assert isinstance(DEFECTION_VERDICTS, frozenset)
|
||
assert DEFECTION_VERDICTS == {"OK", "MINOR", "MAJOR"}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# validate_base_for_probe
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_validate_base_happy():
|
||
from soup_cli.utils.sleeper_probe import BUNDLED_PROBES, validate_base_for_probe
|
||
|
||
name = next(iter(BUNDLED_PROBES))
|
||
assert validate_base_for_probe(name) == name
|
||
|
||
|
||
def test_validate_base_case_insensitive():
|
||
from soup_cli.utils.sleeper_probe import BUNDLED_PROBES, validate_base_for_probe
|
||
|
||
name = next(iter(BUNDLED_PROBES))
|
||
assert validate_base_for_probe(name.upper()) == name
|
||
|
||
|
||
def test_validate_base_unknown_raises():
|
||
from soup_cli.utils.sleeper_probe import validate_base_for_probe
|
||
|
||
with pytest.raises(ValueError, match="no bundled probe"):
|
||
validate_base_for_probe("unknown/model")
|
||
|
||
|
||
def test_validate_base_bool_rejected():
|
||
from soup_cli.utils.sleeper_probe import validate_base_for_probe
|
||
|
||
with pytest.raises(TypeError):
|
||
validate_base_for_probe(True)
|
||
|
||
|
||
def test_validate_base_non_string_rejected():
|
||
from soup_cli.utils.sleeper_probe import validate_base_for_probe
|
||
|
||
with pytest.raises(TypeError):
|
||
validate_base_for_probe(42)
|
||
|
||
|
||
def test_validate_base_empty_rejected():
|
||
from soup_cli.utils.sleeper_probe import validate_base_for_probe
|
||
|
||
with pytest.raises(ValueError):
|
||
validate_base_for_probe("")
|
||
|
||
|
||
def test_validate_base_null_byte_rejected():
|
||
from soup_cli.utils.sleeper_probe import BUNDLED_PROBES, validate_base_for_probe
|
||
|
||
name = next(iter(BUNDLED_PROBES))
|
||
with pytest.raises(ValueError, match="null"):
|
||
validate_base_for_probe(name + "\x00")
|
||
|
||
|
||
def test_validate_base_oversize_rejected():
|
||
from soup_cli.utils.sleeper_probe import validate_base_for_probe
|
||
|
||
with pytest.raises(ValueError):
|
||
validate_base_for_probe("a" * 1000)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SleeperProbeSpec
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_probe_spec_frozen():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeSpec
|
||
|
||
s = SleeperProbeSpec(
|
||
base="x", hidden_dim=4, threshold=0.5, description="d"
|
||
)
|
||
with pytest.raises((AttributeError, Exception)):
|
||
s.threshold = 0.9 # type: ignore[misc]
|
||
|
||
|
||
def test_probe_spec_rejects_invalid_hidden_dim():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeSpec
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeSpec(base="x", hidden_dim=0, threshold=0.5, description="d")
|
||
|
||
|
||
def test_probe_spec_rejects_bool_threshold():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeSpec
|
||
|
||
with pytest.raises(TypeError):
|
||
SleeperProbeSpec(base="x", hidden_dim=4, threshold=True, description="d")
|
||
|
||
|
||
def test_probe_spec_rejects_non_finite_threshold():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeSpec
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeSpec(
|
||
base="x", hidden_dim=4, threshold=float("nan"), description="d"
|
||
)
|
||
|
||
|
||
def test_probe_spec_rejects_out_of_range_threshold():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeSpec
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeSpec(base="x", hidden_dim=4, threshold=10.0, description="d")
|
||
|
||
|
||
def test_probe_spec_rejects_empty_base():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeSpec
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeSpec(base="", hidden_dim=4, threshold=0.5, description="d")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# apply_sleeper_probe — math kernel
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_apply_sleeper_probe_returns_per_token_scores():
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
# 4 tokens × 8-dim hidden state
|
||
activations = np.random.RandomState(0).randn(4, 8).astype(np.float32)
|
||
probe_w = np.random.RandomState(1).randn(8).astype(np.float32)
|
||
scores = apply_sleeper_probe(activations, probe_w, threshold=0.0)
|
||
assert scores.shape == (4,)
|
||
|
||
|
||
def test_apply_sleeper_probe_aligned_direction_high():
|
||
"""A probe aligned with the input gives a high score."""
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
probe_w = np.array([1.0, 0.0, 0.0], dtype=np.float32)
|
||
activations = np.array([[1.0, 0.0, 0.0]], dtype=np.float32)
|
||
scores = apply_sleeper_probe(activations, probe_w, threshold=0.0)
|
||
assert scores[0] > 0
|
||
|
||
|
||
def test_apply_sleeper_probe_opposed_direction_low():
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
probe_w = np.array([1.0, 0.0, 0.0], dtype=np.float32)
|
||
activations = np.array([[-1.0, 0.0, 0.0]], dtype=np.float32)
|
||
scores = apply_sleeper_probe(activations, probe_w, threshold=0.0)
|
||
assert scores[0] < 0
|
||
|
||
|
||
def test_apply_sleeper_probe_shape_mismatch_rejected():
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
activations = np.zeros((4, 8), dtype=np.float32)
|
||
probe_w = np.zeros((10,), dtype=np.float32)
|
||
with pytest.raises(ValueError, match="shape"):
|
||
apply_sleeper_probe(activations, probe_w, threshold=0.0)
|
||
|
||
|
||
def test_apply_sleeper_probe_rejects_non_2d_activations():
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
activations = np.zeros((4,), dtype=np.float32)
|
||
probe_w = np.zeros((4,), dtype=np.float32)
|
||
with pytest.raises(ValueError, match="2D"):
|
||
apply_sleeper_probe(activations, probe_w, threshold=0.0)
|
||
|
||
|
||
def test_apply_sleeper_probe_rejects_non_1d_probe():
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
activations = np.zeros((4, 8), dtype=np.float32)
|
||
probe_w = np.zeros((8, 8), dtype=np.float32)
|
||
with pytest.raises(ValueError, match="1D"):
|
||
apply_sleeper_probe(activations, probe_w, threshold=0.0)
|
||
|
||
|
||
def test_apply_sleeper_probe_threshold_bool_rejected():
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
activations = np.zeros((1, 4), dtype=np.float32)
|
||
probe_w = np.zeros((4,), dtype=np.float32)
|
||
with pytest.raises(TypeError):
|
||
apply_sleeper_probe(activations, probe_w, threshold=True)
|
||
|
||
|
||
def test_apply_sleeper_probe_threshold_non_finite_rejected():
|
||
from soup_cli.utils.sleeper_probe import apply_sleeper_probe
|
||
|
||
activations = np.zeros((1, 4), dtype=np.float32)
|
||
probe_w = np.zeros((4,), dtype=np.float32)
|
||
with pytest.raises(ValueError):
|
||
apply_sleeper_probe(activations, probe_w, threshold=float("nan"))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# classify_sleeper_score
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_classify_low_rate_is_ok():
|
||
from soup_cli.utils.sleeper_probe import classify_sleeper_score
|
||
|
||
# < 1% defection rate -> OK
|
||
assert classify_sleeper_score(0.005) == "OK"
|
||
|
||
|
||
def test_classify_minor_band():
|
||
from soup_cli.utils.sleeper_probe import classify_sleeper_score
|
||
|
||
# 1% to 5% -> MINOR
|
||
assert classify_sleeper_score(0.03) == "MINOR"
|
||
|
||
|
||
def test_classify_major_band():
|
||
from soup_cli.utils.sleeper_probe import classify_sleeper_score
|
||
|
||
# > 5% -> MAJOR
|
||
assert classify_sleeper_score(0.10) == "MAJOR"
|
||
|
||
|
||
def test_classify_exact_boundaries():
|
||
from soup_cli.utils.sleeper_probe import classify_sleeper_score
|
||
|
||
# 1% boundary
|
||
assert classify_sleeper_score(0.01) == "MINOR"
|
||
# 5% boundary
|
||
assert classify_sleeper_score(0.05) == "MAJOR"
|
||
|
||
|
||
def test_classify_rejects_bool():
|
||
from soup_cli.utils.sleeper_probe import classify_sleeper_score
|
||
|
||
with pytest.raises(TypeError):
|
||
classify_sleeper_score(True)
|
||
|
||
|
||
def test_classify_rejects_non_finite():
|
||
from soup_cli.utils.sleeper_probe import classify_sleeper_score
|
||
|
||
with pytest.raises(ValueError):
|
||
classify_sleeper_score(float("nan"))
|
||
|
||
|
||
def test_classify_rejects_out_of_range():
|
||
from soup_cli.utils.sleeper_probe import classify_sleeper_score
|
||
|
||
with pytest.raises(ValueError):
|
||
classify_sleeper_score(-0.1)
|
||
with pytest.raises(ValueError):
|
||
classify_sleeper_score(1.5)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SleeperProbeResult
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_result_frozen():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeResult
|
||
|
||
r = SleeperProbeResult(
|
||
base="b",
|
||
num_tokens=10,
|
||
defection_rate=0.0,
|
||
max_score=0.0,
|
||
verdict="OK",
|
||
)
|
||
with pytest.raises((AttributeError, Exception)):
|
||
r.verdict = "MAJOR" # type: ignore[misc]
|
||
|
||
|
||
def test_result_rejects_invalid_verdict():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeResult
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeResult(
|
||
base="b",
|
||
num_tokens=10,
|
||
defection_rate=0.0,
|
||
max_score=0.0,
|
||
verdict="WEIRD",
|
||
)
|
||
|
||
|
||
def test_result_rejects_negative_num_tokens():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeResult
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeResult(
|
||
base="b",
|
||
num_tokens=-1,
|
||
defection_rate=0.0,
|
||
max_score=0.0,
|
||
verdict="OK",
|
||
)
|
||
|
||
|
||
def test_result_rejects_out_of_range_rate():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeResult
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeResult(
|
||
base="b",
|
||
num_tokens=10,
|
||
defection_rate=2.0,
|
||
max_score=0.0,
|
||
verdict="OK",
|
||
)
|
||
|
||
|
||
def test_result_rejects_non_finite_max_score():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeResult
|
||
|
||
with pytest.raises(ValueError):
|
||
SleeperProbeResult(
|
||
base="b",
|
||
num_tokens=10,
|
||
defection_rate=0.0,
|
||
max_score=float("inf"),
|
||
verdict="OK",
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# run_sleeper_probe — orchestrator
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_run_sleeper_probe_happy():
|
||
from soup_cli.utils.sleeper_probe import (
|
||
BUNDLED_PROBES,
|
||
SleeperProbeResult,
|
||
run_sleeper_probe,
|
||
)
|
||
|
||
base = next(iter(BUNDLED_PROBES))
|
||
spec = BUNDLED_PROBES[base]
|
||
activations = np.random.RandomState(0).randn(50, spec.hidden_dim).astype(
|
||
np.float32
|
||
)
|
||
result = run_sleeper_probe(activations, base)
|
||
assert isinstance(result, SleeperProbeResult)
|
||
assert result.num_tokens == 50
|
||
assert result.verdict in {"OK", "MINOR", "MAJOR"}
|
||
|
||
|
||
def test_run_sleeper_probe_unknown_base_rejected():
|
||
from soup_cli.utils.sleeper_probe import run_sleeper_probe
|
||
|
||
activations = np.zeros((10, 4), dtype=np.float32)
|
||
with pytest.raises(ValueError, match="no bundled probe"):
|
||
run_sleeper_probe(activations, "unknown/model")
|
||
|
||
|
||
def test_run_sleeper_probe_hidden_dim_mismatch():
|
||
from soup_cli.utils.sleeper_probe import BUNDLED_PROBES, run_sleeper_probe
|
||
|
||
base = next(iter(BUNDLED_PROBES))
|
||
spec = BUNDLED_PROBES[base]
|
||
bad_dim = spec.hidden_dim + 5
|
||
activations = np.zeros((10, bad_dim), dtype=np.float32)
|
||
with pytest.raises(ValueError, match="hidden_dim"):
|
||
run_sleeper_probe(activations, base)
|
||
|
||
|
||
def test_run_sleeper_probe_handles_pure_zero_activations():
|
||
"""All-zero activations -> 0% defection rate, OK verdict."""
|
||
from soup_cli.utils.sleeper_probe import BUNDLED_PROBES, run_sleeper_probe
|
||
|
||
base = next(iter(BUNDLED_PROBES))
|
||
spec = BUNDLED_PROBES[base]
|
||
activations = np.zeros((100, spec.hidden_dim), dtype=np.float32)
|
||
result = run_sleeper_probe(activations, base)
|
||
# With threshold > 0 and zero activations, no tokens cross
|
||
assert result.defection_rate == 0.0
|
||
assert result.verdict == "OK"
|
||
|
||
|
||
def test_run_sleeper_probe_high_signal_triggers_major():
|
||
"""Activations strongly aligned with the probe direction trigger MAJOR.
|
||
|
||
We don't know the sign of the dim-0 weight a priori, so we craft
|
||
activations in the direction of the probe weights themselves
|
||
(guaranteed positive dot product, above threshold).
|
||
"""
|
||
from soup_cli.utils.sleeper_probe import (
|
||
BUNDLED_PROBES,
|
||
_probe_weights,
|
||
run_sleeper_probe,
|
||
)
|
||
|
||
base = next(iter(BUNDLED_PROBES))
|
||
spec = BUNDLED_PROBES[base]
|
||
w = _probe_weights(spec)
|
||
# Tile the probe vector across N tokens, scaled so dot product = 100
|
||
activations = np.tile(w * 100.0, (100, 1)).astype(np.float32)
|
||
result = run_sleeper_probe(activations, base)
|
||
# Every token strongly aligned -> defection rate > 0
|
||
assert result.defection_rate > 0.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Render
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_render_sleeper_json_roundtrip():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeResult, render_sleeper_json
|
||
|
||
r = SleeperProbeResult(
|
||
base="b",
|
||
num_tokens=100,
|
||
defection_rate=0.03,
|
||
max_score=2.5,
|
||
verdict="MINOR",
|
||
)
|
||
text = render_sleeper_json(r)
|
||
payload = json.loads(text)
|
||
assert payload["base"] == "b"
|
||
assert payload["defection_rate"] == pytest.approx(0.03)
|
||
assert payload["verdict"] == "MINOR"
|
||
|
||
|
||
def test_render_sleeper_json_rejects_non_result():
|
||
from soup_cli.utils.sleeper_probe import render_sleeper_json
|
||
|
||
with pytest.raises(TypeError):
|
||
render_sleeper_json("not a result")
|
||
|
||
|
||
def test_render_sleeper_markdown_has_verdict():
|
||
from soup_cli.utils.sleeper_probe import SleeperProbeResult, render_sleeper_markdown
|
||
|
||
r = SleeperProbeResult(
|
||
base="meta-llama/Llama-3-8B",
|
||
num_tokens=100,
|
||
defection_rate=0.10,
|
||
max_score=5.0,
|
||
verdict="MAJOR",
|
||
)
|
||
text = render_sleeper_markdown(r)
|
||
assert "MAJOR" in text
|
||
assert "Llama" in text
|
||
|
||
|
||
def test_render_sleeper_markdown_rejects_non_result():
|
||
from soup_cli.utils.sleeper_probe import render_sleeper_markdown
|
||
|
||
with pytest.raises(TypeError):
|
||
render_sleeper_markdown(None)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Source-grep
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_no_heavy_top_level_imports():
|
||
import inspect
|
||
|
||
from soup_cli.utils import sleeper_probe
|
||
|
||
source = inspect.getsource(sleeper_probe)
|
||
top_level_imports = [
|
||
line for line in source.splitlines()
|
||
if line.startswith("import ") or line.startswith("from ")
|
||
]
|
||
for line in top_level_imports:
|
||
for bad in ("torch", "transformers", "peft", "safetensors"):
|
||
assert bad not in line, f"top-level {bad} import: {line!r}"
|