mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-14 21:40:47 +02:00
120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
"""Tests for the abstaining source-pipeline classifier research tool."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from click.testing import CliRunner
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts"))
|
|
|
|
import source_pipeline_classifier as classifier
|
|
|
|
|
|
def test_command_help_exits_cleanly() -> None:
|
|
result = CliRunner().invoke(classifier.main, ["--help"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Usage:" in result.output
|
|
assert "train" in result.output
|
|
assert "predict" in result.output
|
|
|
|
|
|
def _training_rows() -> tuple[np.ndarray, np.ndarray]:
|
|
rng = np.random.default_rng(20260915)
|
|
rows = np.vstack(
|
|
(
|
|
rng.normal((-3.0, 0.0, 0.0), 0.1, size=(12, 3)),
|
|
rng.normal((3.0, 0.0, 0.0), 0.1, size=(12, 3)),
|
|
rng.normal((0.0, 3.0, 0.0), 0.1, size=(12, 3)),
|
|
)
|
|
)
|
|
labels = np.repeat(np.arange(3), 12)
|
|
return rows, labels
|
|
|
|
|
|
def test_fit_and_abstention_separate_provider_rows() -> None:
|
|
features, labels = _training_rows()
|
|
model = classifier.fit_ridge(features, labels, penalty=1.0)
|
|
raw_scores = classifier.scores(features, model)
|
|
margin = classifier.calibrate_margin(raw_scores, labels, unknown_fpr=0.01)
|
|
predictions = classifier.predict_scores(raw_scores, margin)
|
|
|
|
assert np.array_equal(predictions[:24], labels[:24])
|
|
assert np.all(predictions[24:] == classifier.UNKNOWN_INDEX)
|
|
|
|
|
|
def test_model_round_trip_is_pickle_free_and_strict(tmp_path: Path) -> None:
|
|
grid = 4
|
|
feature_count = classifier.FEATURE_WIDTH + 3 * grid * grid
|
|
model = classifier.RidgeModel(
|
|
mean=np.arange(feature_count, dtype=np.float64),
|
|
scale=np.full(feature_count, 2.0),
|
|
weights=np.arange(feature_count * 3, dtype=np.float64).reshape(feature_count, 3),
|
|
margin=0.25,
|
|
grid=grid,
|
|
penalty=100.0,
|
|
)
|
|
path = tmp_path / "model.npz"
|
|
|
|
classifier.save_model(path, model)
|
|
loaded = classifier.load_model(path)
|
|
|
|
assert loaded.grid == grid
|
|
assert loaded.margin == model.margin
|
|
assert np.array_equal(loaded.mean, model.mean)
|
|
assert np.array_equal(loaded.scale, model.scale)
|
|
assert np.array_equal(loaded.weights, model.weights)
|
|
|
|
with np.load(path, allow_pickle=False) as payload:
|
|
assert payload.files
|
|
|
|
|
|
def test_model_loader_rejects_extra_fields(tmp_path: Path) -> None:
|
|
path = tmp_path / "model.npz"
|
|
np.savez(
|
|
path,
|
|
classes=np.asarray(classifier.CLASSES),
|
|
mean=np.zeros(136),
|
|
scale=np.ones(136),
|
|
weights=np.zeros((136, 3)),
|
|
margin=np.asarray(0.1),
|
|
spectral_grid=np.asarray(2),
|
|
penalty=np.asarray(100.0),
|
|
unexpected=np.asarray(1),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="model fields must be exactly"):
|
|
classifier.load_model(path)
|
|
|
|
|
|
def test_model_loader_rejects_nonfinite_margin(tmp_path: Path) -> None:
|
|
grid = 2
|
|
feature_count = classifier.FEATURE_WIDTH + 3 * grid * grid
|
|
path = tmp_path / "model.npz"
|
|
np.savez(
|
|
path,
|
|
classes=np.asarray(classifier.CLASSES),
|
|
mean=np.zeros(feature_count),
|
|
scale=np.ones(feature_count),
|
|
weights=np.zeros((feature_count, 3)),
|
|
margin=np.asarray(np.nan),
|
|
spectral_grid=np.asarray(grid),
|
|
penalty=np.asarray(100.0),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="margin must be finite"):
|
|
classifier.load_model(path)
|
|
|
|
|
|
def test_calibration_requires_unknown_rows() -> None:
|
|
with pytest.raises(ValueError, match="must contain unknown"):
|
|
classifier.calibrate_margin(
|
|
np.zeros((2, 3)),
|
|
np.asarray([0, 1]),
|
|
unknown_fpr=0.01,
|
|
)
|