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

Add files via upload

Browse Source
This commit is contained in:
pliny
2026-03-05 00:50:44 -08:00
committed by GitHub
parent 4cddc6399a
commit 66ea4a6f86
37 changed files with 10756 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+50 -25
View File
@@ -115,6 +115,10 @@ _last_obliterated_label: str = ""
# Counter for unique obliteration save directories
_obliterate_counter: int = 0
# Flag to suppress session_model_dd.change when obliterate programmatically
# sets the dropdown value (prevents wasteful GPU re-allocation on ZeroGPU)
_skip_session_load: bool = False
# ---------------------------------------------------------------------------
# Model presets — 100+ models organized by provider
# ---------------------------------------------------------------------------
@@ -1459,7 +1463,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(),
"", gr.update(), gr.update(), gr.update(), gr.update(),
)
return
@@ -1468,14 +1472,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(),
"", gr.update(), 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(),
"", gr.update(), gr.update(), gr.update(), gr.update(),
)
return
@@ -1486,7 +1490,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(), gr.update(), gr.update()
yield "**Error:** An obliteration is already in progress.", "", gr.update(), gr.update(), gr.update(), gr.update()
return
_state["log"] = []
_state["status"] = "obliterating"
@@ -1638,9 +1642,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(), gr.update(), gr.update()
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
else:
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update()
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
if time.time() - _pipeline_start > _max_pipeline_secs:
log_lines.append("\nTIMEOUT: Pipeline exceeded 45-minute limit.")
break
@@ -1655,7 +1659,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(), gr.update(), gr.update()
yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update()
return
# Success — keep model in memory for chat.
@@ -1757,7 +1761,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(), gr.update(), gr.update()
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
try:
from transformers import BitsAndBytesConfig
bnb_cfg = BitsAndBytesConfig(
@@ -1804,7 +1808,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(), gr.update(), gr.update()
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(
@@ -1861,13 +1865,21 @@ 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."
)
# Update session dropdown directly (don't rely on .then() which can
# fail to fire on ZeroGPU after generator teardown)
# Update BOTH session dropdowns directly (don't rely on .then() which
# fails to fire on ZeroGPU after generator teardown).
# Set skip flag so the .change handler doesn't trigger a wasteful
# GPU re-allocation — the model is already loaded.
global _skip_session_load
_skip_session_load = True
_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, metrics_card
_ab_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, metrics_card, _ab_dd_update
except Exception as e:
# Ensure status never gets stuck on "obliterating"
@@ -1876,7 +1888,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(), gr.update(), gr.update()
yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update()
# ---------------------------------------------------------------------------
@@ -2102,6 +2114,18 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
On ZeroGPU, uses the visitor's GPU quota.
"""
# Skip if the obliterate function just set the dropdown value — the model
# is already loaded and we'd just waste GPU quota re-allocating.
global _skip_session_load
if _skip_session_load:
_skip_session_load = False
if choice and _state.get("status") == "ready":
yield (
f"**Ready!** `{choice}` is loaded — just type in the chat below.",
get_chat_header(),
)
return
if not choice or choice not in _bench_configs:
yield "**Error:** No benchmark result selected. Pick a model from the dropdown first.", ""
return
@@ -3727,6 +3751,7 @@ Pre-configured benchmark configurations for common research questions.
choices=_get_session_model_choices(),
label="Cached Models",
info="Select a model to auto-load it for chat",
allow_custom_value=True,
)
session_load_status = gr.Markdown("")
@@ -3779,6 +3804,7 @@ See exactly how abliteration changes model behavior on the same prompt.
choices=_get_session_model_choices(),
label="Cached Models",
info="Select a model to auto-load it for A/B comparison",
allow_custom_value=True,
)
ab_session_load_status = gr.Markdown("")
@@ -4125,8 +4151,8 @@ Built on the shoulders of:
### Links
- [GitHub](https://github.com/obliteratus-project/OBLITERATUS)
- [Paper](https://github.com/obliteratus-project/OBLITERATUS/tree/main/paper)
- [GitHub](https://github.com/elder-plinius/OBLITERATUS)
- [Paper](https://github.com/elder-plinius/OBLITERATUS/tree/main/paper)
""")
# Wire method dropdown → auto-update advanced settings
@@ -4192,28 +4218,27 @@ Built on the shoulders of:
).then(fn=_get_vram_html, outputs=[vram_display])
# 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.
# Both session_model_dd (4th) and ab_session_model_dd (6th) are direct
# outputs so the dropdowns update 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, session_model_dd, metrics_md],
outputs=[status_md, log_box, chat_status, session_model_dd, metrics_md, ab_session_model_dd],
).then(
fn=lambda: (
gr.update(choices=_get_session_model_choices()),
_get_vram_html(),
),
outputs=[ab_session_model_dd, vram_display],
fn=lambda: _get_vram_html(),
outputs=[vram_display],
)
# Wire session model auto-loading (Chat tab dropdown change)
# Always pass choices + value together so ZeroGPU doesn't hit stale choices
session_model_dd.change(
fn=load_bench_into_chat,
inputs=[session_model_dd],
outputs=[session_load_status, chat_status],
).then(
fn=lambda v: (gr.update(value=v), _get_vram_html()),
fn=lambda v: (gr.update(choices=_get_session_model_choices(), value=v), _get_vram_html()),
inputs=[session_model_dd],
outputs=[ab_session_model_dd, vram_display],
)
@@ -4224,7 +4249,7 @@ Built on the shoulders of:
inputs=[ab_session_model_dd],
outputs=[ab_session_load_status, chat_status],
).then(
fn=lambda v: (gr.update(value=v), _get_vram_html()),
fn=lambda v: (gr.update(choices=_get_session_model_choices(), value=v), _get_vram_html()),
inputs=[ab_session_model_dd],
outputs=[session_model_dd, vram_display],
)