CalvinBackup/OBLITERATUS
1
0
Fork 0
mirror of https://github.com/elder-plinius/OBLITERATUS.git synced 2026-07-13 07:36:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
pliny
2026-03-08 12:07:56 -07:00
committed by GitHub
parent 1065809658
commit 69fa63ac43
14 changed files with 1320 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+219 -43
View File
@@ -98,6 +98,51 @@ def _is_quota_error(exc: BaseException) -> bool:
return True
return False
def _load_model_to_device(
pretrained_path: str,
*,
torch_dtype=None,
trust_remote_code: bool = False,
quantization_config=None,
offload_folder: str | None = None,
low_cpu_mem_usage: bool = False,
token: str | None = None,
) -> AutoModelForCausalLM:
"""Load a causal LM onto the best available device, MPS-safe.
Accelerate's ``device_map="auto"`` is not supported on MPS — models
silently land on CPU. This helper skips ``device_map`` on non-CUDA
backends and explicitly moves the model to the best device after loading.
On CUDA the behaviour is identical to ``device_map="auto"``.
"""
kwargs: dict = {}
if torch_dtype is not None:
kwargs["torch_dtype"] = torch_dtype
if trust_remote_code:
kwargs["trust_remote_code"] = True
if quantization_config is not None:
kwargs["quantization_config"] = quantization_config
if offload_folder is not None:
kwargs["offload_folder"] = offload_folder
if low_cpu_mem_usage:
kwargs["low_cpu_mem_usage"] = True
if token is not None:
kwargs["token"] = token
if dev.supports_device_map_auto():
kwargs["device_map"] = "auto"
model = AutoModelForCausalLM.from_pretrained(pretrained_path, **kwargs)
# On MPS / CPU: model loaded without device_map, move to best device
if not dev.supports_device_map_auto():
target = dev.get_device()
model = model.to(target)
return model
# ---------------------------------------------------------------------------
# Global state
# ---------------------------------------------------------------------------
@@ -164,7 +209,7 @@ def _recover_sessions_from_disk() -> None:
"""
global _last_obliterated_label, _obliterate_counter
found_any = False
for pattern in ("obliterated_*", "obliterated", "bench_*"):
for pattern in ("obliterated_*", "obliterated", "bench_*", "obliteratus_tourney/r*"):
for p in Path("/tmp").glob(pattern):
if not p.is_dir():
continue
@@ -291,6 +336,11 @@ METHODS = {
"optimized (bayesian auto-tuned)": "optimized",
"inverted (semantic refusal inversion)": "inverted",
"nuclear (maximum force combo)": "nuclear",
# Baseline reproductions for benchmarking
"failspy (FailSpy/abliterator baseline)": "failspy",
"gabliteration (Gülmez 2026 baseline)": "gabliteration",
"heretic (p-e-w 2025-2026 baseline)": "heretic",
"rdo (Wollschlager ICML 2025 baseline)": "rdo",
}
# ── Community Hub push ────────────────────────────────────────────────
@@ -316,6 +366,7 @@ def _get_preset_defaults(method_display: str):
cfg = _PRESET_CONFIGS.get(method_key, _PRESET_CONFIGS["advanced"])
return {
"n_directions": cfg.get("n_directions", 4),
"direction_method": cfg.get("direction_method", "svd"),
"regularization": cfg.get("regularization", 0.3),
"refinement_passes": cfg.get("refinement_passes", 2),
"norm_preserve": cfg.get("norm_preserve", True),
@@ -341,6 +392,17 @@ def _get_preset_defaults(method_display: str):
"spectral_cascade": cfg.get("spectral_cascade", False),
"spectral_bands": cfg.get("spectral_bands", 3),
"spectral_threshold": cfg.get("spectral_threshold", 0.05),
# Baseline-specific parameters
"layer_selection": cfg.get("layer_selection", "all"),
"winsorize_activations": cfg.get("winsorize_activations", False),
"winsorize_percentile": cfg.get("winsorize_percentile", 1.0),
"use_kl_optimization": cfg.get("use_kl_optimization", False),
"kl_budget": cfg.get("kl_budget", 0.5),
"float_layer_interpolation": cfg.get("float_layer_interpolation", False),
"rdo_refinement": cfg.get("rdo_refinement", False),
"cot_aware": cfg.get("cot_aware", False),
"bayesian_trials": cfg.get("bayesian_trials", 50),
"n_sae_features": cfg.get("n_sae_features", 64),
}
def _on_method_change(method_display: str):
@@ -348,6 +410,7 @@ def _on_method_change(method_display: str):
d = _get_preset_defaults(method_display)
return (
d["n_directions"],
d["direction_method"],
d["regularization"],
d["refinement_passes"],
d["reflection_strength"],
@@ -374,6 +437,16 @@ def _on_method_change(method_display: str):
d["expert_transplant"],
d["use_wasserstein_optimal"],
d["spectral_cascade"],
d["layer_selection"],
d["winsorize_activations"],
d["winsorize_percentile"],
d["use_kl_optimization"],
d["kl_budget"],
d["float_layer_interpolation"],
d["rdo_refinement"],
d["cot_aware"],
d["bayesian_trials"],
d["n_sae_features"],
)
def _on_dataset_change(dataset_label: str):
@@ -1731,8 +1804,9 @@ def _format_multi_model_results(results: list[dict], context: dict | None = None
def obliterate(model_choice: str, method_choice: str,
prompt_volume_choice: str, dataset_source_choice: str,
custom_harmful: str, custom_harmless: str,
# Advanced params (sliders)
adv_n_directions: int, adv_regularization: float,
# Advanced params (sliders + radio)
adv_n_directions: int, adv_direction_method: str,
adv_regularization: float,
adv_refinement_passes: int, adv_reflection_strength: float,
adv_embed_regularization: float, adv_steering_strength: float,
adv_transplant_blend: float,
@@ -1748,6 +1822,12 @@ def obliterate(model_choice: str, method_choice: str,
adv_project_embeddings: bool, adv_activation_steering: bool,
adv_expert_transplant: bool, adv_wasserstein_optimal: bool,
adv_spectral_cascade: bool,
adv_layer_selection: str, adv_winsorize: bool,
adv_winsorize_percentile: float,
adv_kl_optimization: bool, adv_kl_budget: float,
adv_float_layer_interp: bool, adv_rdo_refinement: bool,
adv_cot_aware: bool,
adv_bayesian_trials: int, adv_n_sae_features: int,
progress=gr.Progress()):
"""Run the full obliteration pipeline, streaming log updates to the UI.
@@ -1906,6 +1986,7 @@ def obliterate(model_choice: str, method_choice: str,
on_log=on_log,
# Advanced overrides from UI
n_directions=int(adv_n_directions),
direction_method=adv_direction_method,
regularization=float(adv_regularization),
refinement_passes=int(adv_refinement_passes),
norm_preserve=adv_norm_preserve,
@@ -1932,6 +2013,15 @@ def obliterate(model_choice: str, method_choice: str,
spectral_bands=int(adv_spectral_bands),
spectral_threshold=float(adv_spectral_threshold),
verify_sample_size=int(adv_verify_sample_size),
layer_selection=adv_layer_selection,
winsorize_activations=adv_winsorize,
winsorize_percentile=float(adv_winsorize_percentile),
use_kl_optimization=adv_kl_optimization,
kl_budget=float(adv_kl_budget),
float_layer_interpolation=adv_float_layer_interp,
rdo_refinement=adv_rdo_refinement,
cot_aware=adv_cot_aware,
n_sae_features=int(adv_n_sae_features),
)
pipeline_ref[0] = pipeline
pipeline.run()
@@ -2103,10 +2193,9 @@ def obliterate(model_choice: str, method_choice: str,
bnb_4bit_quant_type="nf4",
llm_int8_enable_fp32_cpu_offload=True,
)
model_reloaded = AutoModelForCausalLM.from_pretrained(
model_reloaded = _load_model_to_device(
save_dir,
quantization_config=bnb_cfg,
device_map="auto",
trust_remote_code=True,
)
tokenizer_reloaded = AutoTokenizer.from_pretrained(
@@ -2144,9 +2233,8 @@ def obliterate(model_choice: str, method_choice: str,
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
try:
offload_dir = tempfile.mkdtemp(prefix="obliteratus_offload_")
model_reloaded = AutoModelForCausalLM.from_pretrained(
model_reloaded = _load_model_to_device(
save_dir,
device_map="auto",
offload_folder=offload_dir,
torch_dtype=torch.float16,
trust_remote_code=True,
@@ -2307,8 +2395,8 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
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,
model = _load_model_to_device(
checkpoint, torch_dtype=torch.float16,
trust_remote_code=is_preset,
)
tokenizer = AutoTokenizer.from_pretrained(
@@ -2498,8 +2586,8 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
if checkpoint and Path(checkpoint).exists():
is_preset = (_state.get("model_name") or "") in MODELS
try:
model_loaded = AutoModelForCausalLM.from_pretrained(
checkpoint, device_map="auto", torch_dtype=torch.float16,
model_loaded = _load_model_to_device(
checkpoint, torch_dtype=torch.float16,
trust_remote_code=is_preset,
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
@@ -2559,9 +2647,8 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
is_preset = cfg["model_choice"] in MODELS
try:
model_loaded = AutoModelForCausalLM.from_pretrained(
model_loaded = _load_model_to_device(
checkpoint_dir,
device_map="auto",
torch_dtype=torch.float16,
trust_remote_code=is_preset,
)
@@ -2595,10 +2682,9 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
)
yield f"**Loading {choice}** in 4-bit (model too large for fp16)...", ""
progress(0.5, desc="Loading 4-bit...")
model_loaded = AutoModelForCausalLM.from_pretrained(
model_loaded = _load_model_to_device(
checkpoint_dir,
quantization_config=bnb_cfg,
device_map="auto",
trust_remote_code=is_preset,
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
@@ -2740,8 +2826,8 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[
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,
abliterated_model = _load_model_to_device(
checkpoint, torch_dtype=torch.float16,
trust_remote_code=is_preset,
)
tokenizer = AutoTokenizer.from_pretrained(
@@ -2866,10 +2952,9 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[
is_preset = model_name in MODELS
original_response = ""
try:
from transformers import AutoModelForCausalLM as AMCLM
original_model = AMCLM.from_pretrained(
original_model = _load_model_to_device(
model_id, torch_dtype=torch.float16,
device_map="auto", trust_remote_code=is_preset,
trust_remote_code=is_preset,
low_cpu_mem_usage=True,
token=os.environ.get("HF_TOKEN") or None,
)
@@ -3026,6 +3111,9 @@ def strength_sweep(model_choice: str, method_choice: str,
entry["perplexity"] = metrics.get("perplexity")
entry["refusal_rate"] = metrics.get("refusal_rate")
entry["coherence"] = metrics.get("coherence")
entry["kl_divergence"] = metrics.get("kl_divergence")
entry["spectral_cert"] = metrics.get("spectral_certification") or ""
entry["direction_method"] = getattr(pipe, "direction_method", "")
entry["strong_layers"] = len(pipe._strong_layers)
if hasattr(pipe, "handle") and pipe.handle is not None:
pipe.handle.model = None
@@ -3115,17 +3203,21 @@ def _format_sweep_results(results: list[dict]) -> str:
return "*No results yet.*"
lines = ["### Strength Sweep Results", "",
"| Reg | Time | Perplexity | Refusal Rate | Coherence | Error |",
"|-----|------|-----------|-------------|-----------|-------|"]
"| Reg | Dir | Time | PPL | Refusal | Coherence | KL Div | Cert | Error |",
"|-----|-----|------|-----|---------|-----------|--------|------|-------|"]
for r in results:
reg = f"{r['regularization']:.3f}"
ppl = f"{r['perplexity']:.2f}" if r.get("perplexity") is not None else "—"
ref = f"{r['refusal_rate']:.0%}" if r.get("refusal_rate") is not None else "—"
coh = f"{r['coherence']:.0%}" if r.get("coherence") is not None else "—"
kl_val = r.get("kl_divergence")
kl_str = f"{kl_val:.4f}" if kl_val is not None else "—"
cert = r.get("spectral_cert", "") or "—"
dir_m = r.get("direction_method", "") or "—"
err = r.get("error", "")
err_short = (err[:25] + "...") if err and len(err) > 25 else (err or "")
lines.append(f"| {reg} | {r['time_s']}s | {ppl} | {ref} | {coh} | {err_short} |")
lines.append(f"| {reg} | {dir_m} | {r['time_s']}s | {ppl} | {ref} | {coh} | {kl_str} | {cert} | {err_short} |")
return "\n".join(lines)
@@ -3173,8 +3265,8 @@ def _tourney_gpu_wrapper(fn, *args, **kwargs):
return _tourney_gpu_run(fn, *args, **kwargs)
def run_tourney(model_choice, dataset, quantization):
"""Run an elimination tournament across all abliteration methods.
def run_tourney(model_choice, selected_methods, dataset, quantization):
"""Run an elimination tournament across selected abliteration methods.
Each individual method is run inside its own ``@spaces.GPU`` allocation
(up to 5 minutes per method) so the full tournament is not constrained
@@ -3187,6 +3279,10 @@ def run_tourney(model_choice, dataset, quantization):
yield "**Error:** Select a model first.", "", ""
return
if not selected_methods or len(selected_methods) < 3:
yield "**Error:** Select at least 3 methods for a tournament.", "", ""
return
from obliteratus.tourney import (
TourneyRunner, render_bracket_html,
_load_checkpoint, _checkpoint_matches,
@@ -3218,6 +3314,7 @@ def run_tourney(model_choice, dataset, quantization):
hub_repo=None,
dataset_key=dataset_key,
quantization=quant,
methods=list(selected_methods),
on_log=logger,
resume=resume,
)
@@ -3322,18 +3419,27 @@ def run_tourney(model_choice, dataset, quantization):
_ts = datetime.now().strftime("%H:%M")
_short = model_id.split("/")[-1] if "/" in model_id else model_id
_label = f"tourney winner ({winner.method}) on {_short} ({_ts})"
_winner_meta = {
"model_id": model_id,
"model_choice": model_choice,
"method": winner.method,
"dataset_key": dataset_key,
"prompt_volume": 0,
"output_dir": winner.output_dir,
"source": "tourney",
"tourney_score": winner.score,
"tourney_metrics": winner.metrics,
}
with _lock:
_session_models[_label] = {
"model_id": model_id,
"model_choice": model_choice,
"method": winner.method,
"dataset_key": dataset_key,
"prompt_volume": 0,
"output_dir": winner.output_dir,
"source": "tourney",
"tourney_score": winner.score,
"tourney_metrics": winner.metrics,
}
_session_models[_label] = _winner_meta
# Persist so the winner survives ZeroGPU process restarts
_persist_session_meta(winner.output_dir, _label, {
"model_id": model_id,
"model_choice": model_choice,
"method": winner.method,
"dataset_key": dataset_key,
"source": "tourney",
})
yield (
f"**Champion: `{winner.method}`** "
f"(score: {winner.score:.4f})\n"
@@ -3930,7 +4036,13 @@ with gr.Blocks(theme=THEME, css=CSS, js=_JS, title="OBLITERATUS", fill_height=Tr
with gr.Row():
adv_n_directions = gr.Slider(
1, 8, value=_defaults["n_directions"], step=1,
label="Directions", info="Number of refusal directions to extract via SVD",
label="Directions", info="Number of refusal directions to extract",
)
adv_direction_method = gr.Radio(
choices=["diff_means", "svd", "leace"],
value=_defaults["direction_method"],
label="Direction Method",
info="diff_means: simple & robust, svd: multi-direction, leace: optimal erasure",
)
adv_regularization = gr.Slider(
0.0, 1.0, value=_defaults["regularization"], step=0.05,
@@ -3996,10 +4108,52 @@ with gr.Blocks(theme=THEME, css=CSS, js=_JS, title="OBLITERATUS", fill_height=Tr
with gr.Row():
adv_spectral_cascade = gr.Checkbox(value=_defaults["spectral_cascade"], label="Spectral Cascade",
info="DCT frequency decomposition for precision refusal targeting")
gr.Markdown("**Layer Selection & Baseline Options**")
with gr.Row():
adv_layer_selection = gr.Dropdown(
choices=["knee_cosmic", "all", "all_except_first", "middle60", "top_k", "knee"],
value=_defaults["layer_selection"],
label="Layer Selection",
info="Which layers to project refusal directions from",
)
adv_winsorize_percentile = gr.Slider(
0.0, 1.0, value=_defaults["winsorize_percentile"], step=0.01,
label="Winsorize Percentile",
info="Activation clamping quantile (1.0 = disabled, 0.01 = 99th pctile)",
)
adv_kl_budget = gr.Slider(
0.0, 2.0, value=_defaults["kl_budget"], step=0.1,
label="KL Budget",
info="Max KL divergence from base model (Heretic/optimized)",
)
with gr.Row():
adv_winsorize = gr.Checkbox(value=_defaults["winsorize_activations"], label="Winsorize Activations",
info="Clamp outlier activations before direction extraction")
adv_kl_optimization = gr.Checkbox(value=_defaults["use_kl_optimization"], label="KL Optimization",
info="Optimize projection strength to stay within KL budget")
adv_float_layer_interp = gr.Checkbox(value=_defaults["float_layer_interpolation"], label="Float Layer Interpolation",
info="Interpolate between adjacent layers' directions (Heretic)")
adv_rdo_refinement = gr.Checkbox(value=_defaults["rdo_refinement"], label="RDO Refinement",
info="Gradient-based direction refinement (Wollschlager et al.)")
with gr.Row():
adv_cot_aware = gr.Checkbox(value=_defaults["cot_aware"], label="CoT-Aware",
info="Preserve chain-of-thought reasoning during abliteration")
with gr.Row():
adv_bayesian_trials = gr.Slider(
10, 200, value=_defaults["bayesian_trials"], step=10,
label="Bayesian Trials",
info="Optuna TPE optimization trials (Heretic/optimized methods)",
)
adv_n_sae_features = gr.Slider(
16, 256, value=_defaults["n_sae_features"], step=16,
label="SAE Features",
info="Number of SAE features to target (inverted/nuclear methods)",
)
# List of all advanced controls (order must match _on_method_change return)
_adv_controls = [
adv_n_directions, adv_regularization, adv_refinement_passes,
adv_n_directions, adv_direction_method,
adv_regularization, adv_refinement_passes,
adv_reflection_strength, adv_embed_regularization,
adv_steering_strength, adv_transplant_blend,
adv_spectral_bands, adv_spectral_threshold,
@@ -4011,6 +4165,12 @@ with gr.Blocks(theme=THEME, css=CSS, js=_JS, title="OBLITERATUS", fill_height=Tr
adv_project_embeddings, adv_activation_steering,
adv_expert_transplant, adv_wasserstein_optimal,
adv_spectral_cascade,
adv_layer_selection, adv_winsorize,
adv_winsorize_percentile,
adv_kl_optimization, adv_kl_budget,
adv_float_layer_interp, adv_rdo_refinement,
adv_cot_aware,
adv_bayesian_trials, adv_n_sae_features,
]
obliterate_btn = gr.Button(
@@ -4181,7 +4341,8 @@ result = client.predict(
mm_method = gr.Dropdown(
choices=["basic", "advanced", "aggressive",
"spectral_cascade", "informed", "surgical",
"optimized", "inverted", "nuclear"],
"optimized", "inverted", "nuclear",
"failspy", "gabliteration", "heretic", "rdo"],
value="surgical",
label="Abliteration Method",
)
@@ -4550,11 +4711,11 @@ tradeoff point where refusal is minimized with minimal capability damage.
# ── Tab 6: Tourney ────────────────────────────────────────────────
with gr.Tab("Tourney", id="tourney"):
gr.Markdown("""### March Madness Tournament
Pit **all abliteration methods** against each other in elimination rounds.
gr.Markdown("""### Tourney Mode
Pit abliteration methods against each other in elimination rounds.
The winner is saved locally — push it to HuggingFace Hub from the **Push to Hub** tab.
**Round 1 — Qualifiers:** All methods, reduced prompts. Bottom half eliminated.
**Round 1 — Qualifiers:** Selected methods, reduced prompts. Bottom half eliminated.
**Round 2 — Semifinals:** Survivors, full prompts. Bottom half eliminated.
**Round 3 — Finals:** Top contenders, maximum prompts. Champion crowned.
""")
@@ -4566,6 +4727,14 @@ The winner is saved locally — push it to HuggingFace Hub from the **Push to Hu
allow_custom_value=True,
)
from obliteratus.tourney import TOURNEY_METHODS as _ALL_TOURNEY_METHODS
tourney_methods_cb = gr.CheckboxGroup(
choices=_ALL_TOURNEY_METHODS,
value=_ALL_TOURNEY_METHODS,
label="Methods to Compete",
info="Pick at least 3 methods. All selected by default.",
)
with gr.Accordion("Advanced Settings", open=False):
with gr.Row():
tourney_dataset_dd = gr.Dropdown(
@@ -4595,9 +4764,16 @@ The winner is saved locally — push it to HuggingFace Hub from the **Push to Hu
tourney_btn.click(
fn=run_tourney,
inputs=[tourney_model_dd,
inputs=[tourney_model_dd, tourney_methods_cb,
tourney_dataset_dd, tourney_quant_dd],
outputs=[tourney_status, tourney_bracket, tourney_log],
).then(
fn=lambda: (
gr.update(choices=_get_session_model_choices()),
gr.update(choices=_get_session_model_choices()),
_get_vram_html(),
),
outputs=[session_model_dd, ab_session_model_dd, vram_display],
)
# ── Tab 7: Export ─────────────────────────────────────────────────