opsd_utils/mode_router.py

import torch

from opsd_utils.constants import MODE_GRPO, MODE_OPSD, MODE_SFT
from opsd_utils import debug_log as opsd_debug

# Anti-leakage modes: correct → GRPO, wrong → OPSD/OPD, else SFT
_NO_LEAK_MODES = frozenset({"rlsd", "copsd_opd"})


def route_prompt_modes(
    acc_rewards: torch.Tensor,
    num_generations: int,
    opsd_config: dict,
    recoverable_flags: list[bool],
) -> list[int]:
    """
    Route each prompt (not each completion) to GRPO / OPSD / SFT.

    Args:
        acc_rewards: (num_prompts, num_generations)
        recoverable_flags: length num_prompts
    Returns:
        list[int] of length num_prompts with MODE_* values
    """
    threshold = opsd_config.get("gate", {}).get("correct_threshold", 0.5)
    mode_name = opsd_config.get("mode", "dyme")
    enabled = opsd_config.get("enabled", False)

    num_prompts = acc_rewards.shape[0]
    modes: list[int] = []
    opsd_debug.log(
        "mode_router",
        "route_prompt_modes enter",
        num_prompts=num_prompts,
        num_generations=num_generations,
        mode_name=mode_name,
        enabled=enabled,
        threshold=threshold,
        acc_rewards_shape=tuple(acc_rewards.shape),
        recoverable_flags=recoverable_flags,
    )

    for p in range(num_prompts):
        any_correct = (acc_rewards[p] > threshold).any().item()
        recoverable = recoverable_flags[p] if p < len(recoverable_flags) else False

        if not enabled or mode_name == "dyme":
            selected = MODE_GRPO if any_correct else MODE_SFT
        elif mode_name == "opsd_only":
            selected = MODE_OPSD
        elif mode_name == "replace_sft":
            selected = MODE_GRPO if any_correct else MODE_OPSD
        elif mode_name == "opsd_on_wrong":
            if any_correct:
                selected = MODE_GRPO
            elif recoverable:
                selected = MODE_OPSD
            else:
                selected = MODE_SFT
        elif mode_name == "grpo_opsd_joint":
            selected = MODE_GRPO if any_correct else (MODE_OPSD if recoverable else MODE_SFT)
        elif mode_name in _NO_LEAK_MODES:
            if any_correct:
                selected = MODE_GRPO
            elif recoverable:
                selected = MODE_OPSD
            else:
                selected = MODE_SFT
        else:
            # trimode (legacy/leaky): OPSD on correct; wrong → DyME SFT cold-start
            selected = MODE_OPSD if any_correct else MODE_SFT

        modes.append(selected)
        opsd_debug.log(
            "mode_router",
            "prompt routed",
            prompt_index=p,
            any_correct=any_correct,
            recoverable=recoverable,
            selected_mode=opsd_debug.MODE_NAMES.get(selected, selected),
            acc_rewards_row=acc_rewards[p].tolist(),
        )

    opsd_debug.log_mode_summary("mode_router", modes)
    return modes


def expand_modes_to_completions(prompt_modes: list[int], num_generations: int, batch_size: int) -> list[int]:
    """Map per-prompt mode to per-completion mode."""
    completion_modes = []
    for i in range(batch_size):
        batch_id = i // num_generations
        completion_modes.append(prompt_modes[batch_id])
    opsd_debug.log(
        "mode_router",
        "expand_modes_to_completions",
        batch_size=batch_size,
        num_generations=num_generations,
        completion_modes=[opsd_debug.MODE_NAMES.get(m, m) for m in completion_modes],
    )
    return completion_modes


def route_completion_modes(
    acc_rewards: torch.Tensor,
    num_generations: int,
    batch_size: int,
    opsd_config: dict,
    recoverable_flags: list[bool],
    format_rewards: torch.Tensor | None = None,
) -> list[int]:
    """Route each completion individually (TriMode) or expand per-prompt modes."""
    gate = opsd_config.get("gate", {})
    mode_name = opsd_config.get("mode", "dyme")
    per_completion = gate.get("per_completion_opsd", False)
    threshold = gate.get("correct_threshold", 0.5)
    require_format = gate.get("require_format_for_opsd", False)

    if per_completion and mode_name in _NO_LEAK_MODES:
        num_prompts = acc_rewards.shape[0]
        group_has_correct = [
            (acc_rewards[p] > threshold).any().item() for p in range(num_prompts)
        ]
        online_sft_on_all_wrong = gate.get("online_sft_on_all_wrong", True)

        completion_modes: list[int] = []
        for i in range(batch_size):
            prompt_idx = i // num_generations
            gen_idx = i % num_generations
            acc_ok = acc_rewards[prompt_idx, gen_idx].item() > threshold
            fmt_ok = True
            if require_format and format_rewards is not None:
                fmt_ok = format_rewards[prompt_idx, gen_idx].item() > 0
            recoverable = (
                recoverable_flags[prompt_idx] if recoverable_flags and prompt_idx < len(recoverable_flags) else True
            )
            all_wrong_group = not group_has_correct[prompt_idx]

            if acc_ok and fmt_ok:
                selected = MODE_GRPO
            elif (
                online_sft_on_all_wrong
                and all_wrong_group
                and gen_idx == 0
            ):
                # COPSD path C: online SFT cold-start when entire group is wrong
                selected = MODE_SFT
            elif not acc_ok and recoverable:
                selected = MODE_OPSD
            else:
                selected = MODE_SFT
            completion_modes.append(selected)
        opsd_debug.log(
            "mode_router",
            "route_completion_modes rlsd/copsd per-completion",
            batch_size=batch_size,
            num_generations=num_generations,
            mode_name=mode_name,
            require_format_for_opsd=require_format,
            completion_modes=[opsd_debug.MODE_NAMES.get(m, m) for m in completion_modes],
        )
        return completion_modes

    if mode_name == "trimode" and per_completion:
        completion_modes = []
        for i in range(batch_size):
            prompt_idx = i // num_generations
            gen_idx = i % num_generations
            acc_ok = acc_rewards[prompt_idx, gen_idx].item() > threshold
            fmt_ok = True
            if require_format and format_rewards is not None:
                fmt_ok = format_rewards[prompt_idx, gen_idx].item() > 0
            selected = MODE_OPSD if (acc_ok and fmt_ok) else MODE_SFT
            completion_modes.append(selected)
        opsd_debug.log(
            "mode_router",
            "route_completion_modes trimode per-completion",
            batch_size=batch_size,
            num_generations=num_generations,
            per_completion_opsd=True,
            require_format_for_opsd=require_format,
            completion_modes=[opsd_debug.MODE_NAMES.get(m, m) for m in completion_modes],
        )
        return completion_modes

    prompt_modes = route_prompt_modes(
        acc_rewards, num_generations, opsd_config, recoverable_flags
    )
    return expand_modes_to_completions(prompt_modes, num_generations, batch_size)