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, }