code/tests/test_assembly.py · viai957/align at main

align / code /tests /test_assembly.py

Add files using upload-large-folder tool

bcc8d35 verified 16 days ago

4.97 kB

"""Tests for the rollout assembly seam (fake verifier; no model, no engine)."""

import enum

import pytest
import torch

from align.rl.advantages import GroupFilterStats
from align.rl.assembly import assemble_rollout
from align.rl.rewards import CompositeReward


class _Outcome(enum.Enum):
    CORRECT = 1
    WRONG = 2


def _verify(golds, preds):
    return [
        _Outcome.CORRECT if (p is not None and p.strip() == g) else _Outcome.WRONG
        for g, p in zip(golds, preds)
    ]


def _text(answer: str) -> str:
    return f"Step 1: reason.\n\n# Answer\n\n{answer}"


def _mk(n_prompts=2, g=2, answers=None):
    """Two prompts x G completions with controllable answers."""
    prompt_ids = [[1, 2, 3], [4, 5]]
    completions = [[[10, 11], [12, 13, 14]], [[20], [21, 22]]]
    logprobs = [[[-0.1, -0.2], [-0.3, -0.4, -0.5]], [[-0.6], [-0.7, -0.8]]]
    answers = answers or [["7", "8"], ["9", "9"]]
    texts = [[_text(a) for a in row] for a in [None] for row in answers]
    return prompt_ids, completions, logprobs, texts


def test_assembles_shapes_masks_and_alignment():
    prompt_ids, completions, logprobs, texts = _mk(answers=[["7", "8"], ["9", "1"]])
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["7", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
    )
    assert out is not None
    B, T = out.input_ids.shape
    assert B == 4 and T == 6  # longest row: prompt 3 + completion 3
    # row 0: prompt [1,2,3] + completion [10,11]
    assert out.input_ids[0].tolist() == [1, 2, 3, 10, 11, 0]
    assert out.response_mask[0].tolist() == [0, 0, 0, 1, 1, 0]
    # label-aligned engine logprobs: completion token t -> position len(p)+t-1
    assert out.vllm_logprobs[0, 2].item() == pytest.approx(-0.1)
    assert out.vllm_logprobs[0, 3].item() == pytest.approx(-0.2)
    assert out.vllm_logprobs[0, 1].item() == 0.0  # prompt region zero-filled
    # group ids: two kept groups
    assert out.group_ids.tolist() == [0, 0, 1, 1]
    # advantages: group means subtracted -> each group sums to ~0
    a = out.advantages
    assert abs(float(a[0] + a[1])) < 1e-6 and abs(float(a[2] + a[3])) < 1e-6
    # rewards itemised
    assert set(out.rewards) >= {"reward/accuracy", "reward/format", "reward/total"}
    assert out.rewards["reward/accuracy"].tolist() == [1.0, 0.0, 1.0, 0.0]


def test_dead_groups_dropped_whole_with_telemetry():
    # group 0: both wrong (all_zero-ish variance 0 at total level? totals equal)
    # -> dead; group 1: mixed -> kept
    prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["9", "2"]])
    stats = GroupFilterStats(max_possible_score=1.1)  # acc 1.0 + fmt 0.1
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["7", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
        filter_stats=stats,
    )
    assert out is not None
    assert out.kept_groups == 1 and out.dropped_groups == 1
    assert out.input_ids.shape[0] == 2  # only group 1's two rows
    # both-wrong completions still earn the format term (0.1 each), so the dead
    # group's identical totals are NONZERO -> nonzero_std0, not all_zero. With a
    # welded soft-format reward, all_zero is reserved for unformatted garbage.
    assert stats.nonzero_std0 == 1 and stats.kept == 1


def test_all_dead_returns_none_not_error():
    prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["2", "2"]])
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["7", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
    )
    assert out is None


def test_ragged_logprobs_refused():
    prompt_ids, completions, logprobs, texts = _mk()
    logprobs[0][0] = [-0.1]  # 1 logprob for 2 completion tokens
    with pytest.raises(ValueError, match="ragged|logprobs"):
        assemble_rollout(
            prompt_ids,
            completions,
            logprobs,
            texts,
            golds=["7", "9"],
            reward=CompositeReward(format_weight=0.1),
            verify_batch=_verify,
            pad_token_id=0,
        )


def test_unanswerable_text_scores_zero_accuracy():
    prompt_ids, completions, logprobs, _ = _mk()
    texts = [["no answer marker here", _text("8")], [_text("9"), _text("9")]]
    out = assemble_rollout(
        prompt_ids,
        completions,
        logprobs,
        texts,
        golds=["8", "9"],
        reward=CompositeReward(format_weight=0.1),
        verify_batch=_verify,
        pad_token_id=0,
    )
    assert out is not None
    # group 0 kept (one right, one no-answer); the no-answer row scores acc 0
    row0_acc = out.rewards["reward/accuracy"][0].item()
    assert row0_acc == 0.0

	"""Tests for the rollout assembly seam (fake verifier; no model, no engine)."""

	import enum

	import pytest
	import torch

	from align.rl.advantages import GroupFilterStats
	from align.rl.assembly import assemble_rollout
	from align.rl.rewards import CompositeReward


	class _Outcome(enum.Enum):
	CORRECT = 1
	WRONG = 2


	def _verify(golds, preds):
	return [
	_Outcome.CORRECT if (p is not None and p.strip() == g) else _Outcome.WRONG
	for g, p in zip(golds, preds)
	]


	def _text(answer: str) -> str:
	return f"Step 1: reason.\n\n# Answer\n\n{answer}"


	def _mk(n_prompts=2, g=2, answers=None):
	"""Two prompts x G completions with controllable answers."""
	prompt_ids = [[1, 2, 3], [4, 5]]
	completions = [[[10, 11], [12, 13, 14]], [[20], [21, 22]]]
	logprobs = [[[-0.1, -0.2], [-0.3, -0.4, -0.5]], [[-0.6], [-0.7, -0.8]]]
	answers = answers or [["7", "8"], ["9", "9"]]
	texts = [[_text(a) for a in row] for a in [None] for row in answers]
	return prompt_ids, completions, logprobs, texts


	def test_assembles_shapes_masks_and_alignment():
	prompt_ids, completions, logprobs, texts = _mk(answers=[["7", "8"], ["9", "1"]])
	out = assemble_rollout(
	prompt_ids,
	completions,
	logprobs,
	texts,
	golds=["7", "9"],
	reward=CompositeReward(format_weight=0.1),
	verify_batch=_verify,
	pad_token_id=0,
	)
	assert out is not None
	B, T = out.input_ids.shape
	assert B == 4 and T == 6 # longest row: prompt 3 + completion 3
	# row 0: prompt [1,2,3] + completion [10,11]
	assert out.input_ids[0].tolist() == [1, 2, 3, 10, 11, 0]
	assert out.response_mask[0].tolist() == [0, 0, 0, 1, 1, 0]
	# label-aligned engine logprobs: completion token t -> position len(p)+t-1
	assert out.vllm_logprobs[0, 2].item() == pytest.approx(-0.1)
	assert out.vllm_logprobs[0, 3].item() == pytest.approx(-0.2)
	assert out.vllm_logprobs[0, 1].item() == 0.0 # prompt region zero-filled
	# group ids: two kept groups
	assert out.group_ids.tolist() == [0, 0, 1, 1]
	# advantages: group means subtracted -> each group sums to ~0
	a = out.advantages
	assert abs(float(a[0] + a[1])) < 1e-6 and abs(float(a[2] + a[3])) < 1e-6
	# rewards itemised
	assert set(out.rewards) >= {"reward/accuracy", "reward/format", "reward/total"}
	assert out.rewards["reward/accuracy"].tolist() == [1.0, 0.0, 1.0, 0.0]


	def test_dead_groups_dropped_whole_with_telemetry():
	# group 0: both wrong (all_zero-ish variance 0 at total level? totals equal)
	# -> dead; group 1: mixed -> kept
	prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["9", "2"]])
	stats = GroupFilterStats(max_possible_score=1.1) # acc 1.0 + fmt 0.1
	out = assemble_rollout(
	prompt_ids,
	completions,
	logprobs,
	texts,
	golds=["7", "9"],
	reward=CompositeReward(format_weight=0.1),
	verify_batch=_verify,
	pad_token_id=0,
	filter_stats=stats,
	)
	assert out is not None
	assert out.kept_groups == 1 and out.dropped_groups == 1
	assert out.input_ids.shape[0] == 2 # only group 1's two rows
	# both-wrong completions still earn the format term (0.1 each), so the dead
	# group's identical totals are NONZERO -> nonzero_std0, not all_zero. With a
	# welded soft-format reward, all_zero is reserved for unformatted garbage.
	assert stats.nonzero_std0 == 1 and stats.kept == 1


	def test_all_dead_returns_none_not_error():
	prompt_ids, completions, logprobs, texts = _mk(answers=[["1", "1"], ["2", "2"]])
	out = assemble_rollout(
	prompt_ids,
	completions,
	logprobs,
	texts,
	golds=["7", "9"],
	reward=CompositeReward(format_weight=0.1),
	verify_batch=_verify,
	pad_token_id=0,
	)
	assert out is None


	def test_ragged_logprobs_refused():
	prompt_ids, completions, logprobs, texts = _mk()
	logprobs[0][0] = [-0.1] # 1 logprob for 2 completion tokens
	with pytest.raises(ValueError, match="ragged\|logprobs"):
	assemble_rollout(
	prompt_ids,
	completions,
	logprobs,
	texts,
	golds=["7", "9"],
	reward=CompositeReward(format_weight=0.1),
	verify_batch=_verify,
	pad_token_id=0,
	)


	def test_unanswerable_text_scores_zero_accuracy():
	prompt_ids, completions, logprobs, _ = _mk()
	texts = [["no answer marker here", _text("8")], [_text("9"), _text("9")]]
	out = assemble_rollout(
	prompt_ids,
	completions,
	logprobs,
	texts,
	golds=["8", "9"],
	reward=CompositeReward(format_weight=0.1),
	verify_batch=_verify,
	pad_token_id=0,
	)
	assert out is not None
	# group 0 kept (one right, one no-answer); the no-answer row scores acc 0
	row0_acc = out.rewards["reward/accuracy"][0].item()
	assert row0_acc == 0.0