""" Refusal Bench runner — orchestrate a head-to-head comparison of techniques. Pipeline: 1. Split contrastive prompts into extraction set (80%) and eval set (20%). 2. Train ONE harmfulness probe on baseline residuals from the extraction set. The probe is shared across techniques so AUC numbers are directly comparable. 3. Compute baseline refusal-rate + AUC on the eval set (no ablation). 4. For each technique: a. Instantiate; fit on the extraction set. b. Get its ablation hook. c. Generate completions on eval-harmful with the hook active; compute post-ablation refusal rate. d. Extract residuals on eval-harmful + eval-harmless with the hook active; score with the trained probe; compute post-ablation AUC. e. Record Δ refusal-rate and Δ AUC. The two-axis result (Δ refusal-rate vs Δ AUC) is the headline novelty: a technique that drops refusal-rate to ~0 while keeping AUC near baseline has suppressed verbal refusal but left the harmfulness representation intact. This is the Zhao 2507.11878 dissociation, here measured across six techniques on the same model. """ from __future__ import annotations import random import time from dataclasses import asdict, dataclass, field from typing import Callable, List, Optional, Tuple import torch from model import get_model, get_model_name from research import apply_chat_template from .harmfulness_probe import ( evaluate_probe, extract_last_token_residuals, train_probe, ) from .scoring import refusal_rate from .techniques import TECHNIQUES # ----------------------------------------------------------------------------- # Result dataclasses # ----------------------------------------------------------------------------- @dataclass class TechniqueResult: """One row of the bench table.""" name: str paper_url: str layer_used: int refusal_rate_baseline: float refusal_rate_ablated: float delta_refusal_rate: float harmfulness_auc_pre: float harmfulness_auc_post: float delta_auc: float elapsed_seconds: float error: Optional[str] = None @dataclass class BenchResult: """Full bench output. Serializable to JSON via asdict().""" model_name: str layer: int n_extraction_pairs: int n_eval_prompts: int probe_train_auc: float probe_test_auc: float results: List[TechniqueResult] = field(default_factory=list) # ----------------------------------------------------------------------------- # Helpers # ----------------------------------------------------------------------------- def _generate_with_hook( prompt: str, hook_name: Optional[str], hook_fn: Optional[Callable], max_new_tokens: int, temperature: float, ) -> str: """Generate a completion. If hook_name/hook_fn are None, no ablation.""" model = get_model() formatted = apply_chat_template(prompt) tokens = model.to_tokens(formatted) if hook_name is None or hook_fn is None: output = model.generate( tokens, max_new_tokens=max_new_tokens, temperature=temperature, do_sample=True, ) else: with model.hooks(fwd_hooks=[(hook_name, hook_fn)]): output = model.generate( tokens, max_new_tokens=max_new_tokens, temperature=temperature, do_sample=True, ) full_text = model.to_string(output[0]) # Strip the prompt itself so refusal-detection sees only the completion prompt_text = model.to_string(tokens[0]) if full_text.startswith(prompt_text): return full_text[len(prompt_text):] return full_text def _extract_residuals_with_hook( prompts: List[str], extract_layer: int, hook_name: Optional[str], hook_fn: Optional[Callable], ) -> torch.Tensor: """ Last-token residuals at `extract_layer`. If hook_name/hook_fn supplied, they are installed during the forward pass (so the residuals reflect the post-ablation state when extract_layer == ablation layer or downstream of it). """ model = get_model() extract_hook = f"blocks.{extract_layer}.hook_resid_post" residuals: List[torch.Tensor] = [] for prompt in prompts: formatted = apply_chat_template(prompt) if hook_name is None or hook_fn is None: _logits, cache = model.run_with_cache(formatted) else: with model.hooks(fwd_hooks=[(hook_name, hook_fn)]): _logits, cache = model.run_with_cache(formatted) last = cache[extract_hook][:, -1, :].squeeze(0).detach().cpu() residuals.append(last) return torch.stack(residuals, dim=0) def _split( items: List[str], test_fraction: float, rng: random.Random, ) -> Tuple[List[str], List[str]]: """Shuffle then split; minimum 2 in test fold so probe AUC is defined.""" shuffled = items.copy() rng.shuffle(shuffled) n_test = max(2, int(round(len(shuffled) * test_fraction))) return shuffled[n_test:], shuffled[:n_test] # ----------------------------------------------------------------------------- # Main entry point # ----------------------------------------------------------------------------- def run_bench( technique_names: List[str], layer: int, harmful_prompts: List[str], harmless_prompts: List[str], *, test_fraction: float = 0.2, max_new_tokens: int = 32, temperature: float = 0.7, seed: int = 42, ) -> BenchResult: """ Run each named technique on the same data and return scored results. Args: technique_names: keys from refusal_bench.techniques.TECHNIQUES. Unknown names produce an error row instead of crashing the bench. layer: residual-stream layer used for direction extraction AND ablation (techniques like COSMIC may select a different layer via fit() — that's recorded in `layer_used`). harmful_prompts: contrastive prompts the model should refuse. harmless_prompts: contrastive prompts the model should comply with. test_fraction: portion held out for eval (default 0.2). max_new_tokens: completion length budget for refusal-rate measurement. temperature: sampling temperature for generation. seed: split + sampling determinism. Returns: BenchResult — serialize via dataclasses.asdict() for JSON. """ rng = random.Random(seed) extraction_harmful, eval_harmful = _split(harmful_prompts, test_fraction, rng) extraction_harmless, eval_harmless = _split(harmless_prompts, test_fraction, rng) model = get_model() model_name = get_model_name() or "unknown" # ── 1. Train probe on extraction-set baseline residuals ────────────── ext_harmful_resid = extract_last_token_residuals(extraction_harmful, layer) ext_harmless_resid = extract_last_token_residuals(extraction_harmless, layer) probe_info = train_probe(ext_harmful_resid, ext_harmless_resid) probe = probe_info["model"] # ── 2. Baseline (no-ablation) eval-set metrics ─────────────────────── baseline_harmful_resid = extract_last_token_residuals(eval_harmful, layer) baseline_harmless_resid = extract_last_token_residuals(eval_harmless, layer) baseline_residuals = torch.cat([baseline_harmful_resid, baseline_harmless_resid], dim=0) baseline_labels = [1] * len(eval_harmful) + [0] * len(eval_harmless) baseline_eval = evaluate_probe(probe, baseline_residuals, labels=baseline_labels) baseline_auc = baseline_eval["auc"] if baseline_eval["auc"] is not None else 0.5 baseline_completions = [ _generate_with_hook(p, None, None, max_new_tokens, temperature) for p in eval_harmful ] baseline_refusal = refusal_rate(baseline_completions) # ── 3. Per-technique loop ──────────────────────────────────────────── results: List[TechniqueResult] = [] for tname in technique_names: start = time.time() if tname not in TECHNIQUES: results.append(TechniqueResult( name=tname, paper_url="", layer_used=layer, refusal_rate_baseline=baseline_refusal, refusal_rate_ablated=float("nan"), delta_refusal_rate=float("nan"), harmfulness_auc_pre=baseline_auc, harmfulness_auc_post=float("nan"), delta_auc=float("nan"), elapsed_seconds=0.0, error=f"unknown technique: {tname}. Known: {sorted(TECHNIQUES)}", )) continue try: technique = TECHNIQUES[tname]() technique.fit(model, extraction_harmful, extraction_harmless, layer) hook_name, hook_fn = technique.make_ablation_hook() # Refusal rate with ablation hook active ablated_completions = [ _generate_with_hook(p, hook_name, hook_fn, max_new_tokens, temperature) for p in eval_harmful ] ablated_refusal = refusal_rate(ablated_completions) # Post-ablation AUC at the same extract layer abl_harmful_resid = _extract_residuals_with_hook(eval_harmful, layer, hook_name, hook_fn) abl_harmless_resid = _extract_residuals_with_hook(eval_harmless, layer, hook_name, hook_fn) ablated_residuals = torch.cat([abl_harmful_resid, abl_harmless_resid], dim=0) ablated_eval = evaluate_probe(probe, ablated_residuals, labels=baseline_labels) ablated_auc = ablated_eval["auc"] if ablated_eval["auc"] is not None else 0.5 elapsed = time.time() - start results.append(TechniqueResult( name=technique.name, paper_url=technique.paper_url, layer_used=technique._layer if technique._layer is not None else layer, refusal_rate_baseline=baseline_refusal, refusal_rate_ablated=ablated_refusal, delta_refusal_rate=ablated_refusal - baseline_refusal, harmfulness_auc_pre=baseline_auc, harmfulness_auc_post=ablated_auc, delta_auc=ablated_auc - baseline_auc, elapsed_seconds=elapsed, )) except Exception as e: elapsed = time.time() - start results.append(TechniqueResult( name=tname, paper_url="", layer_used=layer, refusal_rate_baseline=baseline_refusal, refusal_rate_ablated=float("nan"), delta_refusal_rate=float("nan"), harmfulness_auc_pre=baseline_auc, harmfulness_auc_post=float("nan"), delta_auc=float("nan"), elapsed_seconds=elapsed, error=f"{type(e).__name__}: {e}", )) return BenchResult( model_name=model_name, layer=layer, n_extraction_pairs=len(extraction_harmful), n_eval_prompts=len(eval_harmful), probe_train_auc=probe_info["train_auc"], probe_test_auc=probe_info["test_auc"], results=results, ) def serialize(result: BenchResult) -> dict: """JSON-friendly dict for HTTP responses.""" return asdict(result)