File size: 3,945 Bytes
6e802a4 | 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 138 | # test_s1mini_render.py: The Path B renderer shim.
import sys
from pathlib import Path
import pytest
torch = pytest.importorskip("torch")
REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "scripts"))
from s1mini_render import FULL_VOCAB_MIN, ForcedSemantic
SEM_BEGIN = 100_000
IM_END = 7
def real_sample(logits, temperature, top_p, top_k):
return torch.tensor([int(logits.argmax())]), "REAL"
def wide(n=152_000):
return torch.zeros(1, 1, n)
def narrow(peak):
x = torch.zeros(1, 1, 1024)
x[0, 0, peak] = 1.0
return x
def test_semantic_calls_forced_fast_calls_passthrough():
f = ForcedSemantic(real_sample, SEM_BEGIN, IM_END)
f.arm([11, 22])
t, _ = f(wide(), None, None, None)
assert int(t) == SEM_BEGIN + 11
t2, tag = f(narrow(peak=333), None, None, None)
assert tag == "REAL" and int(t2) == 333
t3, _ = f(wide(), None, None, None)
assert int(t3) == SEM_BEGIN + 11
def test_ras_double_draw_per_frame():
f = ForcedSemantic(real_sample, SEM_BEGIN, IM_END)
f.arm([5, 6, 7])
got = [int(f(wide(), None, None, None)[0]) for _ in range(6)]
assert got == [SEM_BEGIN + 5, SEM_BEGIN + 5,
SEM_BEGIN + 6, SEM_BEGIN + 6,
SEM_BEGIN + 7, SEM_BEGIN + 7]
def test_exhaustion_returns_im_end():
f = ForcedSemantic(real_sample, SEM_BEGIN, IM_END)
f.arm([9])
f(wide(), None, None, None)
f(wide(), None, None, None)
t, _ = f(wide(), None, None, None)
assert int(t) == IM_END
def test_rearm_resets_counter():
f = ForcedSemantic(real_sample, SEM_BEGIN, IM_END)
f.arm([1, 2])
for _ in range(4):
f(wide(), None, None, None)
f.arm([3])
t, _ = f(wide(), None, None, None)
assert int(t) == SEM_BEGIN + 3
def test_fast_calls_masked_to_residual_range():
from s1mini_render import RES_MAX
f = ForcedSemantic(real_sample, SEM_BEGIN, IM_END)
f.arm([1])
x = torch.zeros(1, 1, 4096)
x[0, 0, 2000] = 5.0
x[0, 0, 700] = 1.0
t, tag = f(x, None, None, None)
assert tag == "REAL" and int(t) == 700
t2, _ = f(narrow(peak=RES_MAX), None, None, None)
assert int(t2) == RES_MAX
def test_fast_temps_indexed_per_level_and_counter_always_advances():
seen = []
def spy(logits, temperature, top_p, top_k):
seen.append(float(temperature))
return torch.tensor([0]), "REAL"
temps = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
f = ForcedSemantic(spy, SEM_BEGIN, IM_END, fast_temps=temps)
f.arm([1, 2])
for _ in range(18):
f(narrow(peak=5), None, None, None)
assert seen == pytest.approx(temps + temps, abs=1e-6), seen
def test_fast_temps_default_none_leaves_temperature_untouched():
seen = []
def spy(logits, temperature, top_p, top_k):
seen.append(temperature)
return torch.tensor([0]), "REAL"
f = ForcedSemantic(spy, SEM_BEGIN, IM_END)
f.arm([1])
f(narrow(peak=5), "SENTINEL", None, None)
assert seen == ["SENTINEL"]
def test_returned_probs_is_none_and_only_index_zero_used():
f = ForcedSemantic(real_sample, SEM_BEGIN, IM_END)
f.arm([4])
t, p = f(wide(), None, None, None)
assert p is None and t.shape == (1,) and t.dtype == torch.long
def test_resume_state_keeps_complete_rows_and_drops_torn_or_unseeded(tmp_path):
import json
from s1mini_render import resume_state
out = tmp_path / "rendered.jsonl"
assert resume_state(out) == (set(), [])
lines = [
json.dumps({"renderer": {"model": "m"}}),
json.dumps({"renderer": {"row": 0, "seed": 5}}),
json.dumps({"renderer": {"row": 2, "seed": 7}}),
json.dumps({"renderer": {"row": 2, "seed": 7}}),
'{"renderer": {"row": 3, "se',
]
out.write_text("\n".join(lines), encoding="utf-8")
done, kept = resume_state(out)
assert done == {0, 2}
assert [json.loads(k)["renderer"]["row"] for k in kept] == [0, 2]
|