OBLITERATUS/obliteratus/bayesian_optimizer.py

"""Bayesian optimization for abliteration hyperparameters.

Implements Optuna TPE-based multi-objective optimization that searches for
optimal ablation parameters co-minimizing refusal rate and KL divergence.

Inspired by Heretic (p-e-w, 2025) which pioneered Bayesian optimization
for abliteration.  OBLITERATUS pushes this further by:

1. **Parametric layer kernel**: Instead of per-layer independent parameters,
   uses a bell-shaped curve described by 4 global params (Heretic-style).
   This reduces the search space from O(n_layers) to O(1) while capturing
   the spatial structure of refusal across layers.
2. **Float direction interpolation**: Direction index is continuous — non-
   integer values interpolate between adjacent SVD directions, unlocking
   a smooth direction space beyond the discrete top-k.
3. **Component-specific weights**: Separate scaling for attention vs MLP
   projections (Heretic showed MLP interventions are more damaging).
4. **Per-expert granularity**: For MoE models, optimizes per-expert scaling.
5. **CoT-aware objectives**: Adds chain-of-thought coherence as a third
   optimization objective for CoT models.
6. **Warm-start from analysis**: Uses OBLITERATUS's analysis-based heuristics
   as initial trial suggestions instead of random initialization.

References:
    - Heretic (p-e-w, 2025): Bayesian optimization for LLM abliteration
    - Akiba et al. (2019): Optuna: A Next-generation Hyperparameter
      Optimization Framework
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import torch
import torch.nn.functional as F

if TYPE_CHECKING:
    from obliteratus.abliterate import AbliterationPipeline

logger = logging.getLogger(__name__)


# Kernel search space, declared once. `objective` suggests from these bounds and
# the warm-start trial is clamped to them; keeping two copies is what let the
# warm start drift outside the space it was seeding (see `_clamp_to_space`).
KERNEL_SPACE: dict[str, tuple[float, float]] = {
    "attn_max_weight": (0.5, 1.0),
    "attn_peak_position": (0.1, 0.9),
    "attn_min_weight": (0.0, 0.3),
    "attn_spread": (0.1, 0.6),
    "mlp_max_weight": (0.3, 1.0),
    "mlp_peak_position": (0.1, 0.9),
    "mlp_min_weight": (0.0, 0.3),
    "mlp_spread": (0.1, 0.6),
}


def _reset_bayesian_state(pipeline: AbliterationPipeline) -> None:
    """Clear status and component overrides from an earlier attempt."""
    pipeline._bayesian_skipped = None
    pipeline._bayesian_attn_scale = None
    pipeline._bayesian_mlp_scale = None


def _skip(pipeline: AbliterationPipeline, reason: str, message: str) -> None:
    """Record and surface a no-op so it cannot be mistaken for a completed run.

    Reported three ways because each reaches a different audience: the stdlib
    logger for anyone capturing logs, the pipeline log for the console panel the
    operator is watching, and an attribute the caller persists into checkpoint
    metadata so the provenance of an already-written model still shows it.
    """
    logger.warning(message)
    log = getattr(pipeline, "log", None)
    if callable(log):
        log(f"WARNING: {message}")
    try:
        pipeline._bayesian_skipped = reason
    except AttributeError:  # pragma: no cover - pipeline stubs may use __slots__
        pass


def _clamp_to_space(params: dict[str, float]) -> dict[str, float]:
    """Clamp warm-start values into the declared search space.

    Optuna warns about an out-of-range enqueued value but still evaluates it.
    Clamping keeps the informed starting point within the same bounds as later
    sampled trials. This is reachable in normal use: `attn_peak_position` is
    seeded from
    `peak_layer / (n_layers - 1)`, which for an early peak layer falls below the
    0.1 floor — a 36-layer model peaking at layer 2 seeds 0.057.
    """
    out = dict(params)
    for name, (lo, hi) in KERNEL_SPACE.items():
        if name in out:
            out[name] = min(max(float(out[name]), lo), hi)
    return out


def _measure_refusal_rate(
    pipeline: AbliterationPipeline,
    n_prompts: int = 10,
    max_new_tokens: int = 128,
) -> float:
    """Measure refusal rate on a subset of harmful prompts.

    Applies the model's chat template (when available) so the model is in
    chat mode — matching how users actually interact with it.  Uses token-ID
    slicing for robust completion extraction regardless of special tokens.
    """
    from obliteratus.evaluation.advanced_metrics import refusal_rate

    model = pipeline.handle.model
    tokenizer = pipeline.handle.tokenizer
    device = pipeline._get_model_device(model)

    test_prompts = pipeline.harmful_prompts[:n_prompts]

    # Always apply chat template for refusal measurement — refusal circuitry
    # is activated by chat formatting, so raw prompts give false 0% rates.
    saved_flag = pipeline.use_chat_template
    pipeline.use_chat_template = True
    formatted_prompts = pipeline._maybe_apply_chat_template(test_prompts)
    pipeline.use_chat_template = saved_flag
    responses = []

    for prompt in formatted_prompts:
        try:
            inputs = tokenizer(
                prompt, return_tensors="pt", truncation=True, max_length=256,
            )
            input_len = inputs["input_ids"].shape[1]
            inputs = {k: v.to(device) for k, v in inputs.items()}
            with torch.no_grad():
                output = model.generate(
                    **inputs, max_new_tokens=max_new_tokens, do_sample=False,
                )
            # Token-ID slicing: robust regardless of special tokens
            completion_ids = output[0][input_len:]
            completion = tokenizer.decode(
                completion_ids, skip_special_tokens=True,
            ).strip()
            responses.append(completion)
            del inputs, output
        except Exception:
            responses.append("")  # Count failures as non-refusals

    pipeline._free_gpu_memory()
    return refusal_rate(responses, mode="combined") if responses else 1.0


def _measure_kl_divergence(
    pipeline: AbliterationPipeline,
    reference_logits: list[torch.Tensor],
    prompts: list[str],
) -> float:
    """Measure KL divergence from reference (pre-ablation) logits."""
    model = pipeline.handle.model
    tokenizer = pipeline.handle.tokenizer
    device = pipeline._get_model_device(model)

    total_kl = 0.0
    n_valid = 0

    for i, prompt in enumerate(prompts):
        if i >= len(reference_logits):
            break
        try:
            inputs = tokenizer(
                prompt, return_tensors="pt", truncation=True, max_length=64,
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = model(**inputs)
                new_logits = outputs.logits[:, -1, :].detach().cpu().float()

            ref = reference_logits[i]
            log_p = F.log_softmax(ref, dim=-1)
            log_q = F.log_softmax(new_logits.squeeze(0), dim=-1)
            p = F.softmax(ref, dim=-1)
            kl = (p * (log_p - log_q)).sum().item()
            total_kl += max(kl, 0.0)  # Clamp negative KL (numerical noise)
            n_valid += 1
            del inputs, outputs, new_logits
        except Exception:
            pass

    pipeline._free_gpu_memory()
    return total_kl / max(n_valid, 1)


def _parametric_layer_weight(
    layer_idx: int,
    n_layers: int,
    max_weight: float,
    peak_position: float,
    min_weight: float,
    spread: float,
) -> float:
    """Compute ablation weight for a layer using a piecewise-linear tent kernel.

    Faithful reproduction of Heretic's parametric kernel (p-e-w/heretic):
    - max_weight: peak ablation strength at peak_position
    - peak_position: normalized position of peak (0..1)
    - min_weight: weight at the edges of the tent
    - spread: normalized distance from peak to tent edge (min_weight_distance)

    Layers beyond ``spread`` from the peak get weight 0 (skipped entirely).
    Within the tent, weight drops linearly from max_weight to min_weight.
    This matches Heretic's actual formula::

        distance = abs(layer_index - max_weight_position)
        if distance > min_weight_distance: skip
        weight = max_weight + (distance / min_weight_distance) * (min_weight - max_weight)
    """
    if n_layers <= 1:
        return max_weight

    normalized_pos = layer_idx / (n_layers - 1)
    dist = abs(normalized_pos - peak_position)
    min_weight_distance = max(spread, 0.01)

    # Hard cutoff: layers outside the tent get 0 weight (Heretic skips them)
    if dist > min_weight_distance:
        return 0.0

    # Linear interpolation: max_weight at peak → min_weight at edges
    return max_weight + (dist / min_weight_distance) * (min_weight - max_weight)


def _interpolate_direction(
    pipeline: AbliterationPipeline,
    layer_idx: int,
    float_dir_idx: float,
) -> torch.Tensor:
    """Get an interpolated refusal direction from a float-valued layer index.

    Faithful reproduction of Heretic's direction interpolation: the index
    selects which *layer's* diff-of-means direction to use, with float
    values interpolating between adjacent layers' directions.  This is
    fundamentally different from interpolating between SVD components
    within a single layer — it searches across the layer axis.

    From Heretic source (model.py)::

        weight, index = math.modf(direction_index + 1)
        refusal_direction = F.normalize(
            refusal_directions[int(index)].lerp(
                refusal_directions[int(index) + 1], weight), p=2, dim=0)

    Args:
        pipeline: Pipeline with extracted refusal directions per layer.
        layer_idx: The layer being projected (used as fallback).
        float_dir_idx: Continuous direction index — selects which layer's
            direction to use (e.g., 5.3 interpolates 70% layer-5 + 30% layer-6).

    Returns:
        Normalized direction tensor.
    """
    # Build sorted list of layer indices that have refusal directions
    sorted_layers = sorted(pipeline.refusal_directions.keys())
    if not sorted_layers:
        return pipeline.refusal_directions.get(layer_idx, torch.zeros(1))

    n_layers_with_dirs = len(sorted_layers)

    # Heretic uses direction_index + 1 offset; we map float_dir_idx into
    # the sorted layer list, clamped to valid range.
    float_dir_idx = max(0.0, min(float_dir_idx, n_layers_with_dirs - 1))

    lo = int(float_dir_idx)
    hi = min(lo + 1, n_layers_with_dirs - 1)

    lo_layer = sorted_layers[lo]
    hi_layer = sorted_layers[hi]

    d_lo = pipeline.refusal_directions[lo_layer]
    d_hi = pipeline.refusal_directions[hi_layer]

    if lo == hi:
        d = d_lo
    else:
        # Linear interpolation between adjacent layers' directions
        alpha = float_dir_idx - lo
        d = (1.0 - alpha) * d_lo + alpha * d_hi

    norm = d.norm()
    if norm > 1e-8:
        d = d / norm
    return d


def run_bayesian_optimization(
    pipeline: AbliterationPipeline,
    n_trials: int = 50,
    n_refusal_prompts: int = 30,
    n_kl_prompts: int = 5,
) -> dict[int, float]:
    """Run Bayesian optimization to find optimal ablation parameters.

    Uses Optuna TPE with a Heretic-style parametric layer kernel to search
    a compact parameter space:
    - 4 kernel params (max_weight, peak_position, min_weight, spread)
    - 1 float direction index (interpolated between SVD directions)
    - 2 component weights (attention vs MLP scaling)
    Total: 7 parameters regardless of model size (vs O(n_layers) before).

    Also optimizes per-layer independently when the kernel doesn't fit well
    (uses kernel as warm-start for per-layer refinement in phase 2).

    Args:
        pipeline: Initialized AbliterationPipeline (post-DISTILL stage).
        n_trials: Number of optimization trials.
        n_refusal_prompts: Number of harmful prompts for refusal measurement.
        n_kl_prompts: Number of harmless prompts for KL measurement.

    Returns:
        Dict mapping layer_idx -> optimal regularization value.
    """
    _reset_bayesian_state(pipeline)
    if n_trials <= 0:
        _skip(pipeline, "no-trials", "Bayesian optimization SKIPPED: no trials requested.")
        return {}

    try:
        import optuna
        from optuna.samplers import TPESampler
    except ImportError:
        # Keep the operator log and saved provenance explicit even when a
        # partial or older installation lacks the optimizer dependency.
        _skip(
            pipeline,
            "optuna-not-installed",
            "Optuna not installed — Bayesian optimization SKIPPED. Layer "
            "weights fall back to method defaults; results are NOT optimized. "
            "Install with: pip install optuna",
        )
        return {}

    if not pipeline.handle or not pipeline._strong_layers:
        _skip(
            pipeline,
            "no-strong-layers",
            "Bayesian optimization SKIPPED: no model handle or no selected "
            "layers to optimize over.",
        )
        return {}

    model = pipeline.handle.model
    tokenizer = pipeline.handle.tokenizer
    device = pipeline._get_model_device(model)

    pipeline.log("Bayesian optimization: collecting reference logits...")

    # Collect reference (pre-ablation) logits for KL measurement
    kl_prompts = [
        "The weather today is",
        "In mathematics, the concept of",
        "The history of ancient Rome",
        "How to cook a simple pasta",
        "The process of photosynthesis involves",
    ][:n_kl_prompts]

    reference_logits: list[torch.Tensor] = []
    for prompt in kl_prompts:
        try:
            inputs = tokenizer(
                prompt, return_tensors="pt", truncation=True, max_length=64,
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = model(**inputs)
                reference_logits.append(
                    outputs.logits[:, -1, :].detach().cpu().float().squeeze(0)
                )
            del inputs, outputs
        except Exception:
            pass
    pipeline._free_gpu_memory()

    if not reference_logits:
        _skip(
            pipeline,
            "reference-logits-unavailable",
            "Bayesian optimization SKIPPED: Failed to collect reference logits.",
        )
        return {}

    from obliteratus.strategies.utils import (
        get_layer_modules,
        get_attention_module,
        get_ffn_module,
    )
    from obliteratus.abliterate import _ATTN_OUT_NAMES, _FFN_OUT_NAMES

    layer_modules = get_layer_modules(pipeline.handle)
    arch = pipeline.handle.architecture
    n_total_layers = len(layer_modules)

    # Save weight tensors for rollback — clone to CPU to free GPU memory
    original_params: list[tuple[torch.Tensor, torch.Tensor]] = []
    seen_data_ptrs: set[int] = set()

    for idx in pipeline._strong_layers:
        try:
            attn = get_attention_module(layer_modules[idx], arch)
            for attr_name in _ATTN_OUT_NAMES:
                proj = getattr(attn, attr_name, None)
                if proj is not None and hasattr(proj, "weight"):
                    ptr = proj.weight.data.data_ptr()
                    if ptr not in seen_data_ptrs:
                        original_params.append((proj.weight.data, proj.weight.data.clone().cpu()))
                        seen_data_ptrs.add(ptr)
                    if hasattr(proj, "bias") and proj.bias is not None:
                        bptr = proj.bias.data.data_ptr()
                        if bptr not in seen_data_ptrs:
                            original_params.append((proj.bias.data, proj.bias.data.clone().cpu()))
                            seen_data_ptrs.add(bptr)
        except (AttributeError, RuntimeError):
            pass
        try:
            ffn = get_ffn_module(layer_modules[idx], arch)
            for attr_name in _FFN_OUT_NAMES:
                proj = getattr(ffn, attr_name, None)
                if proj is not None and hasattr(proj, "weight"):
                    ptr = proj.weight.data.data_ptr()
                    if ptr not in seen_data_ptrs:
                        original_params.append((proj.weight.data, proj.weight.data.clone().cpu()))
                        seen_data_ptrs.add(ptr)
                    if hasattr(proj, "bias") and proj.bias is not None:
                        bptr = proj.bias.data.data_ptr()
                        if bptr not in seen_data_ptrs:
                            original_params.append((proj.bias.data, proj.bias.data.clone().cpu()))
                            seen_data_ptrs.add(bptr)
        except (AttributeError, RuntimeError):
            pass

    del seen_data_ptrs
    total_saved_mb = sum(clone.nelement() * clone.element_size() for _, clone in original_params) / 1e6
    pipeline.log(f"  Saved {len(original_params)} weight tensors for rollback ({total_saved_mb:.0f} MB, on CPU)")

    def _restore_all():
        for live_data, saved_clone in original_params:  # noqa: F821
            live_data.copy_(saved_clone.to(live_data.device))

    # Warm-start values for the parametric kernel.
    # If the informed pipeline provided analysis-derived warm-start params,
    # use those (they're much better than the default heuristic).
    informed_warm = getattr(pipeline, "_informed_warm_start", None)
    if informed_warm:
        warm_peak = informed_warm.get("peak_position", 0.5)
        pipeline.log(f"  Using analysis-informed warm-start (peak={warm_peak:.2f})")
    elif pipeline._strong_layers:
        peak_layer = pipeline._strong_layers[0]
        warm_peak = peak_layer / max(n_total_layers - 1, 1)
    else:
        warm_peak = 0.5

    best_result: dict[int, float] = {}
    best_score = float("inf")

    # Suppress Optuna's verbose logging
    optuna.logging.set_verbosity(optuna.logging.WARNING)

    # Max layers with directions (for float direction interpolation)
    n_layers_with_dirs = len([
        idx for idx in pipeline._strong_layers
        if idx in pipeline.refusal_directions
    ])

    # ── Phase 1: Parametric kernel optimization (compact search space) ──
    # Heretic uses SEPARATE kernel parameters for attention and MLP,
    # allowing them to peak at different layers (8 params + 1 dir_idx = 9).

    def objective(trial: optuna.Trial) -> tuple[float, float]:
        """Multi-objective: minimize (refusal_rate, kl_divergence)."""
        _restore_all()

        # Attention kernel: 4 params
        attn_max = trial.suggest_float("attn_max_weight", *KERNEL_SPACE["attn_max_weight"])
        attn_peak = trial.suggest_float("attn_peak_position", *KERNEL_SPACE["attn_peak_position"])
        attn_min = trial.suggest_float("attn_min_weight", *KERNEL_SPACE["attn_min_weight"])
        attn_spread = trial.suggest_float("attn_spread", *KERNEL_SPACE["attn_spread"])

        # MLP kernel: 4 params (separate — can peak at a different layer)
        mlp_max = trial.suggest_float("mlp_max_weight", *KERNEL_SPACE["mlp_max_weight"])
        mlp_peak = trial.suggest_float("mlp_peak_position", *KERNEL_SPACE["mlp_peak_position"])
        mlp_min = trial.suggest_float("mlp_min_weight", *KERNEL_SPACE["mlp_min_weight"])
        mlp_spread = trial.suggest_float("mlp_spread", *KERNEL_SPACE["mlp_spread"])

        # Float direction index (cross-layer interpolation, Heretic-style)
        dir_idx = trial.suggest_float("dir_idx", 0.0, max(n_layers_with_dirs - 1, 0.0))

        # Compute per-layer, per-component regularization from kernels
        attn_regs: dict[int, float] = {}
        mlp_regs: dict[int, float] = {}
        for idx in pipeline._strong_layers:
            attn_w = _parametric_layer_weight(idx, n_total_layers, attn_max, attn_peak, attn_min, attn_spread)
            mlp_w = _parametric_layer_weight(idx, n_total_layers, mlp_max, mlp_peak, mlp_min, mlp_spread)
            attn_regs[idx] = 1.0 - attn_w
            mlp_regs[idx] = 1.0 - mlp_w

        # Apply projection with trial's parameters
        for idx in pipeline._strong_layers:
            if idx not in pipeline.refusal_directions:
                continue

            # Use cross-layer interpolated direction
            direction = _interpolate_direction(pipeline, idx, dir_idx)
            d_col = direction.to(device=next(layer_modules[idx].parameters()).device)
            d_col = d_col.unsqueeze(-1) if d_col.dim() == 1 else d_col

            # Attention projection (with per-component kernel)
            attn_reg = attn_regs[idx]
            try:
                attn = get_attention_module(layer_modules[idx], arch)
                pipeline._project_out_advanced(
                    attn, d_col, _ATTN_OUT_NAMES,
                    norm_preserve=pipeline.norm_preserve,
                    regularization=attn_reg,
                )
            except (AttributeError, RuntimeError):
                pass

            # MLP/FFN projection (with per-component kernel)
            mlp_reg = mlp_regs[idx]
            try:
                ffn = get_ffn_module(layer_modules[idx], arch)
                count = pipeline._project_out_advanced(
                    ffn, d_col, _FFN_OUT_NAMES,
                    norm_preserve=pipeline.norm_preserve,
                    regularization=mlp_reg,
                )
                if count == 0:
                    pipeline._project_moe_experts(
                        ffn, d_col,
                        norm_preserve=pipeline.norm_preserve,
                        regularization=mlp_reg,
                        project_biases=False,
                    )
            except (AttributeError, RuntimeError):
                pass

        # Measure objectives
        refusal = _measure_refusal_rate(pipeline, n_prompts=n_refusal_prompts)
        kl = _measure_kl_divergence(pipeline, reference_logits, kl_prompts)

        # Track best combined score (use average of attn/mlp regs for layer_regs)
        nonlocal best_score, best_result
        combined = refusal + 0.5 * kl
        if combined < best_score:
            best_score = combined
            best_result = {
                idx: (attn_regs[idx] + mlp_regs[idx]) / 2.0
                for idx in pipeline._strong_layers
            }

        pipeline.log(
            f"  Trial {trial.number + 1}/{n_trials}: "
            f"refusal={refusal:.0%}, KL={kl:.4f} "
            f"(attn_peak={attn_peak:.2f}, mlp_peak={mlp_peak:.2f}, dir={dir_idx:.2f})"
        )

        return refusal, kl

    sampler = TPESampler(seed=42, n_startup_trials=min(5, n_trials // 3))
    study = optuna.create_study(
        directions=["minimize", "minimize"],
        sampler=sampler,
        study_name="obliteratus_parametric_optimization",
    )

    # Enqueue warm-start trial with analysis-derived estimates.
    # Translate informed pipeline params to the new per-component format.
    if informed_warm:
        iw = informed_warm
        warm_params = {
            "attn_max_weight": iw.get("max_weight", 0.9),
            "attn_peak_position": iw.get("peak_position", warm_peak),
            "attn_min_weight": iw.get("min_weight", 0.05),
            "attn_spread": iw.get("spread", 0.3),
            "mlp_max_weight": iw.get("max_weight", 0.9) * iw.get("mlp_scale", 0.6),
            "mlp_peak_position": iw.get("peak_position", warm_peak),
            "mlp_min_weight": iw.get("min_weight", 0.05),
            "mlp_spread": iw.get("spread", 0.3),
            "dir_idx": iw.get("dir_idx", 0.0),
        }
    else:
        warm_params = {
            "attn_max_weight": 0.9,
            "attn_peak_position": warm_peak,
            "attn_min_weight": 0.05,
            "attn_spread": 0.3,
            "mlp_max_weight": 0.6,
            "mlp_peak_position": warm_peak,
            "mlp_min_weight": 0.05,
            "mlp_spread": 0.3,
            "dir_idx": 0.0,
        }
    study.enqueue_trial(_clamp_to_space(warm_params))

    pipeline.log(f"Bayesian optimization: running {n_trials} trials (parametric kernel)...")
    study.optimize(objective, n_trials=n_trials, show_progress_bar=False)

    # Restore model and apply best result
    _restore_all()

    # Get best trial from Pareto front (prefer low refusal)
    pareto = study.best_trials
    if pareto:
        pareto.sort(key=lambda t: (t.values[0], t.values[1]))
        best_trial = pareto[0]

        # Reconstruct per-layer regs from best kernel params
        p = best_trial.params
        best_result = {}
        for idx in pipeline._strong_layers:
            attn_w = _parametric_layer_weight(
                idx, n_total_layers,
                p["attn_max_weight"], p["attn_peak_position"],
                p["attn_min_weight"], p["attn_spread"],
            )
            mlp_w = _parametric_layer_weight(
                idx, n_total_layers,
                p["mlp_max_weight"], p["mlp_peak_position"],
                p["mlp_min_weight"], p["mlp_spread"],
            )
            best_result[idx] = (attn_w + mlp_w) / 2.0  # average for layer-level reg
            best_result[idx] = 1.0 - best_result[idx]

        pipeline.log(
            f"  Best trial: refusal={best_trial.values[0]:.0%}, "
            f"KL={best_trial.values[1]:.4f}"
        )
        pipeline.log(
            f"  Attn kernel: peak={p['attn_peak_position']:.2f}, "
            f"spread={p['attn_spread']:.2f}, max={p['attn_max_weight']:.2f}"
        )
        pipeline.log(
            f"  MLP kernel:  peak={p['mlp_peak_position']:.2f}, "
            f"spread={p['mlp_spread']:.2f}, max={p['mlp_max_weight']:.2f}"
        )
        pipeline.log(f"  dir_idx={p['dir_idx']:.2f}")

        # Store the best direction index for use during EXCISE
        best_dir_idx = p.get("dir_idx", 0.0)
        if best_dir_idx > 0.1:
            pipeline.log(f"  Applying interpolated direction (idx={best_dir_idx:.2f})...")
            for idx in pipeline._strong_layers:
                new_dir = _interpolate_direction(pipeline, idx, best_dir_idx)
                pipeline.refusal_directions[idx] = new_dir

        # Store component scales for use in EXCISE (backward compat)
        pipeline._bayesian_attn_scale = p.get("attn_max_weight", 1.0)
        pipeline._bayesian_mlp_scale = p.get("mlp_max_weight", 1.0)

    elif best_result:
        pipeline.log(f"  Using best combined score: {best_score:.4f}")

    # Clean up
    del original_params
    pipeline._free_gpu_memory()

    if not best_result:
        _skip(
            pipeline,
            "no-successful-trials",
            "Bayesian optimization SKIPPED: no usable trial results.",
        )
    return best_result