diff options
Diffstat (limited to 'tests/test_eval_classification.py')
| -rw-r--r-- | tests/test_eval_classification.py | 16 |
1 files changed, 9 insertions, 7 deletions
diff --git a/tests/test_eval_classification.py b/tests/test_eval_classification.py index 7114166..2cc217b 100644 --- a/tests/test_eval_classification.py +++ b/tests/test_eval_classification.py @@ -12,7 +12,7 @@ import torch from src.config import Config from src.data import sieve_primes -from src.eval import FLAGGED_SIEVE_PREDS, grokking_signature, halting_report, probe_report +from src.eval import _sieve_rank_signature, grokking_signature, halting_report, probe_report DIGITS_CFG = Config(vocab_mode="digits") EOS = DIGITS_CFG.eos_id @@ -58,18 +58,20 @@ def _easy_only(n): return _perfect(n) if (n % 2 == 0 or n % 5 == 0) else 199 -def test_probe_sieve35_classified_p1(): - """A pure {2,3,5,7} sieve must classify P1: errors are predictions of 121/143/169/187/209.""" +def test_probe_sieve35_classified_p5(): + """A pure {2,3,5,7} sieve must classify P5(4): errors match the rank-4 sieve signature.""" r = probe_report(StubModel(_sieve35), DIGITS_CFG) preds = sorted({e["pred"] for e in r["errors"]}) - assert preds == sorted(FLAGGED_SIEVE_PREDS), r["errors"] + expected = sorted(_sieve_rank_signature(4, 101, 200)) + assert preds == expected, r["errors"] assert len(r["errors"]) == 22 # n in 113..120, 139..142, 167..168, 181..186, 199..200 - assert r["code"] == "P1", r + assert r["code"] == "P5(4)", r + assert r["sieve_rank"] == 4 -def test_probe_perfect_classified_p3(): +def test_probe_perfect_classified_p6(): r = probe_report(StubModel(_perfect), DIGITS_CFG) - assert r["acc"] == 1.0 and r["code"] == "P3", r + assert r["acc"] == 1.0 and r["code"] == "P6", r def test_probe_easy_only_classified_p2(): |