File size: 7,349 Bytes
31246a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | #!/usr/bin/env python3
"""Execute the displayed SAE method examples and verify their real outputs.
Examples are extracted from Markdown so the verified code is the code readers
copy. The tiny randomly initialized ViT and synthetic image splits run on CPU;
these checks establish API execution and artifact parity, not semantic quality.
"""
from __future__ import annotations
import argparse
from contextlib import redirect_stdout
from hashlib import sha256
import io
import json
import os
from pathlib import Path
import re
import shutil
import sys
import tempfile
ROOT = Path(__file__).resolve().parents[1]
PAGES = ROOT / "docs" / "methods" / "sae"
EXAMPLE = re.compile(r"<!-- runnable-example -->\s*```python\n(.*?)\n```", re.S)
LINKS = re.compile(r"!?\[[^\]]*\]\(([^\s)]+)\)")
def validate(*, output_dir: Path, write_figures: bool, methods=None):
# Prefer this checkout to a previously installed package.
sys.path.insert(0, str(ROOT / "src"))
os.environ.setdefault("MPLBACKEND", "Agg")
import torch
import transformers
from vlmprobe.sae import list_methods
available = [entry["name"] for entry in list_methods()]
selected = methods or available
if set(selected) - set(available):
raise ValueError("requested method is absent from the native method registry")
output_dir.mkdir(parents=True, exist_ok=True)
record = {"scope": "CPU execution of displayed random-ViT SAE method examples",
"torch": torch.__version__, "transformers": transformers.__version__,
"semantic_quality_evaluated": False, "methods": [], "errors": []}
for name in selected:
page = PAGES / f"{name}.md"
content = page.read_text()
found = EXAMPLE.findall(content)
if len(found) != 1:
raise ValueError(f"{page}: expected exactly one complete runnable example")
source = found[0]
if f"probe.sae.{name}(" not in source:
raise ValueError(f"{page}: example does not use its named method factory")
for forbidden in ("load_snapshot", "from demos.common", "sae.backends", "torch.load("):
if forbidden in source:
raise ValueError(f"{page}: hidden/precomputed example dependency {forbidden}")
for destination in LINKS.findall(content):
if "://" not in destination and not destination.startswith("#"):
target = page.parent / destination.split("#", 1)[0]
# Generated image targets are checked after execution below.
if not target.exists() and not destination.endswith(f"sae-{name}.png"):
raise ValueError(f"{page}: missing relative link {destination}")
old_output = os.environ.get("VLM_PROBING_OUTPUT")
log = io.StringIO()
try:
# Every run is fresh, so documentation verification cannot quietly
# reuse a prior optimizer checkpoint or cached activation result.
with tempfile.TemporaryDirectory(prefix=f"sae-{name}-", dir=output_dir) as run:
run_dir = Path(run)
os.environ["VLM_PROBING_OUTPUT"] = str(run_dir)
namespace = {"__name__": "__main__"}
with redirect_stdout(log):
exec(compile(source, str(page), "exec"), namespace)
result, history = namespace["result"], namespace["history"]
values = result.tensors["features"]
assert values.shape == (16, 32), (name, tuple(values.shape))
assert torch.isfinite(values).all(), name
assert (values != 0).any(), f"{name}: held-out dictionary is entirely inactive"
assert history["completed_steps"] == 24, (name, history["completed_steps"])
assert history["train"]["positions"] == 96, name
assert history["validation"]["positions"] == 32, name
assert history["fingerprints"]["train"] != history["fingerprints"]["validation"], name
assert result.metadata["visual_only"], name
assert result.tensors["positions"][:, 1].tolist() == list(range(1, 17)), name
assert (run_dir / "dictionary" / "model_card.json").is_file(), name
assert (run_dir / "training" / "training_state.pt").is_file(), name
figure = run_dir / "features.png"
assert figure.is_file() and figure.stat().st_size > 1000, name
if write_figures:
target = ROOT / "docs" / "assets" / "tutorials" / f"sae-{name}.png"
target.parent.mkdir(parents=True, exist_ok=True)
shutil.copyfile(figure, target)
row = {"method": name, "success": True,
"code_sha256": sha256(source.encode()).hexdigest(),
"figure_sha256": sha256(figure.read_bytes()).hexdigest(),
"feature_shape": list(values.shape),
"mean_test_l0": result.tensors["l0"].float().mean().item(),
"training_steps": history["completed_steps"],
"training_positions": history["train"]["positions"],
"validation_positions": history["validation"]["positions"],
"validation_mse": history["validation"]["mse"],
"validation_fvu": history["validation"]["fvu"],
"inference_artifact_roundtrip": "identical state; features rtol=1e-5, atol=1e-6"}
record["methods"].append(row)
print(f"PASS {name}: train={row['training_positions']}, validation={row['validation_positions']}, "
f"features={row['feature_shape']}, held-out L0={row['mean_test_l0']:.3f}", flush=True)
except Exception as error:
record["errors"].append({"method": name, "error": repr(error), "stdout": log.getvalue()})
print(f"FAIL {name}: {error}", flush=True)
finally:
if old_output is None:
os.environ.pop("VLM_PROBING_OUTPUT", None)
else:
os.environ["VLM_PROBING_OUTPUT"] = old_output
record["success"] = not record["errors"] and len(record["methods"]) == len(selected)
(output_dir / "validation.json").write_text(json.dumps(record, indent=2) + "\n")
if write_figures and record["success"] and selected == available:
provenance = ROOT / "docs" / "assets" / "tutorials" / "sae-examples-provenance.json"
provenance.write_text(json.dumps(record, indent=2) + "\n")
return record
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--output-dir", type=Path, default=ROOT / "outputs" / "sae_method_guides")
parser.add_argument("--write-figures", action="store_true",
help="Update the documentation figures from the actual executed examples")
parser.add_argument("--methods", nargs="+", help="Optional subset of native method names")
args = parser.parse_args()
record = validate(output_dir=args.output_dir.resolve(), write_figures=args.write_figures, methods=args.methods)
print(f"{len(record['methods'])} passed; {len(record['errors'])} failed")
raise SystemExit(0 if record["success"] else 1)
if __name__ == "__main__":
main()
|