Upload 133 files

obliteratus/abliterate.py

@@ -2984,27 +2984,39 @@ class AbliterationPipeline:
         # before the standard projection loop.  The found values override the
         # static layer_adaptive_strength weights.
         bayesian_regs: dict[int, float] = {}
-        bayesian_trials = getattr(self, "_bayesian_trials", 0) or (
-            METHODS.get(self.method, {}).get("bayesian_trials", 0)
-        )
+        # Use the explicit user override if set (even 0 = disabled);
+        # only fall back to the method default when unset (None).
+        _user_bayesian = getattr(self, "_bayesian_trials", None)
+        if _user_bayesian is not None:
+            bayesian_trials = int(_user_bayesian)
+        else:
+            bayesian_trials = METHODS.get(self.method, {}).get("bayesian_trials", 0)
         if bayesian_trials > 0 and self._strong_layers and self.handle:
             self.log(f"Running Bayesian optimization ({bayesian_trials} trials)...")
-            from obliteratus.bayesian_optimizer import run_bayesian_optimization
-            bayesian_regs = run_bayesian_optimization(
-                self,
-                n_trials=bayesian_trials,
-                n_refusal_prompts=8,
-                n_kl_prompts=5,
-            )
-            if bayesian_regs:
-                self.log(
-                    f"  Bayesian optimization complete: "
-                    f"optimized {len(bayesian_regs)} layer regularizations"
+            try:
+                from obliteratus.bayesian_optimizer import run_bayesian_optimization
+                bayesian_regs = run_bayesian_optimization(
+                    self,
+                    n_trials=bayesian_trials,
+                    n_refusal_prompts=8,
+                    n_kl_prompts=5,
                 )
-                regs_str = ", ".join(
-                    f"{idx}:{reg:.3f}" for idx, reg in sorted(bayesian_regs.items())
+                if bayesian_regs:
+                    self.log(
+                        f"  Bayesian optimization complete: "
+                        f"optimized {len(bayesian_regs)} layer regularizations"
+                    )
+                    regs_str = ", ".join(
+                        f"{idx}:{reg:.3f}" for idx, reg in sorted(bayesian_regs.items())
+                    )
+                    self.log(f"  Optimal regs: {regs_str}")
+            except Exception as e:
+                self.log(
+                    f"  Bayesian optimization failed (non-fatal): {e}\n"
+                    f"  Continuing with standard projection..."
                 )
-                self.log(f"  Optimal regs: {regs_str}")
+                bayesian_regs = {}
+                self._free_gpu_memory()
 
         # ── LoRA-based reversible ablation ──────────────────────────────
         # When enabled, compute LoRA adapters and merge them instead of

obliteratus/bayesian_optimizer.py

@@ -571,9 +571,22 @@ def run_bayesian_optimization(
         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})...")
+            # Snapshot all original directions BEFORE the loop to avoid
+            # order-dependent contamination (each layer's interpolation
+            # should use the original directions, not already-interpolated ones).
+            orig_directions = {
+                idx: pipeline.refusal_directions[idx].clone()
+                for idx in pipeline._strong_layers
+                if idx in pipeline.refusal_directions
+            }
+            saved_refusal_dirs = pipeline.refusal_directions
+            pipeline.refusal_directions = orig_directions
             for idx in pipeline._strong_layers:
                 new_dir = _interpolate_direction(pipeline, idx, best_dir_idx)
-                pipeline.refusal_directions[idx] = new_dir
+                saved_refusal_dirs[idx] = new_dir
+                # Sync refusal_subspaces so EXCISE uses the optimized direction
+                pipeline.refusal_subspaces[idx] = new_dir.unsqueeze(0)
+            pipeline.refusal_directions = saved_refusal_dirs
 
         # Store component scales for use in EXCISE (backward compat)
         pipeline._bayesian_attn_scale = p.get("attn_max_weight", 1.0)