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