Upload 129 files

README.md

@@ -6,7 +6,7 @@ colorTo: gray
 sdk: gradio
 sdk_version: "5.29.0"
 app_file: app.py
-persistent_storage: true
+persistent_storage: large
 pinned: true
 license: agpl-3.0
 tags:

app.py

@@ -1577,16 +1577,6 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
         _ts = datetime.now().strftime("%H:%M")
         _short_model = model_id.split("/")[-1] if "/" in model_id else model_id
         _cache_label = f"{method} on {_short_model} ({_ts})"
-        _last_obliterated_label = _cache_label
-        _session_models[_cache_label] = {
-            "model_id": model_id,
-            "model_choice": model_choice,
-            "method": method,
-            "dataset_key": dataset_key if not use_custom else "custom",
-            "prompt_volume": prompt_volume,
-            "output_dir": save_dir,
-            "source": "obliterate",
-        }
 
         # Preserve activation steering metadata for re-installation after reload
         steering_meta = None
@@ -1601,6 +1591,16 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
                 "steering_strength": pipeline.steering_strength,
             }
         with _lock:
+            _last_obliterated_label = _cache_label
+            _session_models[_cache_label] = {
+                "model_id": model_id,
+                "model_choice": model_choice,
+                "method": method,
+                "dataset_key": dataset_key if not use_custom else "custom",
+                "prompt_volume": prompt_volume,
+                "output_dir": save_dir,
+                "source": "obliterate",
+            }
             _state["steering"] = steering_meta
             _state["output_dir"] = save_dir  # for ZeroGPU checkpoint reload
 
@@ -3813,13 +3813,15 @@ To opt out, set the environment variable `OBLITERATUS_TELEMETRY=0` before launch
             def _load_leaderboard():
                 """Load leaderboard data and format as markdown table."""
                 try:
-                    from obliteratus.telemetry import get_leaderboard_data, is_telemetry_enabled
+                    from obliteratus.telemetry import get_leaderboard_data, is_telemetry_enabled, storage_diagnostic
                     if not is_telemetry_enabled():
                         return "Telemetry is disabled. Remove `OBLITERATUS_TELEMETRY=0` or set it to `1` to re-enable.", ""
 
                     data = get_leaderboard_data()
                     if not data:
-                        return "No benchmark results yet. Run a benchmark to populate the leaderboard!", ""
+                        diag = storage_diagnostic()
+                        storage_info = f"Storage: `{diag['telemetry_dir']}` (persistent={diag['is_persistent']})"
+                        return f"No benchmark results yet. Run a benchmark to populate the leaderboard!\n\n{storage_info}", ""
 
                     # Build markdown table
                     lines = [
@@ -3848,16 +3850,20 @@ To opt out, set the environment variable `OBLITERATUS_TELEMETRY=0` before launch
                     unique_models = len(set(r['model_id'] for r in data))
                     unique_methods = len(set(r['method'] for r in data))
 
-                    # Check data source
+                    # Check data source and storage status
                     from obliteratus.telemetry import _TELEMETRY_REPO
                     source_note = ""
                     if _TELEMETRY_REPO:
                         source_note = f" | Data source: local + [{_TELEMETRY_REPO}](https://huggingface.co/datasets/{_TELEMETRY_REPO})"
 
+                    diag = storage_diagnostic()
+                    persistent_badge = "persistent" if diag["is_persistent"] else "**EPHEMERAL**"
+                    storage_note = f" | Storage: `{diag['telemetry_dir']}` ({persistent_badge})"
+
                     summary = (
                         f"**{total_runs}** total runs across "
                         f"**{unique_models}** models and "
-                        f"**{unique_methods}** methods{source_note}"
+                        f"**{unique_methods}** methods{source_note}{storage_note}"
                     )
                     return table, summary
                 except Exception as e:

hf-spaces/README.md

@@ -7,7 +7,7 @@ sdk: gradio
 sdk_version: "5.29.0"
 app_file: app.py
 hardware: zero-a10g
-persistent_storage: true
+persistent_storage: large
 pinned: true
 license: agpl-3.0
 tags:

obliteratus/abliterate.py

@@ -961,15 +961,23 @@ class AbliterationPipeline:
             n_profiled = sum(1 for v in self._routing_harmful.values() if v)
             self.log(f"  Router profiling complete: {n_profiled} MoE layers profiled")
 
+        empty_layers = []
         for idx in range(n_layers):
             if self._harmful_acts[idx] and self._harmless_acts[idx]:
                 self._harmful_means[idx] = torch.stack(self._harmful_acts[idx]).mean(dim=0)
                 self._harmless_means[idx] = torch.stack(self._harmless_acts[idx]).mean(dim=0)
             else:
                 # Layer produced no activations (hook failure or skipped layer)
+                empty_layers.append(idx)
                 hidden = self._harmful_acts[0][0].shape[-1] if self._harmful_acts.get(0) else 768
                 self._harmful_means[idx] = torch.zeros(1, hidden)
                 self._harmless_means[idx] = torch.zeros(1, hidden)
+        if empty_layers:
+            self.log(
+                f"WARNING: {len(empty_layers)} layers produced no activations "
+                f"(layers {empty_layers[:5]}{'...' if len(empty_layers) > 5 else ''}). "
+                f"These will be skipped during direction extraction."
+            )
 
         # ── Jailbreak-contrastive probing ─────────────────────────────────
         if self.use_jailbreak_contrast:
@@ -1421,6 +1429,31 @@ class AbliterationPipeline:
         norms: dict[int, float] = {}
         n_dirs = self.n_directions
 
+        # ── Small-model direction cap ──────────────────────────────────
+        # On small models, each SVD direction removes a proportionally
+        # larger fraction of weight energy.  With norm preservation, this
+        # amplifies noise in the remaining dimensions.  Cap n_directions
+        # to prevent over-ablation that destroys coherence.
+        hidden_size = self.handle.hidden_size if self.handle else 0
+        total_params = getattr(self.handle, 'total_params', 0) if self.handle else 0
+        if total_params == 0 and self.handle:
+            try:
+                total_params = sum(p.numel() for p in self.handle.model.parameters())
+            except Exception:
+                pass
+        if n_dirs > 1 and (
+            (0 < hidden_size < 2048)
+            or (0 < total_params < 2_000_000_000)
+            or n_layers <= 16
+        ):
+            max_dirs = max(1, min(n_dirs, 2))
+            if max_dirs < n_dirs:
+                self.log(
+                    f"Capped n_directions from {n_dirs} to {max_dirs} for small model "
+                    f"(hidden={hidden_size}, params={total_params / 1e9:.1f}B, layers={n_layers})"
+                )
+                n_dirs = max_dirs
+
         # Optionally use Wasserstein-optimal direction extraction
         wasserstein_extractor = None
         if self.use_wasserstein_optimal:
@@ -1653,7 +1686,7 @@ class AbliterationPipeline:
             if (is_small_by_layers or is_small_by_capacity or is_small_by_params) and len(self._strong_layers) > 0:
                 if is_small_by_layers:
                     max_layer_frac = 0.25
-                    reason = f"≤16 layers"
+                    reason = "≤16 layers"
                 else:
                     max_layer_frac = 0.20
                     reasons = []
@@ -2877,13 +2910,39 @@ class AbliterationPipeline:
         if self.spectral_cascade and self._strong_layers:
             self._apply_spectral_cascade_weights()
 
+        # ── Guard: compound norm amplification ────────────────────────
+        # When true_iterative_refinement is disabled, subsequent passes
+        # re-apply the SAME projection directions without re-probing.
+        # With norm_preserve=True and regularization > 0, this creates
+        # pathological amplification: each pass removes residual refusal
+        # energy (reg% of previous), then norm-restoration rescales the
+        # entire weight matrix UP to compensate, amplifying non-refusal
+        # components.  On small models (< 2B params) where refusal is a
+        # significant fraction of total weight energy, this compounds into
+        # inf perplexity and destroyed coherence.
+        #
+        # Fix: cap to 1 pass when not re-probing + norm-preserving + partial
+        # regularization, since extra passes are purely destructive noise
+        # amplification in this configuration.
+        effective_passes = self.refinement_passes
+        if (effective_passes > 1
+                and not self.true_iterative_refinement
+                and self.norm_preserve
+                and self.regularization > 0):
+            self.log(
+                f"Capping refinement_passes from {effective_passes} to 1: "
+                f"norm_preserve + regularization without re-probing causes "
+                f"compound amplification (directions are not re-extracted)"
+            )
+            effective_passes = 1
+
         # Track previous directions for cosine-similarity early-exit
         _prev_directions: dict[int, torch.Tensor] = {}
 
-        for pass_num in range(self.refinement_passes):
+        for pass_num in range(effective_passes):
             modified_this_pass = 0
-            if self.refinement_passes > 1:
-                self.log(f"Refinement pass {pass_num + 1}/{self.refinement_passes}")
+            if effective_passes > 1:
+                self.log(f"Refinement pass {pass_num + 1}/{effective_passes}")
 
             # True iterative refinement: re-probe and re-distill after first pass
             if pass_num > 0 and self.true_iterative_refinement:
@@ -3439,6 +3498,21 @@ class AbliterationPipeline:
         norms: dict[int, float] = {}
         n_dirs = self.n_directions
 
+        # Small-model direction cap (matching main _distill)
+        hidden_size = self.handle.hidden_size if self.handle else 0
+        total_params = getattr(self.handle, 'total_params', 0) if self.handle else 0
+        if total_params == 0 and self.handle:
+            try:
+                total_params = sum(p.numel() for p in self.handle.model.parameters())
+            except Exception:
+                pass
+        if n_dirs > 1 and (
+            (0 < hidden_size < 2048)
+            or (0 < total_params < 2_000_000_000)
+            or n_layers <= 16
+        ):
+            n_dirs = max(1, min(n_dirs, 2))
+
         # Use Wasserstein-optimal extraction when enabled (matching main _distill)
         wasserstein_extractor = None
         if self.use_wasserstein_optimal:

obliteratus/telemetry.py

@@ -72,31 +72,101 @@ _hub_sync_lock = threading.Lock()
 _hub_repo_created: bool = False
 
 # Locate writable telemetry directory
+def _is_mount_point(path: Path) -> bool:
+    """Check if a path is a mount point (different device from parent)."""
+    try:
+        if not path.exists():
+            return False
+        return path.stat().st_dev != path.parent.stat().st_dev
+    except (OSError, ValueError):
+        return False
+
+
+def _test_writable(d: Path) -> bool:
+    """Test if a directory exists and is writable."""
+    try:
+        d.mkdir(parents=True, exist_ok=True)
+        test_file = d / ".write_test"
+        test_file.write_text("ok")
+        test_file.unlink()
+        return True
+    except (PermissionError, OSError):
+        return False
+
+
 def _telemetry_dir() -> Path:
     """Find a writable directory for telemetry storage.
 
-    Prefers HuggingFace Spaces persistent storage (/data) when available,
-    so that telemetry and leaderboard data survive container rebuilds.
+    Priority order:
+    1. ``OBLITERATUS_DATA_DIR`` env var (explicit override)
+    2. HuggingFace Spaces persistent storage (``/data/obliteratus``)
+       — survives container rebuilds and factory resets
+    3. ``~/.obliteratus`` (local installs)
+    4. ``/tmp/obliteratus_telemetry`` (last resort — does NOT survive rebuilds)
+
+    On HF Spaces, ``/data`` is the persistent storage mount point.  If it
+    exists as a real mount but isn't writable yet (race during boot), we
+    retry briefly before falling through.
     """
-    candidates = [
-        # HF Spaces persistent storage — survives container rebuilds
-        Path("/data/obliteratus"),
-        Path.home() / ".obliteratus",
-        Path("/tmp/obliteratus_telemetry"),
-    ]
-    for d in candidates:
-        try:
-            d.mkdir(parents=True, exist_ok=True)
-            # Test writability
-            test_file = d / ".write_test"
-            test_file.write_text("ok")
-            test_file.unlink()
-            return d
-        except (PermissionError, OSError):
-            continue
-    # Last resort
+    # 1. Explicit override — always wins
+    explicit = os.environ.get("OBLITERATUS_DATA_DIR")
+    if explicit:
+        p = Path(explicit)
+        if _test_writable(p):
+            logger.info("Telemetry storage: %s (OBLITERATUS_DATA_DIR)", p)
+            return p
+        logger.warning(
+            "OBLITERATUS_DATA_DIR=%s is not writable, falling through", explicit
+        )
+
+    # 2. HF Spaces persistent storage at /data
+    if _ON_HF_SPACES:
+        data_root = Path("/data")
+        hf_dir = data_root / "obliteratus"
+        # On Spaces, /data may take a moment to mount after container start.
+        # Retry a few times if the directory exists as a mount point but
+        # isn't writable yet.
+        if data_root.exists():
+            for attempt in range(3):
+                if _test_writable(hf_dir):
+                    if attempt > 0:
+                        logger.info(
+                            "Telemetry storage: %s (HF persistent, ready after %d retries)",
+                            hf_dir, attempt,
+                        )
+                    else:
+                        logger.info("Telemetry storage: %s (HF persistent storage)", hf_dir)
+                    return hf_dir
+                # Brief wait for mount to become ready
+                if attempt < 2:
+                    time.sleep(1)
+            # /data exists but isn't writable — warn loudly
+            is_mount = _is_mount_point(data_root)
+            logger.warning(
+                "/data exists (mount_point=%s) but /data/obliteratus is NOT writable. "
+                "Persistent storage may not be enabled for this Space. "
+                "Data will NOT survive factory rebuilds! "
+                "Enable persistent storage in Space settings or set OBLITERATUS_DATA_DIR.",
+                is_mount,
+            )
+
+    # 3. Home directory (local installs)
+    home_dir = Path.home() / ".obliteratus"
+    if _test_writable(home_dir):
+        logger.info("Telemetry storage: %s (home directory)", home_dir)
+        return home_dir
+
+    # 4. Last resort — /tmp does NOT survive rebuilds
     fallback = Path("/tmp/obliteratus_telemetry")
     fallback.mkdir(parents=True, exist_ok=True)
+    if _ON_HF_SPACES:
+        logger.warning(
+            "Telemetry storage: %s — this is EPHEMERAL and will be lost on rebuild! "
+            "Enable persistent storage in your Space settings.",
+            fallback,
+        )
+    else:
+        logger.info("Telemetry storage: %s (temporary)", fallback)
     return fallback
 
 
@@ -107,6 +177,35 @@ TELEMETRY_FILE = _TELEMETRY_DIR / "telemetry.jsonl"
 _write_lock = threading.Lock()
 
 
+def _is_persistent_storage() -> bool:
+    """Check if the current telemetry directory is on persistent storage."""
+    return str(_TELEMETRY_DIR).startswith("/data")
+
+
+def storage_diagnostic() -> dict[str, Any]:
+    """Return a diagnostic dict about the current storage configuration.
+
+    Useful for debugging persistent storage issues on HF Spaces.
+    """
+    data_root = Path("/data")
+    return {
+        "telemetry_dir": str(_TELEMETRY_DIR),
+        "telemetry_file": str(TELEMETRY_FILE),
+        "telemetry_file_exists": TELEMETRY_FILE.exists(),
+        "telemetry_file_size_bytes": (
+            TELEMETRY_FILE.stat().st_size if TELEMETRY_FILE.exists() else 0
+        ),
+        "is_persistent": _is_persistent_storage(),
+        "on_hf_spaces": _ON_HF_SPACES,
+        "data_dir_exists": data_root.exists(),
+        "data_dir_is_mount": _is_mount_point(data_root),
+        "data_dir_writable": os.access(data_root, os.W_OK) if data_root.exists() else False,
+        "explicit_data_dir": os.environ.get("OBLITERATUS_DATA_DIR", ""),
+        "telemetry_repo": _TELEMETRY_REPO,
+        "telemetry_enabled": is_enabled(),
+    }
+
+
 def disable_telemetry():
     """Disable telemetry collection."""
     global _TELEMETRY_ENABLED, _enabled
@@ -371,6 +470,97 @@ def fetch_hub_records(max_records: int = 10000) -> list[dict[str, Any]]:
         return []
 
 
+# ── Hub restore (warm-start after rebuild) ────────────────────────────
+
+_restore_done = False
+_restore_lock = threading.Lock()
+
+
+def restore_from_hub() -> int:
+    """Download community records from Hub into the local JSONL file.
+
+    This is the critical path for surviving factory rebuilds: even if
+    ``/data`` is wiped or unavailable, we can reconstruct the leaderboard
+    from the central Hub dataset on startup.
+
+    Records already present locally (by ``(session_id, timestamp)`` key)
+    are skipped to avoid duplicates.
+
+    Returns the number of new records restored.
+    """
+    global _restore_done
+    if _restore_done:
+        return 0
+    with _restore_lock:
+        if _restore_done:
+            return 0
+        _restore_done = True
+
+    repo = _TELEMETRY_REPO
+    if not repo:
+        return 0
+
+    try:
+        # Read existing local keys for dedup
+        existing_keys: set[tuple[str, str]] = set()
+        if TELEMETRY_FILE.exists():
+            try:
+                with open(TELEMETRY_FILE) as f:
+                    for line in f:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        try:
+                            r = json.loads(line)
+                            existing_keys.add(
+                                (r.get("session_id", ""), r.get("timestamp", ""))
+                            )
+                        except json.JSONDecodeError:
+                            continue
+            except Exception:
+                pass
+
+        hub_records = fetch_hub_records()
+        if not hub_records:
+            return 0
+
+        new_count = 0
+        with _write_lock:
+            with open(TELEMETRY_FILE, "a") as f:
+                for r in hub_records:
+                    key = (r.get("session_id", ""), r.get("timestamp", ""))
+                    if key in existing_keys:
+                        continue
+                    existing_keys.add(key)
+                    f.write(json.dumps(r, default=str) + "\n")
+                    new_count += 1
+
+        if new_count:
+            logger.info(
+                "Restored %d records from Hub repo %s to local storage at %s",
+                new_count, repo, TELEMETRY_FILE,
+            )
+        return new_count
+    except Exception as e:
+        logger.debug("Hub restore failed: %s", e)
+        return 0
+
+
+def _restore_from_hub_bg() -> None:
+    """Background thread: restore Hub records to local on startup."""
+    try:
+        restore_from_hub()
+    except Exception as e:
+        logger.debug("Background Hub restore failed: %s", e)
+
+
+# Auto-restore on HF Spaces startup (background, non-blocking).
+# This ensures the leaderboard has data even after a factory rebuild.
+if _ON_HF_SPACES and is_enabled() and _TELEMETRY_REPO:
+    _restore_thread = threading.Thread(target=_restore_from_hub_bg, daemon=True)
+    _restore_thread.start()
+
+
 # ── Hardware detection ────────────────────────────────────────────────
 
 def _detect_gpu() -> tuple[str, float]:

tests/test_abliterate.py

@@ -1756,7 +1756,12 @@ class TestDistillBasic:
 
 class TestDistillSVD:
     def test_multi_direction_svd(self, handle):
-        """Advanced method: SVD extracts multiple refusal directions."""
+        """Advanced method: SVD extracts multiple refusal directions.
+
+        Note: on small models (hidden_size < 2048 or < 2B params), n_directions
+        is automatically capped to 2 to prevent over-ablation.  The test model
+        (hidden_size=64, 4 layers) triggers this safeguard.
+        """
         from obliteratus.strategies.utils import get_layer_modules
 
         pipeline = AbliterationPipeline(
@@ -1775,10 +1780,10 @@ class TestDistillSVD:
 
         n_layers = len(get_layer_modules(handle))
         assert len(pipeline.refusal_subspaces) == n_layers
+        # Small-model cap: n_directions capped to 2 for tiny test model
+        expected_dirs = min(2, pipeline.n_directions, 5, handle.hidden_size)
         for idx, subspace in pipeline.refusal_subspaces.items():
-            # Should have min(n_directions, n_prompts, hidden_dim) directions
-            n_dirs = min(pipeline.n_directions, 5, handle.hidden_size)
-            assert subspace.shape[0] == n_dirs
+            assert subspace.shape[0] == expected_dirs
             assert subspace.shape[1] == handle.hidden_size
 
         # Primary direction should still be a unit vector

tests/test_telemetry.py

@@ -2,7 +2,9 @@
 
 import json
 import os
+import tempfile
 from dataclasses import dataclass, field
+from pathlib import Path
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -13,13 +15,17 @@ from obliteratus.telemetry import (
     _extract_excise_details,
     _extract_prompt_counts,
     _extract_analysis_insights,
+    _is_mount_point,
+    _test_writable,
     build_report,
     disable_telemetry,
     enable_telemetry,
     is_enabled,
     maybe_send_informed_report,
     maybe_send_pipeline_report,
+    restore_from_hub,
     send_report,
+    storage_diagnostic,
 )
 
 
@@ -597,3 +603,94 @@ class TestStageDurationTracking:
 
         p._emit("summon", "running", "loading...", duration=0)
         assert p._stage_durations == {}
+
+
+# ── Storage helpers ──────────────────────────────────────────────────────
+
+
+class TestStorageHelpers:
+    """Test persistent storage helper functions."""
+
+    def test_test_writable_valid_dir(self):
+        with tempfile.TemporaryDirectory() as d:
+            assert _test_writable(Path(d) / "subdir")
+
+    def test_test_writable_unwritable(self):
+        # /proc is never writable for arbitrary files
+        assert not _test_writable(Path("/proc/obliteratus_test"))
+
+    def test_is_mount_point_existing_path(self):
+        # Should return a bool without raising for any existing path
+        result = _is_mount_point(Path("/"))
+        assert isinstance(result, bool)
+
+    def test_is_mount_point_nonexistent(self):
+        assert not _is_mount_point(Path("/nonexistent_dir_12345"))
+
+    def test_storage_diagnostic_returns_dict(self):
+        diag = storage_diagnostic()
+        assert isinstance(diag, dict)
+        assert "telemetry_dir" in diag
+        assert "is_persistent" in diag
+        assert "on_hf_spaces" in diag
+        assert "telemetry_enabled" in diag
+        assert "data_dir_exists" in diag
+
+
+# ── Hub restore ──────────────────────────────────────────────────────────
+
+
+class TestHubRestore:
+    """Test Hub-to-local restore functionality."""
+
+    def setup_method(self):
+        _reset_telemetry()
+        # Reset restore state so each test can trigger it
+        import obliteratus.telemetry as t
+        t._restore_done = False
+
+    def test_restore_skips_when_no_repo(self):
+        with patch("obliteratus.telemetry._TELEMETRY_REPO", ""):
+            assert restore_from_hub() == 0
+
+    def test_restore_deduplicates(self):
+        """Records already in local JSONL should not be re-added."""
+        import obliteratus.telemetry as t
+
+        with tempfile.TemporaryDirectory() as d:
+            test_file = Path(d) / "telemetry.jsonl"
+            existing = {"session_id": "abc", "timestamp": "2025-01-01T00:00:00"}
+            test_file.write_text(json.dumps(existing) + "\n")
+
+            old_file = t.TELEMETRY_FILE
+            old_repo = t._TELEMETRY_REPO
+            t.TELEMETRY_FILE = test_file
+            t._TELEMETRY_REPO = "test/repo"
+            t._restore_done = False
+
+            try:
+                hub_records = [
+                    {"session_id": "abc", "timestamp": "2025-01-01T00:00:00"},  # duplicate
+                    {"session_id": "def", "timestamp": "2025-01-02T00:00:00"},  # new
+                ]
+                with patch("obliteratus.telemetry.fetch_hub_records", return_value=hub_records):
+                    count = restore_from_hub()
+                    assert count == 1  # Only the new record
+
+                # Verify file contents
+                lines = test_file.read_text().strip().split("\n")
+                assert len(lines) == 2  # original + 1 new
+            finally:
+                t.TELEMETRY_FILE = old_file
+                t._TELEMETRY_REPO = old_repo
+
+    def test_restore_only_runs_once(self):
+        """Calling restore_from_hub() twice should be a no-op the second time."""
+        import obliteratus.telemetry as t
+        t._restore_done = False
+
+        with patch("obliteratus.telemetry._TELEMETRY_REPO", "test/repo"):
+            with patch("obliteratus.telemetry.fetch_hub_records", return_value=[]):
+                restore_from_hub()
+                # Second call should return 0 immediately
+                assert restore_from_hub() == 0