| |
|
|
| import pathlib |
| import sys |
|
|
| import pytest |
| import torch |
|
|
| ROOT = pathlib.Path(__file__).resolve().parent.parent |
| sys.path.insert(0, str(ROOT)) |
| sys.path.insert(0, str(ROOT / "scripts")) |
|
|
| FISH_CB = [4096] + [1024] * 9 |
|
|
|
|
| def _depth(): |
| from model import DepthModule |
| torch.manual_seed(0) |
| return DepthModule(d_model=16, num_levels=10, codebook_size=1024, |
| d_depth=32, num_layers=1, cross_frame=0, |
| codebook_sizes=FISH_CB).eval() |
|
|
|
|
| def test_off_by_default_is_bit_identical(): |
| m = _depth() |
| h = torch.randn(1, 16) |
| hist = [[3] * 10 for _ in range(10)] |
| torch.manual_seed(7) |
| a = m.generate_frame(h, level_temps=[0.7] * 10, rep_window=hist) |
| torch.manual_seed(7) |
| b = m.generate_frame(h, level_temps=[0.7] * 10, rep_window=hist, |
| ras_tau=0.0) |
| assert a == b |
|
|
|
|
| def test_escape_fires_when_history_repeats(): |
| m = _depth() |
| h = torch.randn(1, 16) |
| torch.manual_seed(11) |
| off = m.generate_frame(h, level_temps=[0.7] * 10) |
| hist = [[off[k]] * 10 for k in range(10)] |
|
|
| torch.manual_seed(11) |
| on = m.generate_frame(h, level_temps=[0.7] * 10, rep_window=hist, |
| ras_tau=0.5, ras_win=10) |
| assert on != off, "escape never fired on a fully repeating history" |
|
|
| torch.manual_seed(11) |
| ctrl = m.generate_frame(h, level_temps=[0.7] * 10, rep_window=hist, |
| ras_tau=0.0) |
| assert ctrl == off, "history alone changed the draw (penalty leaked?)" |
|
|
|
|
| def test_escape_ignores_the_repetition_penalty(): |
| import model as M |
| src = pathlib.Path(M.__file__).read_text(encoding="utf-8") |
| i_snap = src.index("base_logits = logits.clone()") |
| i_pen = src.index("logits[idx] = torch.where(s < 0, s * rp, s / rp)") |
| assert i_snap < i_pen, "base_logits snapshot taken AFTER the penalty" |
| assert "F.softmax(base_logits / t" in src, "escape does not use base_logits" |
|
|
|
|
| def test_composes_with_the_penalty(): |
| m = _depth() |
| hist = [[5] * 10 for _ in range(10)] |
| per = [1.2] * 10 |
| per[1] = 1.6 |
| out = m.generate_frame(torch.randn(1, 16), level_temps=[0.7] * 10, |
| rep_window=hist, rep_penalty=per, |
| ras_tau=0.1, ras_win=10) |
| assert len(out) == 10 |
| assert all(0 <= c < s for c, s in zip(out, FISH_CB)) |
|
|
|
|
| def test_per_level_tau_is_indexed_by_level(): |
| m = _depth() |
| hist = [[9] * 10 for _ in range(10)] |
| tau = [0.0] * 10 |
| tau[1] = 0.1 |
| out = m.generate_frame(torch.randn(1, 16), level_temps=[0.7] * 10, |
| rep_window=hist, ras_tau=tau) |
| assert len(out) == 10 |
|
|
|
|
| def test_greedy_path_is_untouched(): |
| m = _depth() |
| h = torch.randn(1, 16) |
| a = m.generate_frame(h, level_temps=[0.0] * 10) |
| hist = [[a[k]] * 10 for k in range(10)] |
| b = m.generate_frame(h, level_temps=[0.0] * 10, rep_window=hist, |
| ras_tau=0.9) |
| assert a == b, "RAS altered a greedy decode" |
|
|
