scripts/helpfulness_scheme_interface.py

import json
import math
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn.functional as F
from peft import AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer


ROOT = Path("/scratch/users/nus/e1538883/genarm-exp/temp/pku_safe_rlhf_repro")
FIRST_ENTRIES = ROOT / "gpt-evaluation" / "analysis_helpfulness_len_norm" / "dataset_first_entries.json"
BASE_MODEL = "/scratch/users/nus/e1538883/huggingface/alpaca-7b-reproduced"
OLD_HELPFUL = ROOT / "checkpoints" / "helpfulness" / "arm_beta_0p5" / "final_checkpoint"
NEW_LEN_NORM_HELPFUL = ROOT / "checkpoints" / "helpfulness" / "arm_beta_0p5_len_norm_only" / "final_checkpoint"


@dataclass
class SchemeAConfig:
    base_prob_threshold: float = 0.01
    topk: int = 10
    steps: int = 5
    prompt: str | None = None
    response: str | None = None
    adapter_path: str | None = None


@dataclass
class SchemeBConfig:
    topk: int = 10
    steps: int = 5
    prompt: str | None = None
    response: str | None = None
    adapter_path: str | None = None


@dataclass
class PrefixOnlyConfig:
    topk: int = 10
    steps: int = 25
    prompt: str | None = None
    response: str | None = None
    adapter_path: str | None = None


@dataclass
class X1Config:
    tau_p: float = 0.8
    tau_g: float = 0.5


@dataclass
class SchemeBTempSupportConfig:
    support_temperature: float = 2.0
    alpha: float = 1.0
    topk: int = 10
    steps: int = 25
    prompt: str | None = None
    response: str | None = None
    adapter_path: str | None = None


def load_first_helpfulness_winner_pair() -> dict:
    obj = json.load(open(FIRST_ENTRIES, "r", encoding="utf-8"))
    return obj["helpfulness_training_first_formatted"]


def load_base_and_helpful_models(adapter_path: str | Path | None = None):
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    base_model.eval()

    arm_model = AutoPeftModelForCausalLM.from_pretrained(
        str(adapter_path or OLD_HELPFUL),
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    arm_model.eval()
    return tokenizer, base_model, arm_model


def load_base_and_old_helpful_models():
    return load_base_and_helpful_models(OLD_HELPFUL)


def _resolve_prompt_and_response(prompt: str | None, response: str | None) -> tuple[str, str]:
    if prompt is None or response is None:
        row = load_first_helpfulness_winner_pair()
        return row["prompt"], row["chosen"]
    return prompt, response


def _forward_last_logprobs(model, input_ids: list[int]) -> torch.Tensor:
    device = next(model.parameters()).device
    tensor = torch.tensor([input_ids], device=device)
    with torch.no_grad():
        logits = model(input_ids=tensor, use_cache=False).logits[0, -1].float().cpu()
    return F.log_softmax(logits, dim=-1)


def _top_rows_from_allowed(
    values: torch.Tensor,
    allowed_ids: torch.Tensor,
    tokenizer,
    gold_id: int,
    topk: int,
) -> list[dict]:
    allowed_vals = values[allowed_ids]
    k = min(topk, allowed_ids.numel())
    top_vals, top_idx = torch.topk(allowed_vals, k=k)
    top_ids = allowed_ids[top_idx]
    rows = []
    seen_gold = False
    for tok_id, val in zip(top_ids.tolist(), top_vals.tolist()):
        is_gold = int(tok_id) == int(gold_id)
        seen_gold = seen_gold or is_gold
        rows.append(
            {
                "token_id": int(tok_id),
                "token_text": tokenizer.decode([int(tok_id)]),
                "value": float(val),
                "is_gold": bool(is_gold),
                "gold_in_scope": True,
            }
        )
    if not seen_gold:
        rows.append(
            {
                "token_id": int(gold_id),
                "token_text": tokenizer.decode([int(gold_id)]),
                "value": float(values[int(gold_id)].item()),
                "is_gold": True,
                "gold_in_scope": False,
            }
        )
    return rows


def _score_rows_from_allowed(
    values: torch.Tensor,
    allowed_ids: torch.Tensor,
    tokenizer,
    gold_id: int,
    topk: int,
) -> list[dict]:
    allowed_vals = values[allowed_ids]
    k = min(topk, allowed_ids.numel())
    top_vals, top_idx = torch.topk(allowed_vals, k=k)
    top_ids = allowed_ids[top_idx]
    rows = []
    seen_gold = False
    for tok_id, val in zip(top_ids.tolist(), top_vals.tolist()):
        is_gold = int(tok_id) == int(gold_id)
        seen_gold = seen_gold or is_gold
        rows.append(
            {
                "token_id": int(tok_id),
                "token_text": tokenizer.decode([int(tok_id)]),
                "value": float(val),
                "is_gold": bool(is_gold),
                "gold_in_scope": True,
            }
        )
    if not seen_gold:
        rows.append(
            {
                "token_id": int(gold_id),
                "token_text": tokenizer.decode([int(gold_id)]),
                "value": float(values[int(gold_id)].item()),
                "is_gold": True,
                "gold_in_scope": False,
            }
        )
    return rows


def _top_rows_full_vocab(
    values: torch.Tensor,
    tokenizer,
    gold_id: int,
    topk: int,
) -> list[dict]:
    top_vals, top_ids = torch.topk(values, k=min(topk, values.numel()))
    rows = []
    seen_gold = False
    for tok_id, val in zip(top_ids.tolist(), top_vals.tolist()):
        is_gold = int(tok_id) == int(gold_id)
        seen_gold = seen_gold or is_gold
        rows.append(
            {
                "token_id": int(tok_id),
                "token_text": tokenizer.decode([int(tok_id)]),
                "value": float(val),
                "is_gold": bool(is_gold),
                "gold_in_scope": True,
            }
        )
    if not seen_gold:
        rows.append(
            {
                "token_id": int(gold_id),
                "token_text": tokenizer.decode([int(gold_id)]),
                "value": float(values[int(gold_id)].item()),
                "is_gold": True,
                "gold_in_scope": True,
            }
        )
    return rows


def _sparsemax(scores: torch.Tensor) -> torch.Tensor:
    z = scores.float()
    z_sorted, _ = torch.sort(z, descending=True)
    z_cumsum = torch.cumsum(z_sorted, dim=0)
    ks = torch.arange(1, z.numel() + 1, device=z.device, dtype=z.dtype)
    support = 1 + ks * z_sorted > z_cumsum
    k = int(torch.sum(support).item())
    tau = (z_cumsum[k - 1] - 1) / k
    return torch.clamp(z - tau, min=0.0)


def _compute_allowed_ids_scheme_a(prob_base: torch.Tensor, threshold: float) -> torch.Tensor:
    allowed_mask = prob_base >= threshold
    allowed_ids = torch.nonzero(allowed_mask, as_tuple=False).squeeze(-1)
    if allowed_ids.numel() == 0:
        allowed_ids = torch.tensor([int(torch.argmax(prob_base).item())], dtype=torch.long)
    return allowed_ids


def _compute_allowed_ids_scheme_b(logp_base: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    base_sparse_probs = _sparsemax(logp_base)
    allowed_mask = base_sparse_probs > 0
    allowed_ids = torch.nonzero(allowed_mask, as_tuple=False).squeeze(-1)
    if allowed_ids.numel() == 0:
        allowed_ids = torch.tensor([int(torch.argmax(logp_base).item())], dtype=torch.long)
    return allowed_ids, base_sparse_probs


def _compute_allowed_ids_scheme_b_temp(logp_base: torch.Tensor, support_temperature: float) -> tuple[torch.Tensor, torch.Tensor]:
    if support_temperature <= 0:
        raise ValueError(f"support_temperature must be > 0, got {support_temperature}")
    scaled_base_sparse_probs = _sparsemax(logp_base / support_temperature)
    allowed_mask = scaled_base_sparse_probs > 0
    allowed_ids = torch.nonzero(allowed_mask, as_tuple=False).squeeze(-1)
    if allowed_ids.numel() == 0:
        allowed_ids = torch.tensor([int(torch.argmax(logp_base).item())], dtype=torch.long)
    return allowed_ids, scaled_base_sparse_probs


def _compute_x1_gate(prob_base: torch.Tensor, x1_config: X1Config) -> dict:
    top_vals, top_ids = torch.topk(prob_base, k=2)
    p1 = float(top_vals[0].item())
    p2 = float(top_vals[1].item())
    non_forking = (p1 >= x1_config.tau_p) or ((p1 - p2) >= x1_config.tau_g)
    return {
        "non_forking": bool(non_forking),
        "p1": p1,
        "p2": p2,
        "gap": p1 - p2,
        "top1_token_id": int(top_ids[0].item()),
        "top2_token_id": int(top_ids[1].item()),
    }


def _build_score_with_optional_x1(
    logp_base: torch.Tensor,
    logp_arm: torch.Tensor,
    prob_base: torch.Tensor,
    x1_config: X1Config | None = None,
) -> tuple[torch.Tensor, dict]:
    if x1_config is None:
        return logp_base + logp_arm, {
            "x1_enabled": False,
            "non_forking": False,
            "score_mode": "base_plus_arm",
        }

    gate = _compute_x1_gate(prob_base, x1_config)
    if gate["non_forking"]:
        score_s = logp_base
        score_mode = "base_only_non_forking"
    else:
        score_s = logp_base + logp_arm
        score_mode = "base_plus_arm_forking"
    gate.update(
        {
            "x1_enabled": True,
            "score_mode": score_mode,
            "tau_p": x1_config.tau_p,
            "tau_g": x1_config.tau_g,
        }
    )
    return score_s, gate


def scheme_a_supervised_old_helpful(config: SchemeAConfig) -> dict:
    tokenizer, base_model, arm_model = load_base_and_helpful_models(config.adapter_path)
    prompt, response = _resolve_prompt_and_response(config.prompt, config.response)

    prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    response_ids = tokenizer(response, add_special_tokens=False)["input_ids"]

    steps = []
    for step_idx, gold_id in enumerate(response_ids[: config.steps]):
        prefix_ids = response_ids[:step_idx]
        full_prefix = prompt_ids + prefix_ids

        logp_base = _forward_last_logprobs(base_model, full_prefix)
        logp_arm = _forward_last_logprobs(arm_model, full_prefix)
        prob_base = torch.exp(logp_base)

        allowed_mask = prob_base >= config.base_prob_threshold
        allowed_ids = torch.nonzero(allowed_mask, as_tuple=False).squeeze(-1)
        if allowed_ids.numel() == 0:
            allowed_ids = torch.tensor([int(torch.argmax(prob_base).item())], dtype=torch.long)

        score_s = logp_base + logp_arm
        filtered_score = score_s[allowed_ids]
        filtered_probs = torch.softmax(filtered_score, dim=-1)
        prob_arm = torch.exp(logp_arm)

        if int(gold_id) in set(allowed_ids.tolist()):
            gold_allowed_index = (allowed_ids == int(gold_id)).nonzero(as_tuple=False).item()
            gold_final_prob = float(filtered_probs[gold_allowed_index].item())
            gold_in_scope = True
        else:
            gold_final_prob = 0.0
            gold_in_scope = False

        prob_rows = []
        top_prob_vals, top_prob_idx = torch.topk(filtered_probs, k=min(config.topk, allowed_ids.numel()))
        top_prob_ids = allowed_ids[top_prob_idx]
        seen_gold_prob = False
        for tok_id, val in zip(top_prob_ids.tolist(), top_prob_vals.tolist()):
            is_gold = int(tok_id) == int(gold_id)
            seen_gold_prob = seen_gold_prob or is_gold
            prob_rows.append(
                {
                    "token_id": int(tok_id),
                    "token_text": tokenizer.decode([int(tok_id)]),
                    "value": float(val),
                    "is_gold": bool(is_gold),
                    "gold_in_scope": True,
                }
            )
        if not seen_gold_prob:
            prob_rows.append(
                {
                    "token_id": int(gold_id),
                    "token_text": tokenizer.decode([int(gold_id)]),
                    "value": gold_final_prob,
                    "is_gold": True,
                    "gold_in_scope": False,
                }
            )
        score_rows = _score_rows_from_allowed(score_s, allowed_ids, tokenizer, int(gold_id), config.topk)

        steps.append(
            {
                "step_index": step_idx + 1,
                "gold_token_id": int(gold_id),
                "gold_token_text": tokenizer.decode([int(gold_id)]),
                "prefix_response_token_ids": prefix_ids,
                "prefix_response_text": tokenizer.decode(prefix_ids),
                "allowed_token_count": int(allowed_ids.numel()),
                "gold_token_base_prob": float(prob_base[int(gold_id)].item()),
                "gold_token_in_scope": gold_in_scope,
                "gold_logp_base": float(logp_base[int(gold_id)].item()),
                "gold_logp_arm": float(logp_arm[int(gold_id)].item()),
                "gold_score_s": float(score_s[int(gold_id)].item()),
                "gold_prob_arm": float(prob_arm[int(gold_id)].item()),
                "gold_final_prob_after_filter_softmax": gold_final_prob,
                "logp_base_top10_allowed": _top_rows_from_allowed(logp_base, allowed_ids, tokenizer, int(gold_id), config.topk),
                "logp_arm_top10_allowed": _top_rows_from_allowed(logp_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
                "score_s_top10_allowed": score_rows,
                "prob_base_top10_allowed": _top_rows_from_allowed(prob_base, allowed_ids, tokenizer, int(gold_id), config.topk),
                "prob_arm_top10_allowed": _top_rows_from_allowed(prob_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
                "prob_score_top10_allowed": prob_rows,
            }
        )

    return {
        "scheme": "A",
        "threshold": config.base_prob_threshold,
        "prompt": prompt,
        "response": response,
        "response_token_count": len(response_ids),
        "steps": steps,
    }


def scheme_b_supervised_old_helpful(config: SchemeBConfig) -> dict:
    tokenizer, base_model, arm_model = load_base_and_helpful_models(config.adapter_path)
    prompt, response = _resolve_prompt_and_response(config.prompt, config.response)

    prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    response_ids = tokenizer(response, add_special_tokens=False)["input_ids"]

    steps = []
    for step_idx, gold_id in enumerate(response_ids[: config.steps]):
        prefix_ids = response_ids[:step_idx]
        full_prefix = prompt_ids + prefix_ids

        logp_base = _forward_last_logprobs(base_model, full_prefix)
        logp_arm = _forward_last_logprobs(arm_model, full_prefix)
        score_s = logp_base + logp_arm
        prob_base = torch.exp(logp_base)
        prob_arm = torch.exp(logp_arm)

        base_sparse_probs = _sparsemax(logp_base)
        allowed_mask = base_sparse_probs > 0
        allowed_ids = torch.nonzero(allowed_mask, as_tuple=False).squeeze(-1)
        if allowed_ids.numel() == 0:
            allowed_ids = torch.tensor([int(torch.argmax(logp_base).item())], dtype=torch.long)

        filtered_score = score_s[allowed_ids]
        filtered_probs = torch.softmax(filtered_score, dim=-1)

        if int(gold_id) in set(allowed_ids.tolist()):
            gold_allowed_index = (allowed_ids == int(gold_id)).nonzero(as_tuple=False).item()
            gold_final_prob = float(filtered_probs[gold_allowed_index].item())
            gold_in_scope = True
        else:
            gold_final_prob = 0.0
            gold_in_scope = False

        prob_rows = []
        top_prob_vals, top_prob_idx = torch.topk(filtered_probs, k=min(config.topk, allowed_ids.numel()))
        top_prob_ids = allowed_ids[top_prob_idx]
        seen_gold_prob = False
        for tok_id, val in zip(top_prob_ids.tolist(), top_prob_vals.tolist()):
            is_gold = int(tok_id) == int(gold_id)
            seen_gold_prob = seen_gold_prob or is_gold
            prob_rows.append(
                {
                    "token_id": int(tok_id),
                    "token_text": tokenizer.decode([int(tok_id)]),
                    "value": float(val),
                    "is_gold": bool(is_gold),
                    "gold_in_scope": True,
                }
            )
        if not seen_gold_prob:
            prob_rows.append(
                {
                    "token_id": int(gold_id),
                    "token_text": tokenizer.decode([int(gold_id)]),
                    "value": gold_final_prob,
                    "is_gold": True,
                    "gold_in_scope": False,
                }
            )
        score_rows = _score_rows_from_allowed(score_s, allowed_ids, tokenizer, int(gold_id), config.topk)

        steps.append(
            {
                "step_index": step_idx + 1,
                "gold_token_id": int(gold_id),
                "gold_token_text": tokenizer.decode([int(gold_id)]),
                "prefix_response_token_ids": prefix_ids,
                "prefix_response_text": tokenizer.decode(prefix_ids),
                "allowed_token_count": int(allowed_ids.numel()),
                "gold_token_sparsemax_base_weight": float(base_sparse_probs[int(gold_id)].item()),
                "gold_token_in_scope": gold_in_scope,
                "gold_logp_base": float(logp_base[int(gold_id)].item()),
                "gold_logp_arm": float(logp_arm[int(gold_id)].item()),
                "gold_score_s": float(score_s[int(gold_id)].item()),
                "gold_prob_base": float(prob_base[int(gold_id)].item()),
                "gold_prob_arm": float(prob_arm[int(gold_id)].item()),
                "gold_final_prob_after_support_softmax": gold_final_prob,
                "logp_base_top10_allowed": _top_rows_from_allowed(logp_base, allowed_ids, tokenizer, int(gold_id), config.topk),
                "logp_arm_top10_allowed": _top_rows_from_allowed(logp_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
                "score_s_top10_allowed": score_rows,
                "prob_base_top10_allowed": _top_rows_from_allowed(prob_base, allowed_ids, tokenizer, int(gold_id), config.topk),
                "prob_arm_top10_allowed": _top_rows_from_allowed(prob_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
                "prob_score_top10_allowed": prob_rows,
            }
        )

    return {
        "scheme": "B",
        "support_rule": "sparsemax(logp_base) > 0",
        "prompt": prompt,
        "response": response,
        "response_token_count": len(response_ids),
        "steps": steps,
    }


def scheme_a_supervised_old_helpful_x1(config: SchemeAConfig, x1_config: X1Config | None = None) -> dict:
    tokenizer, base_model, arm_model = load_base_and_helpful_models(config.adapter_path)
    prompt, response = _resolve_prompt_and_response(config.prompt, config.response)

    prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    response_ids = tokenizer(response, add_special_tokens=False)["input_ids"]

    x1_config = x1_config or X1Config()
    steps = []
    for step_idx, gold_id in enumerate(response_ids[: config.steps]):
        prefix_ids = response_ids[:step_idx]
        full_prefix = prompt_ids + prefix_ids

        logp_base = _forward_last_logprobs(base_model, full_prefix)
        logp_arm = _forward_last_logprobs(arm_model, full_prefix)
        prob_base = torch.exp(logp_base)
        prob_arm = torch.exp(logp_arm)
        allowed_ids = _compute_allowed_ids_scheme_a(prob_base, config.base_prob_threshold)
        score_s, gate = _build_score_with_optional_x1(logp_base, logp_arm, prob_base, x1_config)

        filtered_score = score_s[allowed_ids]
        filtered_probs = torch.softmax(filtered_score, dim=-1)

        if int(gold_id) in set(allowed_ids.tolist()):
            gold_allowed_index = (allowed_ids == int(gold_id)).nonzero(as_tuple=False).item()
            gold_final_prob = float(filtered_probs[gold_allowed_index].item())
            gold_in_scope = True
        else:
            gold_final_prob = 0.0
            gold_in_scope = False

        prob_rows = []
        top_prob_vals, top_prob_idx = torch.topk(filtered_probs, k=min(config.topk, allowed_ids.numel()))
        top_prob_ids = allowed_ids[top_prob_idx]
        seen_gold_prob = False
        for tok_id, val in zip(top_prob_ids.tolist(), top_prob_vals.tolist()):
            is_gold = int(tok_id) == int(gold_id)
            seen_gold_prob = seen_gold_prob or is_gold
            prob_rows.append(
                {
                    "token_id": int(tok_id),
                    "token_text": tokenizer.decode([int(tok_id)]),
                    "value": float(val),
                    "is_gold": bool(is_gold),
                    "gold_in_scope": True,
                }
            )
        if not seen_gold_prob:
            prob_rows.append(
                {
                    "token_id": int(gold_id),
                    "token_text": tokenizer.decode([int(gold_id)]),
                    "value": gold_final_prob,
                    "is_gold": True,
                    "gold_in_scope": False,
                }
            )

        step_payload = {
            "step_index": step_idx + 1,
            "gold_token_id": int(gold_id),
            "gold_token_text": tokenizer.decode([int(gold_id)]),
            "prefix_response_token_ids": prefix_ids,
            "prefix_response_text": tokenizer.decode(prefix_ids),
            "allowed_token_count": int(allowed_ids.numel()),
            "gold_token_base_prob": float(prob_base[int(gold_id)].item()),
            "gold_token_in_scope": gold_in_scope,
            "gold_logp_base": float(logp_base[int(gold_id)].item()),
            "gold_logp_arm": float(logp_arm[int(gold_id)].item()),
            "gold_score_s": float(score_s[int(gold_id)].item()),
            "gold_prob_arm": float(prob_arm[int(gold_id)].item()),
            "gold_final_prob_after_filter_softmax": gold_final_prob,
            "logp_base_top10_allowed": _top_rows_from_allowed(logp_base, allowed_ids, tokenizer, int(gold_id), config.topk),
            "logp_arm_top10_allowed": _top_rows_from_allowed(logp_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
            "score_s_top10_allowed": _score_rows_from_allowed(score_s, allowed_ids, tokenizer, int(gold_id), config.topk),
            "prob_base_top10_allowed": _top_rows_from_allowed(prob_base, allowed_ids, tokenizer, int(gold_id), config.topk),
            "prob_arm_top10_allowed": _top_rows_from_allowed(prob_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
            "prob_score_top10_allowed": prob_rows,
        }
        step_payload.update(gate)
        steps.append(step_payload)

    return {
        "scheme": "A",
        "plugin": "X1",
        "threshold": config.base_prob_threshold,
        "prompt": prompt,
        "response": response,
        "response_token_count": len(response_ids),
        "steps": steps,
    }


def scheme_b_supervised_old_helpful_x1(config: SchemeBConfig, x1_config: X1Config | None = None) -> dict:
    tokenizer, base_model, arm_model = load_base_and_helpful_models(config.adapter_path)
    prompt, response = _resolve_prompt_and_response(config.prompt, config.response)

    prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    response_ids = tokenizer(response, add_special_tokens=False)["input_ids"]

    x1_config = x1_config or X1Config()
    steps = []
    for step_idx, gold_id in enumerate(response_ids[: config.steps]):
        prefix_ids = response_ids[:step_idx]
        full_prefix = prompt_ids + prefix_ids

        logp_base = _forward_last_logprobs(base_model, full_prefix)
        logp_arm = _forward_last_logprobs(arm_model, full_prefix)
        prob_base = torch.exp(logp_base)
        prob_arm = torch.exp(logp_arm)
        allowed_ids, base_sparse_probs = _compute_allowed_ids_scheme_b(logp_base)
        score_s, gate = _build_score_with_optional_x1(logp_base, logp_arm, prob_base, x1_config)

        filtered_score = score_s[allowed_ids]
        filtered_probs = torch.softmax(filtered_score, dim=-1)

        if int(gold_id) in set(allowed_ids.tolist()):
            gold_allowed_index = (allowed_ids == int(gold_id)).nonzero(as_tuple=False).item()
            gold_final_prob = float(filtered_probs[gold_allowed_index].item())
            gold_in_scope = True
        else:
            gold_final_prob = 0.0
            gold_in_scope = False

        prob_rows = []
        top_prob_vals, top_prob_idx = torch.topk(filtered_probs, k=min(config.topk, allowed_ids.numel()))
        top_prob_ids = allowed_ids[top_prob_idx]
        seen_gold_prob = False
        for tok_id, val in zip(top_prob_ids.tolist(), top_prob_vals.tolist()):
            is_gold = int(tok_id) == int(gold_id)
            seen_gold_prob = seen_gold_prob or is_gold
            prob_rows.append(
                {
                    "token_id": int(tok_id),
                    "token_text": tokenizer.decode([int(tok_id)]),
                    "value": float(val),
                    "is_gold": bool(is_gold),
                    "gold_in_scope": True,
                }
            )
        if not seen_gold_prob:
            prob_rows.append(
                {
                    "token_id": int(gold_id),
                    "token_text": tokenizer.decode([int(gold_id)]),
                    "value": gold_final_prob,
                    "is_gold": True,
                    "gold_in_scope": False,
                }
            )

        step_payload = {
            "step_index": step_idx + 1,
            "gold_token_id": int(gold_id),
            "gold_token_text": tokenizer.decode([int(gold_id)]),
            "prefix_response_token_ids": prefix_ids,
            "prefix_response_text": tokenizer.decode(prefix_ids),
            "allowed_token_count": int(allowed_ids.numel()),
            "gold_token_sparsemax_base_weight": float(base_sparse_probs[int(gold_id)].item()),
            "gold_token_in_scope": gold_in_scope,
            "gold_logp_base": float(logp_base[int(gold_id)].item()),
            "gold_logp_arm": float(logp_arm[int(gold_id)].item()),
            "gold_score_s": float(score_s[int(gold_id)].item()),
            "gold_prob_base": float(prob_base[int(gold_id)].item()),
            "gold_prob_arm": float(prob_arm[int(gold_id)].item()),
            "gold_final_prob_after_support_softmax": gold_final_prob,
            "logp_base_top10_allowed": _top_rows_from_allowed(logp_base, allowed_ids, tokenizer, int(gold_id), config.topk),
            "logp_arm_top10_allowed": _top_rows_from_allowed(logp_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
            "score_s_top10_allowed": _score_rows_from_allowed(score_s, allowed_ids, tokenizer, int(gold_id), config.topk),
            "prob_base_top10_allowed": _top_rows_from_allowed(prob_base, allowed_ids, tokenizer, int(gold_id), config.topk),
            "prob_arm_top10_allowed": _top_rows_from_allowed(prob_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
            "prob_score_top10_allowed": prob_rows,
        }
        step_payload.update(gate)
        steps.append(step_payload)

    return {
        "scheme": "B",
        "plugin": "X1",
        "support_rule": "sparsemax(logp_base) > 0",
        "prompt": prompt,
        "response": response,
        "response_token_count": len(response_ids),
        "steps": steps,
    }


def prefix_only_old_helpful_winner(config: PrefixOnlyConfig) -> dict:
    tokenizer, base_model, arm_model = load_base_and_helpful_models(config.adapter_path)
    prompt, response = _resolve_prompt_and_response(config.prompt, config.response)

    prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    response_ids = tokenizer(response, add_special_tokens=False)["input_ids"]

    steps = []
    for step_idx, gold_id in enumerate(response_ids[: config.steps]):
        prefix_ids = response_ids[:step_idx]
        full_prefix = prompt_ids + prefix_ids

        logp_base = _forward_last_logprobs(base_model, full_prefix)
        logp_arm = _forward_last_logprobs(arm_model, full_prefix)
        score_s = logp_base + logp_arm
        prob_base = torch.exp(logp_base)
        prob_arm = torch.exp(logp_arm)
        prob_score = torch.softmax(score_s, dim=-1)

        steps.append(
            {
                "step_index": step_idx + 1,
                "gold_token_id": int(gold_id),
                "gold_token_text": tokenizer.decode([int(gold_id)]),
                "prefix_response_token_ids": prefix_ids,
                "prefix_response_text": tokenizer.decode(prefix_ids),
                "gold_logp_base": float(logp_base[int(gold_id)].item()),
                "gold_logp_arm": float(logp_arm[int(gold_id)].item()),
                "gold_score_s": float(score_s[int(gold_id)].item()),
                "gold_prob_base": float(prob_base[int(gold_id)].item()),
                "gold_prob_arm": float(prob_arm[int(gold_id)].item()),
                "gold_prob_score": float(prob_score[int(gold_id)].item()),
                "logp_base_top10": _top_rows_full_vocab(logp_base, tokenizer, int(gold_id), config.topk),
                "logp_arm_top10": _top_rows_full_vocab(logp_arm, tokenizer, int(gold_id), config.topk),
                "score_s_top10": _top_rows_full_vocab(score_s, tokenizer, int(gold_id), config.topk),
                "prob_base_top10": _top_rows_full_vocab(prob_base, tokenizer, int(gold_id), config.topk),
                "prob_arm_top10": _top_rows_full_vocab(prob_arm, tokenizer, int(gold_id), config.topk),
                "prob_score_top10": _top_rows_full_vocab(prob_score, tokenizer, int(gold_id), config.topk),
            }
        )

    return {
        "mode": "prefix_only_old_helpful_winner",
        "prompt": prompt,
        "response": response,
        "response_token_count": len(response_ids),
        "steps": steps,
    }


def scheme_b_temp_support_supervised_old_helpful(config: SchemeBTempSupportConfig) -> dict:
    tokenizer, base_model, arm_model = load_base_and_helpful_models(config.adapter_path)
    prompt, response = _resolve_prompt_and_response(config.prompt, config.response)

    prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
    response_ids = tokenizer(response, add_special_tokens=False)["input_ids"]

    steps = []
    for step_idx, gold_id in enumerate(response_ids[: config.steps]):
        prefix_ids = response_ids[:step_idx]
        full_prefix = prompt_ids + prefix_ids

        logp_base = _forward_last_logprobs(base_model, full_prefix)
        logp_arm = _forward_last_logprobs(arm_model, full_prefix)
        prob_base = torch.exp(logp_base)
        prob_arm = torch.exp(logp_arm)

        allowed_ids, scaled_base_sparse_probs = _compute_allowed_ids_scheme_b_temp(
            logp_base, config.support_temperature
        )
        score_s = logp_base + config.alpha * logp_arm
        filtered_score = score_s[allowed_ids]
        filtered_probs = torch.softmax(filtered_score, dim=-1)

        if int(gold_id) in set(allowed_ids.tolist()):
            gold_allowed_index = (allowed_ids == int(gold_id)).nonzero(as_tuple=False).item()
            gold_final_prob = float(filtered_probs[gold_allowed_index].item())
            gold_in_scope = True
        else:
            gold_final_prob = 0.0
            gold_in_scope = False

        prob_rows = []
        top_prob_vals, top_prob_idx = torch.topk(filtered_probs, k=min(config.topk, allowed_ids.numel()))
        top_prob_ids = allowed_ids[top_prob_idx]
        seen_gold_prob = False
        for tok_id, val in zip(top_prob_ids.tolist(), top_prob_vals.tolist()):
            is_gold = int(tok_id) == int(gold_id)
            seen_gold_prob = seen_gold_prob or is_gold
            prob_rows.append(
                {
                    "token_id": int(tok_id),
                    "token_text": tokenizer.decode([int(tok_id)]),
                    "value": float(val),
                    "is_gold": bool(is_gold),
                    "gold_in_scope": True,
                }
            )
        if not seen_gold_prob:
            prob_rows.append(
                {
                    "token_id": int(gold_id),
                    "token_text": tokenizer.decode([int(gold_id)]),
                    "value": gold_final_prob,
                    "is_gold": True,
                    "gold_in_scope": False,
                }
            )

        steps.append(
            {
                "step_index": step_idx + 1,
                "gold_token_id": int(gold_id),
                "gold_token_text": tokenizer.decode([int(gold_id)]),
                "prefix_response_token_ids": prefix_ids,
                "prefix_response_text": tokenizer.decode(prefix_ids),
                "allowed_token_count": int(allowed_ids.numel()),
                "support_temperature": float(config.support_temperature),
                "alpha": float(config.alpha),
                "gold_token_sparsemax_base_weight_scaled": float(scaled_base_sparse_probs[int(gold_id)].item()),
                "gold_token_in_scope": gold_in_scope,
                "gold_logp_base": float(logp_base[int(gold_id)].item()),
                "gold_logp_arm": float(logp_arm[int(gold_id)].item()),
                "gold_score_s": float(score_s[int(gold_id)].item()),
                "gold_prob_base": float(prob_base[int(gold_id)].item()),
                "gold_prob_arm": float(prob_arm[int(gold_id)].item()),
                "gold_final_prob_after_support_softmax": gold_final_prob,
                "logp_base_top10_allowed": _top_rows_from_allowed(logp_base, allowed_ids, tokenizer, int(gold_id), config.topk),
                "logp_arm_top10_allowed": _top_rows_from_allowed(logp_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
                "score_s_top10_allowed": _score_rows_from_allowed(score_s, allowed_ids, tokenizer, int(gold_id), config.topk),
                "prob_base_top10_allowed": _top_rows_from_allowed(prob_base, allowed_ids, tokenizer, int(gold_id), config.topk),
                "prob_arm_top10_allowed": _top_rows_from_allowed(prob_arm, allowed_ids, tokenizer, int(gold_id), config.topk),
                "prob_score_top10_allowed": prob_rows,
            }
        )

    support_counts = [step["allowed_token_count"] for step in steps]
    support_counts_sorted = sorted(support_counts)
    mid = len(support_counts_sorted) // 2
    if len(support_counts_sorted) % 2 == 1:
        median_support = float(support_counts_sorted[mid])
    else:
        median_support = 0.5 * (support_counts_sorted[mid - 1] + support_counts_sorted[mid])

    return {
        "scheme": "B_temp_support",
        "support_rule": "support(sparsemax(logp_base / T_support))",
        "support_temperature": float(config.support_temperature),
        "alpha": float(config.alpha),
        "prompt": prompt,
        "response": response,
        "response_token_count": len(response_ids),
        "support_stats": {
            "steps_analyzed": len(steps),
            "min": int(min(support_counts)),
            "max": int(max(support_counts)),
            "mean": float(sum(support_counts) / len(support_counts)),
            "median": float(median_support),
        },
        "steps": steps,
    }