File size: 2,698 Bytes
1f71c7d | 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 | """Tests for the benchmark harness (fast, small-D)."""
from __future__ import annotations
import json
import sys
from pathlib import Path
import numpy as np
import pytest
# make bench/ importable
BENCH_DIR = Path(__file__).resolve().parent.parent / "bench"
sys.path.insert(0, str(BENCH_DIR))
import gen_datasets # noqa: E402
import run_benchmark as bench # noqa: E402
from ensemble import Expert # noqa: E402
@pytest.fixture
def small_expert():
text = "the mitochondria is the powerhouse of the cell. " * 20
return Expert.from_text(text, domain="bio", D=2000, seed=0)
class TestDatasets:
def test_gen_facts_deterministic(self):
a = gen_datasets.gen_facts()
b = gen_datasets.gen_facts()
assert a == b
def test_gen_qa_splits_disjoint(self):
train, holdout = gen_datasets.gen_qa_splits()
train_qs = {q for q, _ in train}
holdout_qs = {q for q, _ in holdout}
# holdout questions must NOT appear in train
assert holdout_qs.isdisjoint(train_qs)
def test_gen_all(self, tmp_path):
paths = gen_datasets.generate_all(tmp_path)
assert (tmp_path / "facts.txt").exists()
assert (tmp_path / "qa_train.json").exists()
assert (tmp_path / "qa_holdout.json").exists()
assert (tmp_path / "prose.txt").exists()
class TestMetrics:
def test_next_token_accuracy_in_range(self, small_expert):
acc = bench._next_token_accuracy(small_expert, "the mitochondria is the powerhouse", max_eval=50)
assert 0.0 <= acc <= 1.0
def test_next_token_accuracy_repetitive_high(self, small_expert):
# on the exact training text, accuracy should be decent (memorized)
text = "the mitochondria is the powerhouse of the cell. "
acc = bench._next_token_accuracy(small_expert, text * 5, max_eval=80)
assert acc > 0.3 # repeated text is well learned
def test_ram_mb_positive(self, small_expert):
ram = bench._ram_mb(small_expert)
assert ram > 0
def test_qa_exact_recall(self):
pairs = [("what is x", "x is the answer")] * 4
e = Expert.from_qa_pairs(pairs, domain="t", D=5000, seed=0)
recall, n = bench._qa_exact_recall(e, [("what is x", "x is the answer")])
assert n == 1
assert 0.0 <= recall <= 1.0
def test_query_latency_positive(self, small_expert):
lat = bench._measure_query_latency(small_expert, ["what is"], k=3)
assert lat > 0
class TestScalingPresets:
def test_scales_defined(self):
from ensemble.expert import SCALES
assert "1b" in SCALES
assert SCALES["1b"]["D"] == 100_000
assert SCALES["small"]["D"] == 10_000
|