Upload 80 files

app.py

@@ -73,6 +73,29 @@ METHODS = {
     "aggressive (maximum removal)": "aggressive",
 }
 
+# Models that need 4bit quantization to fit on a T4 16GB
+_NEEDS_QUANTIZATION = {
+    "openai/gpt-oss-20b",
+}
+
+
+def _should_quantize(model_id: str) -> str | None:
+    """Return '4bit' if the model needs quantization for available GPU, else None."""
+    if model_id in _NEEDS_QUANTIZATION:
+        return "4bit"
+    # Auto-detect from config for custom model IDs
+    try:
+        from obliteratus.models.loader import _estimate_model_memory_gb, _available_gpu_memory_gb
+        from transformers import AutoConfig
+        config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
+        est_gb = _estimate_model_memory_gb(config, torch.float16)
+        gpu_gb = _available_gpu_memory_gb()
+        if gpu_gb > 0 and est_gb > gpu_gb * 0.85:
+            return "4bit"
+    except Exception:
+        pass
+    return None
+
 
 # ---------------------------------------------------------------------------
 # Obliteration
@@ -118,6 +141,8 @@ def obliterate(model_choice: str, method_choice: str, progress=gr.Progress()):
         idx = stage_order.get(stage_key, 0)
         progress((idx + 1) / 6, desc=f"{stage_key.upper()}")
 
+    quantization = _should_quantize(model_id)
+
     def run_pipeline():
         try:
             from obliteratus.abliterate import AbliterationPipeline
@@ -127,6 +152,7 @@ def obliterate(model_choice: str, method_choice: str, progress=gr.Progress()):
                 device="auto",
                 dtype="float16",
                 method=method,
+                quantization=quantization,
                 on_stage=on_stage,
                 on_log=on_log,
             )
@@ -137,6 +163,8 @@ def obliterate(model_choice: str, method_choice: str, progress=gr.Progress()):
 
     log_lines.append(f"Target: {model_id}")
     log_lines.append(f"Method: {method}")
+    if quantization:
+        log_lines.append(f"Quantization: {quantization} (auto-detected for GPU fit)")
     log_lines.append("")
 
     worker = threading.Thread(target=run_pipeline, daemon=True)

obliteratus/abliterate.py

@@ -252,6 +252,7 @@ class AbliterationPipeline:
         use_chat_template: bool | None = None,
         use_whitened_svd: bool | None = None,
         true_iterative_refinement: bool | None = None,
+        quantization: str | None = None,
         harmful_prompts: list[str] | None = None,
         harmless_prompts: list[str] | None = None,
         on_stage: Callable[[StageResult], None] | None = None,
@@ -278,6 +279,7 @@ class AbliterationPipeline:
         self.use_chat_template = use_chat_template if use_chat_template is not None else method_cfg.get("use_chat_template", False)
         self.use_whitened_svd = use_whitened_svd if use_whitened_svd is not None else method_cfg.get("use_whitened_svd", False)
         self.true_iterative_refinement = true_iterative_refinement if true_iterative_refinement is not None else method_cfg.get("true_iterative_refinement", False)
+        self.quantization = quantization
 
         self.handle: ModelHandle | None = None
         self.refusal_directions: dict[int, torch.Tensor] = {}  # per-layer primary direction
@@ -339,7 +341,7 @@ class AbliterationPipeline:
             device=self.device,
             dtype=self.dtype,
             trust_remote_code=self.trust_remote_code,
-            quantization=getattr(self, "quantization", None),
+            quantization=self.quantization,
         )
 
         summary = self.handle.summary()

obliteratus/cli.py

@@ -71,6 +71,10 @@ def main(argv: list[str] | None = None):
         p.add_argument("--n-directions", type=int, default=None, help="Override: number of SVD directions to extract")
         p.add_argument("--regularization", type=float, default=None, help="Override: fraction to preserve (0.0-1.0)")
         p.add_argument("--refinement-passes", type=int, default=None, help="Override: number of iterative passes")
+        p.add_argument(
+            "--quantization", type=str, default=None, choices=["4bit", "8bit"],
+            help="Load model with quantization (4bit or 8bit). Requires bitsandbytes.",
+        )
 
     abl_parser = subparsers.add_parser(
         "obliterate",
@@ -323,6 +327,7 @@ def _cmd_abliterate(args):
         n_directions=args.n_directions,
         regularization=args.regularization,
         refinement_passes=args.refinement_passes,
+        quantization=args.quantization,
         on_stage=on_stage,
         on_log=on_log,
     )