File size: 1,121 Bytes
2abcc30 | 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 | import torch
from local_train.rl import gold_continuation, group_advantages
def test_gold_continuation_strips_duplicated_think():
prompt = "<|im_start|>assistant\n<think>\n"
completion = "<think>\nedit now\n</think>\n\n```bash\nsed -i 's/a/b/' foo.py\n```\n"
assert gold_continuation(prompt, completion).startswith("edit now")
assert gold_continuation("plain", completion) == completion
def test_group_advantages_prefers_higher_reward_in_pair():
rewards = torch.tensor([0.2, 0.8, 0.1, 0.9])
adv = group_advantages(rewards, num_generations=2)
assert adv[1] > adv[0]
assert adv[3] > adv[2]
def test_group_advantages_falls_back_when_groups_tie():
# Each pair is tied, but the two groups differ — batch baseline should fire.
rewards = torch.tensor([0.1, 0.1, 0.9, 0.9])
adv = group_advantages(rewards, num_generations=2)
assert adv[0] < 0
assert adv[2] > 0
def test_group_advantages_all_equal_is_zero():
rewards = torch.tensor([0.0, 0.0, 0.0, 0.0])
adv = group_advantages(rewards, num_generations=2)
assert torch.allclose(adv, torch.zeros_like(adv))
|