Upload 130 files

app.py

@@ -95,6 +95,9 @@ _state: dict = {
     "log": [],
     # Activation steering metadata (survives model reload)
     "steering": None,  # dict with refusal_directions, strong_layers, steering_strength
+    # Checkpoint directory for ZeroGPU reload (model tensors may become stale
+    # after GPU deallocation — this path lets chat_respond reload from disk)
+    "output_dir": None,
 }
 _lock = threading.Lock()
 
@@ -1328,7 +1331,7 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
             f"   or locally: `export HF_TOKEN=hf_...`\n\n"
             f"Get your token at [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)\n\n"
             f"Alternatively, choose a non-gated model (those without the \U0001f512 icon).",
-            "", gr.update(),
+            "", gr.update(), gr.update(),
         )
         return
 
@@ -1337,14 +1340,14 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
         if not re.match(r'^[a-zA-Z0-9_-]+/[a-zA-Z0-9_.-]+$', push_to_hub):
             yield (
                 "**Error:** Invalid Hub repo format. Use `username/model-name`.",
-                "", gr.update(),
+                "", gr.update(), gr.update(),
             )
             return
         if not os.environ.get("HF_TOKEN"):
             yield (
                 "**Error:** HF_TOKEN not set. Push to Hub requires a write token. "
                 "Set it via `export HF_TOKEN=hf_...` or in your Space secrets.",
-                "", gr.update(),
+                "", gr.update(), gr.update(),
             )
             return
 
@@ -1355,7 +1358,7 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
     _clear_gpu()
     with _lock:
         if _state["status"] == "obliterating":
-            yield "**Error:** An obliteration is already in progress.", "", gr.update()
+            yield "**Error:** An obliteration is already in progress.", "", gr.update(), gr.update()
             return
         _state["log"] = []
         _state["status"] = "obliterating"
@@ -1506,9 +1509,9 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
         status_msg = f"**Obliterating\u2026** ({_elapsed()})"
         if len(log_lines) > last_yielded[0]:
             last_yielded[0] = len(log_lines)
-            yield status_msg, "\n".join(log_lines), gr.update()
+            yield status_msg, "\n".join(log_lines), gr.update(), gr.update()
         else:
-            yield status_msg, "\n".join(log_lines), gr.update()
+            yield status_msg, "\n".join(log_lines), gr.update(), gr.update()
         if time.time() - _pipeline_start > _max_pipeline_secs:
             log_lines.append("\nTIMEOUT: Pipeline exceeded 45-minute limit.")
             break
@@ -1523,7 +1526,7 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
         err_msg = str(error_ref[0]) or repr(error_ref[0])
         log_lines.append(f"\nERROR: {err_msg}")
         _state["log"] = log_lines
-        yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header()
+        yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header(), gr.update()
         return
 
     # Success — keep model in memory for chat.
@@ -1596,6 +1599,7 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
             }
         with _lock:
             _state["steering"] = steering_meta
+            _state["output_dir"] = save_dir  # for ZeroGPU checkpoint reload
 
         if can_generate:
             # Model fits — use it directly (steering hooks already installed)
@@ -1624,7 +1628,7 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
             if bnb_available:
                 log_lines.append("\nModel too large for chat at float16 — reloading in 4-bit...")
                 last_yielded[0] = len(log_lines)
-                yield status_msg, "\n".join(log_lines), gr.update()
+                yield status_msg, "\n".join(log_lines), gr.update(), gr.update()
                 try:
                     from transformers import BitsAndBytesConfig
                     bnb_cfg = BitsAndBytesConfig(
@@ -1671,7 +1675,7 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
                     else "Falling back to CPU offload..."
                 )
                 last_yielded[0] = len(log_lines)
-                yield status_msg, "\n".join(log_lines), gr.update()
+                yield status_msg, "\n".join(log_lines), gr.update(), gr.update()
                 try:
                     offload_dir = tempfile.mkdtemp(prefix="obliteratus_offload_")
                     model_reloaded = AutoModelForCausalLM.from_pretrained(
@@ -1725,7 +1729,13 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
                 f"**{model_choice}** liberated with `{method}` method. "
                 f"Saved to `{save_dir}`. Chat requires a larger GPU."
             )
-        yield status_msg, "\n".join(log_lines), get_chat_header()
+        # Update session dropdown directly (don't rely on .then() which can
+        # fail to fire on ZeroGPU after generator teardown)
+        _dd_update = gr.update(
+            choices=_get_session_model_choices(),
+            value=_last_obliterated_label or None,
+        )
+        yield status_msg, "\n".join(log_lines), get_chat_header(), _dd_update
 
     except Exception as e:
         # Ensure status never gets stuck on "obliterating"
@@ -1734,7 +1744,7 @@ def obliterate(model_choice: str, method_choice: str, hub_repo: str,
         err_msg = str(e) or repr(e)
         log_lines.append(f"\nERROR (post-pipeline): {err_msg}")
         _state["log"] = log_lines
-        yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header()
+        yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header(), gr.update()
 
 
 # ---------------------------------------------------------------------------
@@ -1798,14 +1808,41 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
         return
 
     # ZeroGPU safety: ensure model is on GPU if available.
-    # Between GPU allocations, ZeroGPU may have moved the model to CPU/meta.
-    # The @spaces.GPU decorator guarantees a GPU is available here, so move if needed.
+    # Between GPU allocations, ZeroGPU may have moved the model to CPU/meta,
+    # or tensors may be stale from a previous GPU context.
+    # The @spaces.GPU decorator guarantees a GPU is available here.
+    _needs_reload = False
     try:
         dev = next(model.parameters()).device
         if torch.cuda.is_available() and dev.type != "cuda":
             model.to("cuda")
     except (StopIteration, RuntimeError):
-        pass
+        _needs_reload = True
+
+    # If model tensors are stale/meta, reload from the saved checkpoint
+    if _needs_reload and _ZEROGPU_AVAILABLE:
+        checkpoint = _state.get("output_dir")
+        if checkpoint and Path(checkpoint).exists():
+            try:
+                is_preset = (_state.get("model_name") or "") in MODELS
+                model = AutoModelForCausalLM.from_pretrained(
+                    checkpoint, device_map="auto", torch_dtype=torch.float16,
+                    trust_remote_code=is_preset,
+                )
+                tokenizer = AutoTokenizer.from_pretrained(
+                    checkpoint, trust_remote_code=is_preset,
+                )
+                if tokenizer.pad_token is None:
+                    tokenizer.pad_token = tokenizer.eos_token
+                with _lock:
+                    _state["model"] = model
+                    _state["tokenizer"] = tokenizer
+            except Exception:
+                yield "Model failed to reload from checkpoint. Try re-obliterating."
+                return
+        else:
+            yield "Model tensors are stale (ZeroGPU). Re-obliterate to create a fresh checkpoint."
+            return
 
     # Sanitize inputs to prevent resource exhaustion
     system_prompt = (system_prompt or "")[:4096]
@@ -1984,6 +2021,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
                 _state["tokenizer"] = tokenizer_loaded
                 _state["steering"] = None
                 _state["status"] = "ready"
+                _state["output_dir"] = checkpoint_dir
             progress(1.0, desc="Ready!")
             yield (
                 f"**Loaded!** `{choice}` is ready in the Chat tab (loaded from checkpoint).",
@@ -2019,6 +2057,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
                     _state["tokenizer"] = tokenizer_loaded
                     _state["steering"] = None
                     _state["status"] = "ready"
+                    _state["output_dir"] = checkpoint_dir
                 progress(1.0, desc="Ready!")
                 yield (
                     f"**Loaded!** `{choice}` is ready in the Chat tab (4-bit from checkpoint).",
@@ -2093,6 +2132,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
         _state["tokenizer"] = pipeline.handle.tokenizer
         _state["steering"] = None
         _state["status"] = "ready"
+        _state["output_dir"] = "/tmp/obliterated"  # re-abliteration fallback path
 
     pipeline_ref[0] = None
 
@@ -2132,13 +2172,35 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[
                "#### Abliterated")
         return
 
-    # ZeroGPU safety: ensure model is on GPU if available
+    # ZeroGPU safety: ensure model is on GPU if available.
+    # If tensors are stale from a prior GPU context, reload from checkpoint.
+    _needs_reload = False
     try:
         dev = next(abliterated_model.parameters()).device
         if torch.cuda.is_available() and dev.type != "cuda":
             abliterated_model.to("cuda")
     except (StopIteration, RuntimeError):
-        pass
+        _needs_reload = True
+
+    if _needs_reload and _ZEROGPU_AVAILABLE:
+        checkpoint = _state.get("output_dir")
+        if checkpoint and Path(checkpoint).exists():
+            try:
+                is_preset = (model_name or "") in MODELS
+                abliterated_model = AutoModelForCausalLM.from_pretrained(
+                    checkpoint, device_map="auto", torch_dtype=torch.float16,
+                    trust_remote_code=is_preset,
+                )
+                tokenizer = AutoTokenizer.from_pretrained(
+                    checkpoint, trust_remote_code=is_preset,
+                )
+                if tokenizer.pad_token is None:
+                    tokenizer.pad_token = tokenizer.eos_token
+                with _lock:
+                    _state["model"] = abliterated_model
+                    _state["tokenizer"] = tokenizer
+            except Exception:
+                pass  # Fall through — will fail at generation with a clear error
 
     # Build header strings showing model name on each side
     header_left = f"#### Original (Pre-Abliteration)\n`{model_name}`"
@@ -3935,17 +3997,16 @@ Built on the shoulders of:
     )
 
     # Wire obliterate button (after all tabs so chat_status is defined)
+    # session_model_dd is a direct output (4th) so the dropdown updates
+    # reliably even on ZeroGPU where .then() may not fire after generator teardown.
     obliterate_btn.click(
         fn=obliterate,
         inputs=[model_dd, method_dd, hub_repo, prompt_vol_dd, dataset_dd,
                 custom_harmful_tb, custom_harmless_tb] + _adv_controls,
-        outputs=[status_md, log_box, chat_status],
+        outputs=[status_md, log_box, chat_status, session_model_dd],
     ).then(
-        fn=lambda: (
-            gr.update(choices=_get_session_model_choices(), value=_last_obliterated_label or None),
-            _get_vram_html(),
-        ),
-        outputs=[session_model_dd, vram_display],
+        fn=lambda: _get_vram_html(),
+        outputs=[vram_display],
     )
 
     # Wire session model loading (Chat tab)

obliteratus/abliterate.py

@@ -1595,6 +1595,38 @@ class AbliterationPipeline:
             else:
                 self._strong_layers = knee_layers
 
+        # ── Small-model safeguards ────────────────────────────────────
+        # Models with few layers (≤16) are highly sensitive to ablation.
+        # Projecting early layers (0, 1) or too many layers relative to
+        # total can destroy the model's ability to produce coherent output.
+        #
+        # Guard 1: Exclude the first 2 layers (layers 0 and 1) — these
+        #   encode fundamental token representations, not refusal.
+        #   The COSMIC cosine-similarity metric often selects layer 0
+        #   because it naturally has divergent harmful/harmless representations
+        #   (token-level differences, not refusal-specific).
+        # Guard 2: Cap selected layers to at most 25% of total for small
+        #   models (≤16 layers).  Even 30% ablation of a 12-layer model
+        #   causes decoherence.
+        if self._strong_layers and n_layers > 0:
+            min_safe_layer = min(2, n_layers // 4)  # layers 0..(min_safe-1) are off-limits
+            early_excluded = [idx for idx in self._strong_layers if idx < min_safe_layer]
+            if early_excluded:
+                self._strong_layers = [idx for idx in self._strong_layers if idx >= min_safe_layer]
+                self.log(
+                    f"Excluded early layers {early_excluded} from ablation "
+                    f"(first {min_safe_layer} layers encode fundamental representations)"
+                )
+
+            if n_layers <= 16 and len(self._strong_layers) > 0:
+                max_small_model_layers = max(1, int(n_layers * 0.25))
+                if len(self._strong_layers) > max_small_model_layers:
+                    self._strong_layers = self._strong_layers[:max_small_model_layers]
+                    self.log(
+                        f"Capped to {max_small_model_layers} layers for small model "
+                        f"(25% of {n_layers} layers)"
+                    )
+
         # Cap layer count for inversion modes — reflecting too many weak-signal
         # layers destroys coherence.  Limit to top 40% of total layers.
         if self.invert_refusal and len(self._strong_layers) > 0:
@@ -3443,6 +3475,15 @@ class AbliterationPipeline:
         sorted_layers = sorted(norms.items(), key=lambda x: x[1], reverse=True)
         self._strong_layers = self._select_layers_knee(sorted_layers)
 
+        # Apply small-model safeguards (matching _distill)
+        if self._strong_layers and n_layers > 0:
+            min_safe_layer = min(2, n_layers // 4)
+            self._strong_layers = [idx for idx in self._strong_layers if idx >= min_safe_layer]
+            if n_layers <= 16 and len(self._strong_layers) > 0:
+                max_small = max(1, int(n_layers * 0.25))
+                if len(self._strong_layers) > max_small:
+                    self._strong_layers = self._strong_layers[:max_small]
+
         # Re-apply jailbreak-contrastive blending with data-driven alpha
         if self.use_jailbreak_contrast and self._jailbreak_means:
             for idx in self._strong_layers:
@@ -3593,8 +3634,10 @@ class AbliterationPipeline:
                 inputs = {k: v.to(device) for k, v in inputs.items()}
                 with torch.no_grad():
                     outputs = model(**inputs, labels=inputs["input_ids"])
-                    total_loss += outputs.loss.item() * inputs["input_ids"].shape[1]
-                    n_tokens += inputs["input_ids"].shape[1]
+                    loss_val = outputs.loss.item()
+                    if not math.isnan(loss_val) and not math.isinf(loss_val):
+                        total_loss += loss_val * inputs["input_ids"].shape[1]
+                        n_tokens += inputs["input_ids"].shape[1]
                 del inputs, outputs
         except Exception:
             pass
@@ -3860,7 +3903,9 @@ class AbliterationPipeline:
             original_norm = saved_norms[param_name]
             if original_norm > 0:
                 new_norm = param.data.norm().item()
-                if new_norm > 0 and abs(new_norm - original_norm) > 1e-6:
+                if math.isnan(new_norm) or math.isinf(new_norm) or new_norm == 0:
+                    continue  # Skip — weight is degenerate after projection
+                if abs(new_norm - original_norm) > 1e-6:
                     param.data.mul_(original_norm / new_norm)
 
     @staticmethod
@@ -3896,6 +3941,10 @@ class AbliterationPipeline:
             W, is_quantized = AbliterationPipeline._dequantize_weight(proj)
             d = direction.to(device=W.device, dtype=W.dtype)
 
+            # Skip projection if weight or direction contains NaN/Inf
+            if not torch.isfinite(W).all() or not torch.isfinite(d).all():
+                continue
+
             if W.shape[-1] == d.shape[0]:
                 # Standard Linear: W is (out_features, hidden_dim)
                 original_norm_sq = W.pow(2).sum().item() if norm_preserve else 0.0
@@ -3905,6 +3954,15 @@ class AbliterationPipeline:
                 W.sub_(d.T * (scale * coeff))      # in-place rank-1 update
                 del coeff
 
+                # Verify projection didn't produce NaN (can happen with
+                # degenerate weights on small models)
+                if not torch.isfinite(W).all():
+                    # Revert: re-load from the Linear module's original data
+                    # This shouldn't normally happen but guards against
+                    # numerical catastrophe
+                    W.copy_(proj.weight.data if not is_quantized else W)
+                    continue
+
                 # Analytical norm: ||W'||² = ||W||² - scale(2-scale)||coeff||²
                 if norm_preserve and original_norm_sq > 0:
                     new_norm_sq = max(0.0, original_norm_sq - scale * (2 - scale) * coeff_norm_sq)
@@ -3926,6 +3984,11 @@ class AbliterationPipeline:
                 W.sub_((scale * d) * coeff)        # in-place rank-1 update
                 del coeff
 
+                # Verify projection didn't produce NaN
+                if not torch.isfinite(W).all():
+                    W.copy_(proj.weight.data if not is_quantized else W)
+                    continue
+
                 # Analytical norm: ||W'||² = ||W||² - scale(2-scale)||coeff||²
                 if norm_preserve and original_norm_sq > 0:
                     new_norm_sq = max(0.0, original_norm_sq - scale * (2 - scale) * coeff_norm_sq)
@@ -4902,24 +4965,42 @@ class AbliterationPipeline:
         self.log("Measuring perplexity on reference texts...")
         total_loss = 0.0
         n_tokens = 0
+        has_nan_loss = False
         for text in reference_texts:
             inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=self.max_seq_length or 256)
             inputs = {k: v.to(device) for k, v in inputs.items()}
             with torch.no_grad():
                 outputs = model(**inputs, labels=inputs["input_ids"])
+                loss_val = outputs.loss.item()
                 seq_len = inputs["input_ids"].shape[1]
-                total_loss += outputs.loss.item() * seq_len
-                n_tokens += seq_len
+                if math.isnan(loss_val) or math.isinf(loss_val):
+                    has_nan_loss = True
+                else:
+                    total_loss += loss_val * seq_len
+                    n_tokens += seq_len
             del inputs, outputs
         self._free_gpu_memory()
 
-        avg_loss = total_loss / n_tokens if n_tokens > 0 else float("inf")
-        try:
-            perplexity = math.exp(min(avg_loss, 100.0))  # clamp to avoid OverflowError
-        except OverflowError:
+        if has_nan_loss and n_tokens == 0:
+            # All reference texts produced NaN loss — model is completely broken
             perplexity = float("inf")
+            self.log("  Perplexity: inf (model produces NaN outputs — weights may be destroyed)")
+        elif has_nan_loss:
+            # Some texts produced NaN — compute from valid ones but warn
+            avg_loss = total_loss / n_tokens
+            try:
+                perplexity = math.exp(min(avg_loss, 100.0))
+            except OverflowError:
+                perplexity = float("inf")
+            self.log(f"  Perplexity: {perplexity:.2f} (WARNING: some reference texts produced NaN loss)")
+        else:
+            avg_loss = total_loss / n_tokens if n_tokens > 0 else float("inf")
+            try:
+                perplexity = math.exp(min(avg_loss, 100.0))  # clamp to avoid OverflowError
+            except OverflowError:
+                perplexity = float("inf")
+            self.log(f"  Perplexity: {perplexity:.2f}")
         self._quality_metrics["perplexity"] = perplexity
-        self.log(f"  Perplexity: {perplexity:.2f}")
 
         # 2. Generation coherence test
         test_prompts = [
@@ -5181,24 +5262,35 @@ class AbliterationPipeline:
                 post_logits = torch.cat(all_post_logits, dim=0)
                 pre_logits = self._baseline_first_token_logits[:post_logits.shape[0]]
 
-                # Use F.kl_div for numerical stability
-                log_p = torch.nn.functional.log_softmax(pre_logits.float(), dim=-1)
-                log_q = torch.nn.functional.log_softmax(post_logits.float(), dim=-1)
-                kl_per_prompt = torch.nn.functional.kl_div(
-                    log_q, log_p, log_target=True, reduction="none"
-                ).sum(dim=-1).clamp(min=0.0)
-                kl_divergence = kl_per_prompt.mean().item()
-
-                self._quality_metrics["kl_divergence"] = kl_divergence
-                if kl_divergence < 0.2:
-                    kl_label = "excellent"
-                elif kl_divergence < 0.5:
-                    kl_label = "good"
-                elif kl_divergence < 1.0:
-                    kl_label = "moderate"
+                # Check for NaN/Inf in post-ablation logits (model may be broken)
+                if torch.isnan(post_logits).any() or torch.isinf(post_logits).any():
+                    self.log("  KL divergence: inf (model produces NaN/Inf logits — weights may be destroyed)")
+                    kl_divergence = float("inf")
+                    self._quality_metrics["kl_divergence"] = kl_divergence
                 else:
-                    kl_label = "high"
-                self.log(f"  First-token KL divergence: {kl_divergence:.4f} ({kl_label})")
+                    # Use F.kl_div for numerical stability
+                    log_p = torch.nn.functional.log_softmax(pre_logits.float(), dim=-1)
+                    log_q = torch.nn.functional.log_softmax(post_logits.float(), dim=-1)
+                    kl_per_prompt = torch.nn.functional.kl_div(
+                        log_q, log_p, log_target=True, reduction="none"
+                    ).sum(dim=-1).clamp(min=0.0)
+                    kl_divergence = kl_per_prompt.mean().item()
+
+                    # Guard against NaN from numerical issues in KL computation
+                    if math.isnan(kl_divergence) or math.isinf(kl_divergence):
+                        kl_divergence = float("inf")
+                        self.log("  First-token KL divergence: inf (numerical overflow — model may be severely damaged)")
+                    else:
+                        if kl_divergence < 0.2:
+                            kl_label = "excellent"
+                        elif kl_divergence < 0.5:
+                            kl_label = "good"
+                        elif kl_divergence < 1.0:
+                            kl_label = "moderate"
+                        else:
+                            kl_label = "high"
+                        self.log(f"  First-token KL divergence: {kl_divergence:.4f} ({kl_label})")
+                    self._quality_metrics["kl_divergence"] = kl_divergence
             except Exception as e:
                 self.log(f"  KL divergence computation failed (non-fatal): {e}")
                 self._quality_metrics["kl_divergence"] = None