File size: 3,003 Bytes
0e8fbd9 | 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 | # test_ras.py: Repetition Aware Sampling.
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"
|