digital-garbageandClaude Opus 5.5
837a4215a4
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
46 lines
1.7 KiB
Python
46 lines
1.7 KiB
Python
"""The judge: needs enough ratings, learns their order, and its push helps."""
|
|
|
|
import torch
|
|
|
|
from modules.system.taste import value
|
|
|
|
|
|
def _rated(n=12):
|
|
torch.manual_seed(0)
|
|
liked_dir = torch.randn(value.DIM)
|
|
descs, rewards = [], []
|
|
for i in range(n):
|
|
sign = 1.0 if i % 2 else -1.0
|
|
descs.append(sign * liked_dir + 0.3 * torch.randn(value.DIM))
|
|
rewards.append(sign)
|
|
return descs, rewards, liked_dir
|
|
|
|
|
|
def test_too_few_or_all_alike_is_no_judge():
|
|
descs, rewards, _ = _rated(9)
|
|
assert value.Judge.train_on(descs, rewards) is None
|
|
descs, _r, _ = _rated(12)
|
|
assert value.Judge.train_on(descs, [1.0] * 12) is None
|
|
|
|
|
|
def test_its_push_raises_the_score_and_points_at_liked():
|
|
descs, rewards, liked_dir = _rated()
|
|
judge = value.Judge.train_on(descs, rewards)
|
|
video = torch.randn(1, 4, 2, 8, 8)
|
|
moved = judge.nudge(video, 0.05)
|
|
score = lambda v: float(sum(net(value.describe(v).unsqueeze(0)) for net in judge.nets))
|
|
assert score(moved) > score(video)
|
|
shift = value.describe(moved) - value.describe(video)
|
|
assert torch.nn.functional.cosine_similarity(shift, liked_dir, dim=0) > 0.3
|
|
assert torch.allclose((moved - video).norm(), 0.05 * video.norm(), rtol=1e-3)
|
|
|
|
|
|
def test_the_judge_trains_and_steers_under_comfyuis_inference_mode():
|
|
# ComfyUI runs every prompt under torch.inference_mode(): weights made there cannot be trained.
|
|
import torch
|
|
from modules.system.taste import value as v
|
|
with torch.inference_mode():
|
|
judge = v.Judge.train_on([torch.randn(v.DIM) for _ in range(v.MIN_SAMPLES + 2)], [i % 2 for i in range(v.MIN_SAMPLES + 2)])
|
|
assert judge is not None
|
|
assert judge.gradient(torch.randn(1, 4, 2, 8, 8)).shape == (1, 4, 2, 8, 8)
|