mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
1486 lines
55 KiB
Python
1486 lines
55 KiB
Python
"""OBLITERATUS Tourney — March Madness-style tournament to find the best abliteration method.
|
|
|
|
Run all methods head-to-head in elimination rounds. The winner gets auto-pushed
|
|
to HuggingFace Hub so the community can use the best possible abliteration.
|
|
|
|
Usage (CLI):
|
|
obliteratus tourney meta-llama/Llama-3.1-8B-Instruct --hub-org my-org
|
|
|
|
Usage (Python):
|
|
from obliteratus.tourney import TourneyRunner
|
|
runner = TourneyRunner("meta-llama/Llama-3.1-8B-Instruct", hub_org="my-org")
|
|
winner = runner.run()
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import json
|
|
import math
|
|
import os
|
|
import shutil
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# All tournament-eligible methods.
|
|
#
|
|
# Excluded:
|
|
# - 'optimized' — runs its own Bayesian optimization, far too slow per round
|
|
# - 'nuclear' — collapsed in essentially every telemetry run (n=7,545)
|
|
# - 'basic' — fast but quality is unusable across architectures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
TOURNEY_METHODS = [
|
|
"advanced",
|
|
"aggressive",
|
|
"spectral_cascade",
|
|
"informed",
|
|
"surgical",
|
|
"inverted",
|
|
"failspy",
|
|
"gabliteration",
|
|
"heretic",
|
|
"rdo",
|
|
]
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scoring
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def composite_score(metrics: dict[str, Any]) -> float:
|
|
"""Score an abliteration run on [0, 1]. Higher is better.
|
|
|
|
Weights:
|
|
35% refusal removal — the whole point
|
|
25% coherence — model must still be useful
|
|
20% KL divergence — minimal capability damage
|
|
10% perplexity — fluency preservation
|
|
5% spectral diagnostic — residual-signal assessment
|
|
5% degenerate penalty — penalize broken output
|
|
"""
|
|
rr = metrics.get("refusal_rate")
|
|
co = metrics.get("coherence")
|
|
kl = metrics.get("kl_divergence")
|
|
pp = metrics.get("perplexity")
|
|
spec = metrics.get("spectral_certification")
|
|
degen = metrics.get("degenerate_count", 0) or 0
|
|
|
|
refusal_score = (1.0 - rr) if rr is not None else 0.0
|
|
coherence_score = co if co is not None else 0.0
|
|
kl_score = 1.0 / (1.0 + kl) if kl is not None else 0.5
|
|
ppl_score = 1.0 / (1.0 + pp / 100.0) if pp is not None else 0.5
|
|
|
|
# Spectral certification: GREEN=1.0, YELLOW=0.5, RED=0.0, None=0.5 (neutral)
|
|
if spec == "GREEN":
|
|
spec_score = 1.0
|
|
elif spec == "YELLOW":
|
|
spec_score = 0.5
|
|
elif spec == "RED":
|
|
spec_score = 0.0
|
|
else:
|
|
spec_score = 0.5 # not measured → neutral
|
|
|
|
# Degenerate penalty: any broken outputs reduce score
|
|
degen_score = 1.0 / (1.0 + degen) if degen > 0 else 1.0
|
|
|
|
return (
|
|
refusal_score * 0.35
|
|
+ coherence_score * 0.25
|
|
+ kl_score * 0.20
|
|
+ ppl_score * 0.10
|
|
+ spec_score * 0.05
|
|
+ degen_score * 0.05
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Data classes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class Contender:
|
|
"""A single method's result in the tournament."""
|
|
|
|
method: str
|
|
score: float = 0.0
|
|
metrics: dict[str, Any] = field(default_factory=dict)
|
|
output_dir: str = ""
|
|
time_s: float = 0.0
|
|
error: str | None = None
|
|
round_eliminated: int = 0 # 0 = still alive / winner
|
|
direction_method: str = "" # which direction extraction was used
|
|
spectral_cert: str = "" # GREEN/YELLOW/RED/INCONCLUSIVE/""
|
|
|
|
|
|
@dataclass
|
|
class TourneyRound:
|
|
"""One round of the tournament."""
|
|
|
|
round_num: int
|
|
name: str
|
|
contenders: list[Contender] = field(default_factory=list)
|
|
prompt_volume: int = 0
|
|
advanced_to: list[str] = field(default_factory=list)
|
|
eliminated: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class TourneyResult:
|
|
"""Full tournament results."""
|
|
|
|
model: str
|
|
winner: Contender | None = None
|
|
rounds: list[TourneyRound] = field(default_factory=list)
|
|
total_time_s: float = 0.0
|
|
hub_repo: str | None = None
|
|
timestamp: str = ""
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"model": self.model,
|
|
"winner": {
|
|
"method": self.winner.method,
|
|
"score": self.winner.score,
|
|
"metrics": self.winner.metrics,
|
|
"time_s": self.winner.time_s,
|
|
} if self.winner else None,
|
|
"rounds": [
|
|
{
|
|
"round": r.round_num,
|
|
"name": r.name,
|
|
"prompt_volume": r.prompt_volume,
|
|
"contenders": [
|
|
{
|
|
"method": c.method,
|
|
"score": c.score,
|
|
"metrics": c.metrics,
|
|
"time_s": c.time_s,
|
|
"error": c.error,
|
|
"direction_method": c.direction_method,
|
|
"spectral_cert": c.spectral_cert,
|
|
}
|
|
for c in sorted(r.contenders, key=lambda x: x.score, reverse=True)
|
|
],
|
|
"advanced": r.advanced_to,
|
|
"eliminated": r.eliminated,
|
|
}
|
|
for r in self.rounds
|
|
],
|
|
"total_time_s": self.total_time_s,
|
|
"hub_repo": self.hub_repo,
|
|
"timestamp": self.timestamp,
|
|
}
|
|
|
|
|
|
CHECKPOINT_FILENAME = "tourney_checkpoint.json"
|
|
|
|
|
|
def _save_checkpoint(
|
|
output_dir: Path,
|
|
result: TourneyResult,
|
|
current_round_num: int,
|
|
current_round_name: str,
|
|
current_round_volume: int,
|
|
current_round_advance: int,
|
|
current_round_verify: int,
|
|
completed_methods: list[Contender],
|
|
remaining_methods: list[str],
|
|
alive: list[str],
|
|
model_name: str,
|
|
dataset_key: str,
|
|
quantization: str | None,
|
|
methods: list[str],
|
|
) -> Path:
|
|
"""Save tournament progress so it can be resumed after quota exhaustion."""
|
|
checkpoint = {
|
|
"version": 1,
|
|
"model": model_name,
|
|
"dataset_key": dataset_key,
|
|
"quantization": quantization,
|
|
"methods": methods,
|
|
"alive": alive,
|
|
"completed_rounds": [
|
|
{
|
|
"round_num": r.round_num,
|
|
"name": r.name,
|
|
"prompt_volume": r.prompt_volume,
|
|
"advanced_to": r.advanced_to,
|
|
"eliminated": r.eliminated,
|
|
"contenders": [
|
|
{
|
|
"method": c.method,
|
|
"score": c.score,
|
|
"metrics": c.metrics,
|
|
"output_dir": c.output_dir,
|
|
"time_s": c.time_s,
|
|
"error": c.error,
|
|
"round_eliminated": c.round_eliminated,
|
|
"direction_method": c.direction_method,
|
|
"spectral_cert": c.spectral_cert,
|
|
}
|
|
for c in r.contenders
|
|
],
|
|
}
|
|
for r in result.rounds
|
|
],
|
|
"interrupted_round": {
|
|
"round_num": current_round_num,
|
|
"name": current_round_name,
|
|
"prompt_volume": current_round_volume,
|
|
"advance_count": current_round_advance,
|
|
"verify_sample_size": current_round_verify,
|
|
"completed_methods": [
|
|
{
|
|
"method": c.method,
|
|
"score": c.score,
|
|
"metrics": c.metrics,
|
|
"output_dir": c.output_dir,
|
|
"time_s": c.time_s,
|
|
"error": c.error,
|
|
"round_eliminated": c.round_eliminated,
|
|
"direction_method": c.direction_method,
|
|
"spectral_cert": c.spectral_cert,
|
|
}
|
|
for c in completed_methods
|
|
],
|
|
"remaining_methods": remaining_methods,
|
|
},
|
|
"timestamp": datetime.now().isoformat(),
|
|
}
|
|
path = output_dir / CHECKPOINT_FILENAME
|
|
path.write_text(json.dumps(checkpoint, indent=2))
|
|
return path
|
|
|
|
|
|
def _load_checkpoint(output_dir: Path) -> dict | None:
|
|
"""Load a tournament checkpoint if one exists. Returns None if absent or corrupt."""
|
|
path = output_dir / CHECKPOINT_FILENAME
|
|
if not path.exists():
|
|
return None
|
|
try:
|
|
data = json.loads(path.read_text())
|
|
if data.get("version") != 1:
|
|
return None
|
|
return data
|
|
except (json.JSONDecodeError, KeyError):
|
|
return None
|
|
|
|
|
|
def _checkpoint_matches(
|
|
checkpoint: dict,
|
|
model_name: str,
|
|
dataset_key: str,
|
|
quantization: str | None,
|
|
) -> bool:
|
|
"""Check if a checkpoint is for the same model/dataset/quantization config."""
|
|
return (
|
|
checkpoint.get("model") == model_name
|
|
and checkpoint.get("dataset_key") == dataset_key
|
|
and checkpoint.get("quantization") == quantization
|
|
)
|
|
|
|
|
|
def _restore_rounds(checkpoint: dict) -> tuple[TourneyResult, list[Contender], list[str], dict]:
|
|
"""Restore completed rounds and interrupted round state from checkpoint.
|
|
|
|
Returns:
|
|
(result_with_completed_rounds, partial_contenders, remaining_methods, interrupted_round_spec)
|
|
"""
|
|
result = TourneyResult(
|
|
model=checkpoint["model"],
|
|
timestamp=checkpoint.get("timestamp", ""),
|
|
)
|
|
|
|
for rnd_data in checkpoint.get("completed_rounds", []):
|
|
rnd = TourneyRound(
|
|
round_num=rnd_data["round_num"],
|
|
name=rnd_data["name"],
|
|
prompt_volume=rnd_data.get("prompt_volume", 0),
|
|
advanced_to=rnd_data.get("advanced_to", []),
|
|
eliminated=rnd_data.get("eliminated", []),
|
|
)
|
|
for c_data in rnd_data.get("contenders", []):
|
|
rnd.contenders.append(Contender(
|
|
method=c_data["method"],
|
|
score=c_data.get("score", 0.0),
|
|
metrics=c_data.get("metrics", {}),
|
|
output_dir=c_data.get("output_dir", ""),
|
|
time_s=c_data.get("time_s", 0.0),
|
|
error=c_data.get("error"),
|
|
round_eliminated=c_data.get("round_eliminated", 0),
|
|
direction_method=c_data.get("direction_method", ""),
|
|
spectral_cert=c_data.get("spectral_cert", ""),
|
|
))
|
|
result.rounds.append(rnd)
|
|
|
|
ir = checkpoint.get("interrupted_round", {})
|
|
partial_contenders = []
|
|
for c_data in ir.get("completed_methods", []):
|
|
partial_contenders.append(Contender(
|
|
method=c_data["method"],
|
|
score=c_data.get("score", 0.0),
|
|
metrics=c_data.get("metrics", {}),
|
|
output_dir=c_data.get("output_dir", ""),
|
|
time_s=c_data.get("time_s", 0.0),
|
|
error=c_data.get("error"),
|
|
round_eliminated=c_data.get("round_eliminated", 0),
|
|
))
|
|
|
|
remaining = ir.get("remaining_methods", [])
|
|
|
|
return result, partial_contenders, remaining, ir
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Bracket renderer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def render_bracket(result: TourneyResult) -> str:
|
|
"""Render the tournament bracket as a markdown string."""
|
|
lines = []
|
|
lines.append(f"# OBLITERATUS TOURNEY — {result.model}")
|
|
lines.append("")
|
|
lines.append(f"**Winner: `{result.winner.method}`** "
|
|
f"(score: {result.winner.score:.4f})" if result.winner else "**No winner**")
|
|
lines.append(f"Total time: {result.total_time_s / 60:.1f} minutes")
|
|
if result.hub_repo:
|
|
lines.append(f"Pushed to: [{result.hub_repo}](https://huggingface.co/{result.hub_repo})")
|
|
lines.append("")
|
|
|
|
for rnd in result.rounds:
|
|
lines.append(f"## Round {rnd.round_num}: {rnd.name}")
|
|
lines.append(f"*{len(rnd.contenders)} contenders, {rnd.prompt_volume} prompt pairs*")
|
|
lines.append("")
|
|
lines.append("| Rank | Method | Dir | Score | Refusal | Coherence | KL Div | PPL | Cert | Time |")
|
|
lines.append("|------|--------|-----|-------|---------|-----------|--------|-----|------|------|")
|
|
|
|
sorted_contenders = sorted(rnd.contenders, key=lambda x: x.score, reverse=True)
|
|
for i, c in enumerate(sorted_contenders, 1):
|
|
if c.error:
|
|
lines.append(
|
|
f"| {i} | {c.method} | — | ERROR | — | — | — | — | — | {c.time_s:.0f}s |"
|
|
)
|
|
continue
|
|
m = c.metrics
|
|
# Only annotate elimination for non-final rounds
|
|
if c.method in rnd.advanced_to:
|
|
marker = ""
|
|
elif rnd.round_num < len(result.rounds):
|
|
marker = " *out*"
|
|
else:
|
|
marker = ""
|
|
rr = f"{m.get('refusal_rate', 0):.1%}" if m.get('refusal_rate') is not None else "—"
|
|
co = f"{m.get('coherence', 0):.3f}" if m.get('coherence') is not None else "—"
|
|
kl_val = m.get('kl_divergence')
|
|
kl_str = f"{kl_val:.4f}" if kl_val is not None else "—"
|
|
pp = f"{m.get('perplexity', 0):.1f}" if m.get('perplexity') is not None else "—"
|
|
dir_m = c.direction_method or m.get("direction_method", "—")
|
|
cert = c.spectral_cert or "—"
|
|
lines.append(
|
|
f"| {i} | **{c.method}**{marker} | {dir_m} | {c.score:.4f} "
|
|
f"| {rr} | {co} | {kl_str} | {pp} | {cert} | {c.time_s:.0f}s |"
|
|
)
|
|
lines.append("")
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
def render_bracket_html(result: TourneyResult) -> str:
|
|
"""Render the tournament bracket as a styled HTML bracket visualization."""
|
|
import html as html_mod
|
|
|
|
model_short = result.model.split("/")[-1] if "/" in result.model else result.model
|
|
|
|
# ── CSS ──────────────────────────────────────────────────────────────
|
|
css = """
|
|
<style>
|
|
.tourney-wrap {
|
|
font-family: 'Segoe UI', system-ui, -apple-system, sans-serif;
|
|
color: #e0e0e0;
|
|
max-width: 100%;
|
|
overflow-x: auto;
|
|
}
|
|
.tourney-header {
|
|
text-align: center;
|
|
padding: 18px 20px;
|
|
margin-bottom: 20px;
|
|
background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
|
|
border-radius: 12px;
|
|
border: 1px solid #333;
|
|
}
|
|
.tourney-header h2 {
|
|
margin: 0 0 4px 0;
|
|
font-size: 1.4em;
|
|
color: #fff;
|
|
letter-spacing: 1px;
|
|
}
|
|
.tourney-header .model-name {
|
|
font-size: 0.85em;
|
|
color: #8892b0;
|
|
font-family: 'Courier New', monospace;
|
|
}
|
|
.tourney-header .champion-box {
|
|
margin-top: 14px;
|
|
padding: 12px 18px;
|
|
background: linear-gradient(135deg, #2d1f00 0%, #3d2a00 100%);
|
|
border: 1px solid #f0c040;
|
|
border-radius: 8px;
|
|
display: inline-block;
|
|
}
|
|
.tourney-header .champion-box .trophy { font-size: 1.4em; }
|
|
.tourney-header .champion-box .champ-name {
|
|
font-size: 1.15em;
|
|
font-weight: 700;
|
|
color: #f0c040;
|
|
font-family: 'Courier New', monospace;
|
|
}
|
|
.tourney-header .champion-box .champ-score {
|
|
font-size: 0.85em;
|
|
color: #cca030;
|
|
margin-top: 2px;
|
|
}
|
|
.tourney-header .no-winner {
|
|
margin-top: 14px;
|
|
padding: 10px 16px;
|
|
background: #2a1a1a;
|
|
border: 1px solid #cc4444;
|
|
border-radius: 8px;
|
|
display: inline-block;
|
|
color: #ff6b6b;
|
|
font-weight: 600;
|
|
}
|
|
.tourney-header .time-info {
|
|
font-size: 0.78em;
|
|
color: #666;
|
|
margin-top: 8px;
|
|
}
|
|
|
|
/* ── Bracket flow ── */
|
|
.bracket-flow {
|
|
display: flex;
|
|
gap: 12px;
|
|
align-items: stretch;
|
|
padding: 4px 0;
|
|
}
|
|
.round-col {
|
|
flex: 1;
|
|
min-width: 200px;
|
|
max-width: 340px;
|
|
}
|
|
.round-title {
|
|
text-align: center;
|
|
font-size: 0.82em;
|
|
font-weight: 700;
|
|
text-transform: uppercase;
|
|
letter-spacing: 1.5px;
|
|
color: #8892b0;
|
|
padding: 6px 0 8px 0;
|
|
border-bottom: 2px solid #333;
|
|
margin-bottom: 8px;
|
|
}
|
|
.round-subtitle {
|
|
text-align: center;
|
|
font-size: 0.7em;
|
|
color: #555;
|
|
margin-top: 2px;
|
|
}
|
|
|
|
/* ── Method cards ── */
|
|
.method-card {
|
|
padding: 8px 10px;
|
|
margin: 4px 0;
|
|
border-radius: 6px;
|
|
border-left: 3px solid #444;
|
|
background: #1c1c2e;
|
|
transition: all 0.2s;
|
|
}
|
|
.method-card.advanced {
|
|
border-left-color: #4ecca3;
|
|
background: #1a2e28;
|
|
}
|
|
.method-card.champion {
|
|
border-left-color: #f0c040;
|
|
background: #2d2a1a;
|
|
box-shadow: 0 0 8px rgba(240, 192, 64, 0.15);
|
|
}
|
|
.method-card.eliminated {
|
|
border-left-color: #cc4444;
|
|
background: #1e1a1a;
|
|
opacity: 0.7;
|
|
}
|
|
.method-card.errored {
|
|
border-left-color: #ff4444;
|
|
background: #2a1a1a;
|
|
opacity: 0.6;
|
|
}
|
|
.card-top {
|
|
display: flex;
|
|
justify-content: space-between;
|
|
align-items: center;
|
|
}
|
|
.card-rank {
|
|
font-size: 0.7em;
|
|
color: #666;
|
|
font-weight: 700;
|
|
min-width: 18px;
|
|
}
|
|
.card-name {
|
|
font-weight: 600;
|
|
font-size: 0.88em;
|
|
font-family: 'Courier New', monospace;
|
|
flex: 1;
|
|
margin: 0 6px;
|
|
}
|
|
.card-score {
|
|
font-weight: 700;
|
|
font-size: 0.88em;
|
|
font-family: 'Courier New', monospace;
|
|
}
|
|
.card-score.good { color: #4ecca3; }
|
|
.card-score.mid { color: #f0c040; }
|
|
.card-score.bad { color: #cc4444; }
|
|
.card-metrics {
|
|
display: flex;
|
|
gap: 8px;
|
|
margin-top: 4px;
|
|
flex-wrap: wrap;
|
|
}
|
|
.metric {
|
|
font-size: 0.68em;
|
|
color: #777;
|
|
}
|
|
.metric .val {
|
|
color: #aaa;
|
|
font-family: 'Courier New', monospace;
|
|
}
|
|
.card-badge {
|
|
font-size: 0.65em;
|
|
font-weight: 700;
|
|
padding: 1px 5px;
|
|
border-radius: 3px;
|
|
text-transform: uppercase;
|
|
letter-spacing: 0.5px;
|
|
}
|
|
.badge-adv { background: #1a3a2e; color: #4ecca3; }
|
|
.badge-out { background: #2a1a1a; color: #cc6666; }
|
|
.badge-champ { background: #3d2a00; color: #f0c040; }
|
|
.badge-err { background: #2a1a1a; color: #ff6666; }
|
|
|
|
/* ── Arrow column ── */
|
|
.arrow-col {
|
|
display: flex;
|
|
flex-direction: column;
|
|
justify-content: center;
|
|
align-items: center;
|
|
width: 30px;
|
|
min-width: 30px;
|
|
color: #444;
|
|
font-size: 1.2em;
|
|
}
|
|
</style>
|
|
"""
|
|
|
|
# ── Header ───────────────────────────────────────────────────────────
|
|
header_parts = [
|
|
'<div class="tourney-header">',
|
|
'<h2>OBLITERATUS TOURNEY</h2>',
|
|
f'<div class="model-name">{html_mod.escape(model_short)}</div>',
|
|
]
|
|
|
|
if result.winner and not result.winner.error:
|
|
w = result.winner
|
|
m = w.metrics or {}
|
|
rr = f"{m.get('refusal_rate', 0):.1%}" if m.get("refusal_rate") is not None else "—"
|
|
co = f"{m.get('coherence', 0):.3f}" if m.get("coherence") is not None else "—"
|
|
header_parts.append('<div class="champion-box">')
|
|
header_parts.append(
|
|
f'<span class="trophy">🏆</span> '
|
|
f'<span class="champ-name">{html_mod.escape(w.method)}</span>'
|
|
)
|
|
dir_m = w.direction_method or "—"
|
|
cert = w.spectral_cert or "—"
|
|
header_parts.append(
|
|
f'<div class="champ-score">'
|
|
f'Score: {w.score:.4f} | Refusal: {rr} | '
|
|
f'Coherence: {co} | Dir: {html_mod.escape(dir_m)} | Cert: {html_mod.escape(cert)}'
|
|
f'</div>'
|
|
)
|
|
header_parts.append("</div>")
|
|
else:
|
|
header_parts.append('<div class="no-winner">No winner determined</div>')
|
|
|
|
if result.total_time_s:
|
|
header_parts.append(
|
|
f'<div class="time-info">{result.total_time_s / 60:.1f} min total</div>'
|
|
)
|
|
header_parts.append("</div>")
|
|
|
|
# ── Bracket columns ──────────────────────────────────────────────────
|
|
bracket_parts = ['<div class="bracket-flow">']
|
|
n_rounds = len(result.rounds)
|
|
|
|
for ri, rnd in enumerate(result.rounds):
|
|
if ri > 0:
|
|
bracket_parts.append('<div class="arrow-col">▶</div>')
|
|
|
|
bracket_parts.append('<div class="round-col">')
|
|
bracket_parts.append(
|
|
f'<div class="round-title">{html_mod.escape(rnd.name)}'
|
|
f'<div class="round-subtitle">{rnd.prompt_volume} pairs</div></div>'
|
|
)
|
|
|
|
sorted_c = sorted(rnd.contenders, key=lambda c: c.score, reverse=True)
|
|
is_final = ri == n_rounds - 1
|
|
|
|
for rank, c in enumerate(sorted_c, 1):
|
|
if c.error:
|
|
css_cls = "errored"
|
|
badge = '<span class="card-badge badge-err">ERR</span>'
|
|
elif is_final and rank == 1 and result.winner and not result.winner.error:
|
|
css_cls = "champion"
|
|
badge = '<span class="card-badge badge-champ">★ CHAMP</span>'
|
|
elif c.method in (rnd.advanced_to or []):
|
|
css_cls = "advanced"
|
|
badge = '<span class="card-badge badge-adv">ADV</span>'
|
|
else:
|
|
css_cls = "eliminated"
|
|
badge = '<span class="card-badge badge-out">OUT</span>'
|
|
|
|
# Score color
|
|
if c.error:
|
|
score_html = '<span class="card-score bad">ERR</span>'
|
|
elif c.score >= 0.7:
|
|
score_html = f'<span class="card-score good">{c.score:.4f}</span>'
|
|
elif c.score >= 0.4:
|
|
score_html = f'<span class="card-score mid">{c.score:.4f}</span>'
|
|
else:
|
|
score_html = f'<span class="card-score bad">{c.score:.4f}</span>'
|
|
|
|
# Compact metrics
|
|
m = c.metrics or {}
|
|
metric_spans = []
|
|
if not c.error:
|
|
dm = c.direction_method or m.get("direction_method", "")
|
|
if dm:
|
|
metric_spans.append(
|
|
f'<span class="metric">dir <span class="val">{html_mod.escape(dm)}</span></span>'
|
|
)
|
|
rr = m.get("refusal_rate")
|
|
if rr is not None:
|
|
metric_spans.append(
|
|
f'<span class="metric">ref <span class="val">{rr:.0%}</span></span>'
|
|
)
|
|
co = m.get("coherence")
|
|
if co is not None:
|
|
metric_spans.append(
|
|
f'<span class="metric">coh <span class="val">{co:.3f}</span></span>'
|
|
)
|
|
sc = c.spectral_cert or m.get("spectral_certification", "")
|
|
if sc:
|
|
cert_color = {"GREEN": "#4ecca3", "YELLOW": "#f0c040", "RED": "#cc4444"}.get(sc, "#777")
|
|
metric_spans.append(
|
|
f'<span class="metric">cert <span class="val" style="color:{cert_color}">{html_mod.escape(sc)}</span></span>'
|
|
)
|
|
kl = m.get("kl_divergence")
|
|
if kl is not None:
|
|
metric_spans.append(
|
|
f'<span class="metric">kl <span class="val">{kl:.4f}</span></span>'
|
|
)
|
|
pp = m.get("perplexity")
|
|
if pp is not None:
|
|
metric_spans.append(
|
|
f'<span class="metric">ppl <span class="val">{pp:.1f}</span></span>'
|
|
)
|
|
metrics_html = "".join(metric_spans)
|
|
|
|
bracket_parts.append(f'<div class="method-card {css_cls}">')
|
|
bracket_parts.append(
|
|
f'<div class="card-top">'
|
|
f'<span class="card-rank">#{rank}</span>'
|
|
f'<span class="card-name">{html_mod.escape(c.method)}</span>'
|
|
f'{score_html}'
|
|
f'{badge}'
|
|
f'</div>'
|
|
)
|
|
if metrics_html:
|
|
bracket_parts.append(f'<div class="card-metrics">{metrics_html}</div>')
|
|
bracket_parts.append("</div>")
|
|
|
|
bracket_parts.append("</div>")
|
|
|
|
bracket_parts.append("</div>")
|
|
|
|
return css + '<div class="tourney-wrap">' + "\n".join(header_parts + bracket_parts) + "</div>"
|
|
|
|
|
|
def generate_model_card(result: TourneyResult) -> str:
|
|
"""Generate a HuggingFace model card for the tournament winner."""
|
|
w = result.winner
|
|
if not w:
|
|
return ""
|
|
|
|
short_model = result.model.split("/")[-1] if "/" in result.model else result.model
|
|
bracket = render_bracket(result)
|
|
|
|
return f"""---
|
|
language: en
|
|
tags:
|
|
- obliteratus
|
|
- abliteration
|
|
- uncensored
|
|
- tourney
|
|
base_model: {result.model}
|
|
---
|
|
|
|
# {short_model} — Obliterated (Tourney Winner)
|
|
|
|
This model was abliterated using the **`{w.method}`** method, selected by an
|
|
automated [OBLITERATUS](https://github.com/elder-plinius/OBLITERATUS) tournament
|
|
that pitted **{len(TOURNEY_METHODS)} abliteration techniques** against each other
|
|
in elimination rounds.
|
|
|
|
## Winning Method: `{w.method}`
|
|
|
|
| Metric | Value |
|
|
|--------|-------|
|
|
| Composite Score | **{w.score:.4f}** |
|
|
| Direction Method | {w.direction_method or 'N/A'} |
|
|
| Refusal Rate | {f'{w.metrics["refusal_rate"]:.1%}' if w.metrics.get('refusal_rate') is not None else 'N/A'} |
|
|
| Coherence | {f'{w.metrics["coherence"]:.3f}' if w.metrics.get('coherence') is not None else 'N/A'} |
|
|
| KL Divergence | {f'{w.metrics["kl_divergence"]:.4f}' if w.metrics.get('kl_divergence') is not None else 'N/A'} |
|
|
| Perplexity | {f'{w.metrics["perplexity"]:.1f}' if w.metrics.get('perplexity') is not None else 'N/A'} |
|
|
| Spectral Cert | {w.spectral_cert or 'N/A'} |
|
|
|
|
## How to Use
|
|
|
|
```python
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
model = AutoModelForCausalLM.from_pretrained("{result.hub_repo or 'this-repo'}")
|
|
tokenizer = AutoTokenizer.from_pretrained("{result.hub_repo or 'this-repo'}")
|
|
```
|
|
|
|
## Full Tournament Bracket
|
|
|
|
{bracket}
|
|
|
|
---
|
|
|
|
*Generated by [OBLITERATUS](https://github.com/elder-plinius/OBLITERATUS) tourney on {result.timestamp}*
|
|
"""
|
|
|
|
|
|
def _noop_log(msg: str) -> None:
|
|
"""Picklable no-op log callback (lambdas can't be pickled by ZeroGPU)."""
|
|
pass
|
|
|
|
|
|
def _noop_round(r: TourneyRound) -> None:
|
|
"""Picklable no-op round callback."""
|
|
pass
|
|
|
|
|
|
class _MethodLogger:
|
|
"""Picklable per-method log adapter that prefixes messages.
|
|
|
|
ZeroGPU pickles bound methods (and their ``self``) when shipping work to
|
|
the GPU worker process. Plain lambdas like
|
|
``lambda msg: self.log(f" [{method}] {msg}")`` can't survive that, so
|
|
this small class replaces them.
|
|
"""
|
|
|
|
def __init__(self, parent_log: Callable[[str], None], method: str):
|
|
self._parent = parent_log
|
|
self._method = method
|
|
|
|
def __call__(self, msg: str):
|
|
self._parent(f" [{self._method}] {msg}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core runner
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TourneyRunner:
|
|
"""Run a March Madness-style tournament across all abliteration methods.
|
|
|
|
Round 1 — Qualifiers: All methods, reduced prompts. Bottom half eliminated.
|
|
Round 2 — Semifinals: Survivors, full prompts. Bottom half eliminated.
|
|
Round 3 — Finals: Top contenders, maximum prompts + extended verify.
|
|
Winner — Auto-pushed to HuggingFace Hub (if hub_org is set).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_name: str,
|
|
hub_org: str | None = None,
|
|
hub_repo: str | None = None,
|
|
device: str = "auto",
|
|
dtype: str = "float16",
|
|
dataset_key: str = "builtin",
|
|
quantization: str | None = None,
|
|
methods: list[str] | None = None,
|
|
output_dir: str = "/tmp/obliteratus_tourney",
|
|
on_log: Callable[[str], None] | None = None,
|
|
on_round: Callable[[TourneyRound], None] | None = None,
|
|
resume: bool = False,
|
|
):
|
|
self.model_name = model_name
|
|
self.hub_org = hub_org
|
|
self.hub_repo = hub_repo
|
|
self.device = device
|
|
self.dtype = dtype
|
|
self.dataset_key = dataset_key
|
|
self.quantization = quantization
|
|
self.methods = methods or list(TOURNEY_METHODS)
|
|
self.output_dir = Path(output_dir)
|
|
self.resume = resume
|
|
|
|
# When resuming, preserve the output directory (it contains
|
|
# checkpoints and model saves from completed methods).
|
|
if not resume:
|
|
if self.output_dir.exists():
|
|
shutil.rmtree(self.output_dir, ignore_errors=True)
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
self._on_log = on_log or _noop_log
|
|
self._on_round = on_round or _noop_round
|
|
|
|
def log(self, msg: str):
|
|
self._on_log(msg)
|
|
|
|
def _load_prompts(self, volume: int) -> tuple[list[str], list[str]]:
|
|
from obliteratus.prompts import load_dataset_source
|
|
harmful, harmless = load_dataset_source(self.dataset_key)
|
|
n = min(volume, len(harmful), len(harmless))
|
|
return harmful[:n], harmless[:n]
|
|
|
|
def _run_method(
|
|
self,
|
|
method: str,
|
|
harmful: list[str],
|
|
harmless: list[str],
|
|
save_dir: str,
|
|
verify_sample_size: int = 30,
|
|
) -> Contender:
|
|
"""Run a single abliteration method and return its Contender result."""
|
|
t0 = time.time()
|
|
contender = Contender(method=method)
|
|
|
|
try:
|
|
# Use informed pipeline for 'informed' method
|
|
method_log = _MethodLogger(self._on_log, method)
|
|
|
|
if method == "informed":
|
|
from obliteratus.informed_pipeline import InformedAbliterationPipeline
|
|
pipeline = InformedAbliterationPipeline(
|
|
model_name=self.model_name,
|
|
output_dir=save_dir,
|
|
device=self.device,
|
|
dtype=self.dtype,
|
|
quantization=self.quantization,
|
|
trust_remote_code=True,
|
|
harmful_prompts=harmful,
|
|
harmless_prompts=harmless,
|
|
on_log=method_log,
|
|
)
|
|
pipeline.run_informed()
|
|
else:
|
|
from obliteratus.abliterate import AbliterationPipeline
|
|
pipeline = AbliterationPipeline(
|
|
model_name=self.model_name,
|
|
output_dir=save_dir,
|
|
device=self.device,
|
|
dtype=self.dtype,
|
|
method=method,
|
|
quantization=self.quantization,
|
|
trust_remote_code=True,
|
|
harmful_prompts=harmful,
|
|
harmless_prompts=harmless,
|
|
verify_sample_size=verify_sample_size,
|
|
on_log=method_log,
|
|
)
|
|
pipeline.run()
|
|
|
|
contender.metrics = dict(pipeline._quality_metrics)
|
|
contender.score = composite_score(contender.metrics)
|
|
contender.output_dir = save_dir
|
|
contender.direction_method = getattr(pipeline, "direction_method", "")
|
|
contender.spectral_cert = contender.metrics.get("spectral_certification", "") or ""
|
|
|
|
# Free pipeline to reclaim GPU
|
|
del pipeline
|
|
except Exception as e:
|
|
# Re-raise GPU quota / expired-token errors so the
|
|
# tournament aborts immediately rather than letting every
|
|
# remaining method fail for the same reason.
|
|
if self._is_quota_error(e):
|
|
raise
|
|
import traceback
|
|
contender.error = f"{type(e).__name__}: {e}"
|
|
contender.score = -1.0 # errors sort to bottom
|
|
self.log(f" [{method}] ERROR: {contender.error}")
|
|
self.log(f" [{method}] TRACEBACK:\n{traceback.format_exc()}")
|
|
finally:
|
|
# Always clean up GPU between methods — including when
|
|
# re-raising quota errors, to avoid leaking the pipeline.
|
|
gc.collect()
|
|
try:
|
|
from obliteratus import device as dev
|
|
dev.empty_cache()
|
|
except Exception:
|
|
pass
|
|
|
|
contender.time_s = time.time() - t0
|
|
return contender
|
|
|
|
def _run_round(
|
|
self,
|
|
round_num: int,
|
|
name: str,
|
|
methods: list[str],
|
|
prompt_volume: int,
|
|
advance_count: int,
|
|
verify_sample_size: int = 30,
|
|
) -> TourneyRound:
|
|
"""Execute one round of the tournament."""
|
|
self.log("")
|
|
self.log("=" * 60)
|
|
self.log(f"ROUND {round_num}: {name}")
|
|
self.log(f" {len(methods)} contenders | {prompt_volume} prompt pairs | "
|
|
f"top {advance_count} advance")
|
|
self.log("=" * 60)
|
|
|
|
harmful, harmless = self._load_prompts(prompt_volume)
|
|
|
|
rnd = TourneyRound(
|
|
round_num=round_num,
|
|
name=name,
|
|
prompt_volume=prompt_volume,
|
|
)
|
|
|
|
for i, method in enumerate(methods, 1):
|
|
self.log(f"\n[{i}/{len(methods)}] Running: {method}")
|
|
save_dir = str(self.output_dir / f"r{round_num}_{method}")
|
|
contender = self._run_method(
|
|
method, harmful, harmless, save_dir, verify_sample_size,
|
|
)
|
|
rnd.contenders.append(contender)
|
|
self.log(
|
|
f" {method}: score={contender.score:.4f} "
|
|
f"(refusal={contender.metrics.get('refusal_rate', '?')}, "
|
|
f"coherence={contender.metrics.get('coherence', '?')}) "
|
|
f"[{contender.time_s:.0f}s]"
|
|
)
|
|
|
|
# Free checkpoint for non-finalists as we go (save disk)
|
|
# We'll keep them until we know who advances
|
|
|
|
# Rank by score
|
|
ranked = sorted(rnd.contenders, key=lambda c: c.score, reverse=True)
|
|
rnd.advanced_to = [c.method for c in ranked[:advance_count]]
|
|
rnd.eliminated = [c.method for c in ranked[advance_count:]]
|
|
|
|
# Mark eliminated
|
|
for c in ranked[advance_count:]:
|
|
c.round_eliminated = round_num
|
|
|
|
self.log(f"\n{'─' * 40}")
|
|
self.log(f"Round {round_num} results:")
|
|
for i, c in enumerate(ranked, 1):
|
|
status = "ADVANCE" if c.method in rnd.advanced_to else "OUT"
|
|
self.log(f" {i}. {c.method}: {c.score:.4f} [{status}]")
|
|
|
|
# Clean up eliminated checkpoints to free disk
|
|
for c in ranked[advance_count:]:
|
|
if c.output_dir and Path(c.output_dir).exists():
|
|
shutil.rmtree(c.output_dir, ignore_errors=True)
|
|
|
|
self._on_round(rnd)
|
|
return rnd
|
|
|
|
def run(self) -> TourneyResult:
|
|
"""Execute the full tournament. Returns TourneyResult with winner."""
|
|
t_start = time.time()
|
|
result = TourneyResult(
|
|
model=self.model_name,
|
|
timestamp=datetime.now().isoformat(),
|
|
)
|
|
|
|
n_methods = len(self.methods)
|
|
self.log("OBLITERATUS TOURNEY")
|
|
self.log(f"Model: {self.model_name}")
|
|
self.log(f"Contenders: {n_methods} methods")
|
|
self.log(f"Dataset: {self.dataset_key}")
|
|
|
|
# Pre-flight disk space check
|
|
try:
|
|
disk = shutil.disk_usage(self.output_dir)
|
|
free_gb = disk.free / 1e9
|
|
self.log(f"Disk space: {free_gb:.1f} GB free on {self.output_dir}")
|
|
if free_gb < 5.0:
|
|
self.log(
|
|
f"WARNING: Low disk space ({free_gb:.1f} GB free). "
|
|
f"Tournament may fail saving checkpoints."
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# ── Round 1: Qualifiers — all methods, reduced prompts ────────
|
|
r1_advance = max(2, math.ceil(n_methods / 2))
|
|
r1 = self._run_round(
|
|
round_num=1,
|
|
name="Qualifiers",
|
|
methods=self.methods,
|
|
prompt_volume=64, # fast qualifier round
|
|
advance_count=r1_advance,
|
|
verify_sample_size=20,
|
|
)
|
|
result.rounds.append(r1)
|
|
alive = list(r1.advanced_to)
|
|
|
|
if len(alive) <= 1:
|
|
# Only 1 survivor — they win
|
|
pass
|
|
else:
|
|
# ── Round 2: Semifinals — survivors, full prompts ─────────
|
|
r2_advance = max(2, math.ceil(len(alive) / 2))
|
|
r2 = self._run_round(
|
|
round_num=2,
|
|
name="Semifinals",
|
|
methods=alive,
|
|
prompt_volume=128,
|
|
advance_count=r2_advance,
|
|
verify_sample_size=30,
|
|
)
|
|
result.rounds.append(r2)
|
|
alive = list(r2.advanced_to)
|
|
|
|
if len(alive) > 2:
|
|
# ── Round 3: Finals — top contenders, max prompts ─────
|
|
r3 = self._run_round(
|
|
round_num=3,
|
|
name="Finals",
|
|
methods=alive,
|
|
prompt_volume=256,
|
|
advance_count=1,
|
|
verify_sample_size=50,
|
|
)
|
|
result.rounds.append(r3)
|
|
alive = list(r3.advanced_to)
|
|
elif len(alive) == 2:
|
|
# Head-to-head final
|
|
r3 = self._run_round(
|
|
round_num=3,
|
|
name="Championship",
|
|
methods=alive,
|
|
prompt_volume=256,
|
|
advance_count=1,
|
|
verify_sample_size=50,
|
|
)
|
|
result.rounds.append(r3)
|
|
alive = list(r3.advanced_to)
|
|
|
|
# ── Determine winner ──────────────────────────────────────────
|
|
last_round = result.rounds[-1]
|
|
ranked = sorted(last_round.contenders, key=lambda c: c.score, reverse=True)
|
|
# Only crown a winner if they completed without error
|
|
winner = ranked[0] if ranked and not ranked[0].error else None
|
|
result.winner = winner
|
|
result.total_time_s = time.time() - t_start
|
|
|
|
# Clean up non-winner finalist dirs to free disk
|
|
for c in ranked[1:]:
|
|
if c.output_dir and Path(c.output_dir).exists():
|
|
shutil.rmtree(c.output_dir, ignore_errors=True)
|
|
|
|
self.log("")
|
|
self.log("=" * 60)
|
|
if winner:
|
|
self.log(f"CHAMPION: {winner.method} (score: {winner.score:.4f})")
|
|
else:
|
|
n_errors = sum(1 for c in ranked if c.error)
|
|
self.log(f"NO WINNER — {n_errors}/{len(ranked)} methods errored")
|
|
self.log(f"Total tournament time: {result.total_time_s / 60:.1f} minutes")
|
|
self.log("=" * 60)
|
|
|
|
# ── Save tournament results ───────────────────────────────────
|
|
results_path = self.output_dir / "tourney_results.json"
|
|
results_path.write_text(json.dumps(result.to_dict(), indent=2))
|
|
self.log(f"Results saved to {results_path}")
|
|
|
|
bracket_path = self.output_dir / "tourney_bracket.md"
|
|
bracket_path.write_text(render_bracket(result))
|
|
self.log(f"Bracket saved to {bracket_path}")
|
|
|
|
# ── Push winner to HuggingFace Hub ────────────────────────────
|
|
if winner and winner.output_dir and (self.hub_org or self.hub_repo):
|
|
self._push_winner(result)
|
|
|
|
return result
|
|
|
|
@staticmethod
|
|
def _is_quota_error(exc: BaseException) -> bool:
|
|
msg = str(exc).lower()
|
|
if "exceeded" in msg and "gpu quota" in msg:
|
|
return True
|
|
if "expired" in msg and "zerogpu" in msg:
|
|
return True
|
|
return False
|
|
|
|
def _run_one_method(self, method, harmful, harmless, save_dir, verify_sz, gpu_wrapper):
|
|
"""Run a single method, optionally inside a gpu_wrapper."""
|
|
if gpu_wrapper is not None:
|
|
return gpu_wrapper(
|
|
self._run_method, method, harmful, harmless,
|
|
save_dir, verify_sz,
|
|
)
|
|
return self._run_method(
|
|
method, harmful, harmless, save_dir, verify_sz,
|
|
)
|
|
|
|
def run_iter(self, gpu_wrapper=None):
|
|
"""Generator version of run() — yields (status, result_so_far) after each method.
|
|
|
|
Supports automatic resume: if ``self.resume`` is True and a valid
|
|
checkpoint exists from a previous quota-interrupted run with the
|
|
same model/dataset/quantization, completed rounds and methods are
|
|
restored and execution continues from the interruption point.
|
|
|
|
When a GPU quota error occurs, a checkpoint is saved to disk and
|
|
the exception is re-raised. The caller can catch it and inform
|
|
the user that clicking **Run** again will resume automatically.
|
|
|
|
Args:
|
|
gpu_wrapper: Optional callable ``gpu_wrapper(fn, *args, **kw)``
|
|
that executes *fn* inside a GPU context. On ZeroGPU Spaces
|
|
this should be a ``@spaces.GPU``-decorated function so each
|
|
method gets its own GPU allocation (up to 5 min each).
|
|
|
|
Yields:
|
|
(status_msg: str, result: TourneyResult | None)
|
|
"""
|
|
|
|
t_start = time.time()
|
|
resuming = False
|
|
checkpoint = None
|
|
partial_contenders: list[Contender] = []
|
|
resume_remaining: list[str] = []
|
|
resume_round_spec: dict = {}
|
|
|
|
# ── Try to resume from checkpoint ────────────────────────────
|
|
if self.resume:
|
|
checkpoint = _load_checkpoint(self.output_dir)
|
|
if checkpoint and _checkpoint_matches(
|
|
checkpoint, self.model_name, self.dataset_key, self.quantization
|
|
):
|
|
resuming = True
|
|
result, partial_contenders, resume_remaining, resume_round_spec = (
|
|
_restore_rounds(checkpoint)
|
|
)
|
|
n_completed_rounds = len(result.rounds)
|
|
n_completed_methods = len(partial_contenders)
|
|
self.log("OBLITERATUS TOURNEY — RESUMING")
|
|
self.log(f"Restored {n_completed_rounds} completed round(s), "
|
|
f"{n_completed_methods} method(s) from interrupted round")
|
|
yield (
|
|
f"**Resuming tournament** — {n_completed_rounds} round(s) "
|
|
f"and {n_completed_methods} method(s) restored from checkpoint.",
|
|
result,
|
|
)
|
|
|
|
# Determine alive list from checkpoint
|
|
alive = list(checkpoint.get("alive", self.methods))
|
|
|
|
# Remove the checkpoint file now that we've loaded it
|
|
ckpt_path = self.output_dir / CHECKPOINT_FILENAME
|
|
if ckpt_path.exists():
|
|
ckpt_path.unlink()
|
|
else:
|
|
# Checkpoint doesn't match current config — start fresh
|
|
checkpoint = None
|
|
|
|
n_methods = len(self.methods)
|
|
|
|
if not resuming:
|
|
result = TourneyResult(
|
|
model=self.model_name,
|
|
timestamp=datetime.now().isoformat(),
|
|
)
|
|
alive = list(self.methods)
|
|
|
|
self.log("OBLITERATUS TOURNEY")
|
|
self.log(f"Model: {self.model_name}")
|
|
self.log(f"Contenders: {n_methods} methods")
|
|
self.log(f"Dataset: {self.dataset_key}")
|
|
|
|
# Pre-flight disk space check
|
|
try:
|
|
disk = shutil.disk_usage(self.output_dir)
|
|
free_gb = disk.free / 1e9
|
|
self.log(f"Disk space: {free_gb:.1f} GB free on {self.output_dir}")
|
|
if free_gb < 5.0:
|
|
msg = (
|
|
f"Low disk space: only {free_gb:.1f} GB free. "
|
|
f"Tournament needs space for multiple model checkpoints. "
|
|
f"Free up space or use quantization to reduce checkpoint sizes."
|
|
)
|
|
self.log(f"WARNING: {msg}")
|
|
yield (f"**Warning:** {msg}", result)
|
|
except Exception:
|
|
pass
|
|
|
|
# Build round schedule
|
|
rounds_schedule: list[tuple] = []
|
|
|
|
if resuming and resume_round_spec:
|
|
# We have an interrupted round to finish — schedule it first,
|
|
# then let the dynamic scheduling add subsequent rounds.
|
|
skip_completed_rounds = len(result.rounds)
|
|
else:
|
|
skip_completed_rounds = 0
|
|
|
|
# Always build the full schedule starting from round 1.
|
|
# Completed rounds will be skipped below.
|
|
r1_advance = max(2, math.ceil(n_methods / 2))
|
|
rounds_schedule.append((1, "Qualifiers", self.methods, 64, r1_advance, 20))
|
|
|
|
for round_spec in rounds_schedule:
|
|
round_num, name, methods, volume, advance_count, verify_sz = round_spec
|
|
|
|
# Skip rounds that were already completed in the checkpoint
|
|
if resuming and round_num <= skip_completed_rounds:
|
|
# Re-derive alive and schedule next rounds from completed data
|
|
completed_rnd = result.rounds[round_num - 1]
|
|
alive = list(completed_rnd.advanced_to)
|
|
if round_num == 1 and len(alive) > 1:
|
|
r2_advance = max(2, math.ceil(len(alive) / 2))
|
|
rounds_schedule.append((2, "Semifinals", alive, 128, r2_advance, 30))
|
|
elif round_num == 2 and len(alive) > 1:
|
|
r3_name = "Championship" if len(alive) == 2 else "Finals"
|
|
rounds_schedule.append((3, r3_name, alive, 256, 1, 50))
|
|
self.log(f"\nSkipping completed Round {round_num}: {name}")
|
|
yield (
|
|
f"**Round {round_num} ({name}):** already completed (restored from checkpoint)",
|
|
result,
|
|
)
|
|
continue
|
|
|
|
# For the interrupted round, merge checkpoint data
|
|
is_interrupted_round = (
|
|
resuming
|
|
and resume_round_spec
|
|
and round_num == resume_round_spec.get("round_num")
|
|
)
|
|
|
|
if is_interrupted_round:
|
|
# Use the interrupted round's parameters
|
|
volume = resume_round_spec.get("prompt_volume", volume)
|
|
advance_count = resume_round_spec.get("advance_count", advance_count)
|
|
verify_sz = resume_round_spec.get("verify_sample_size", verify_sz)
|
|
methods = list(
|
|
[c.method for c in partial_contenders] + resume_remaining
|
|
)
|
|
|
|
self.log("")
|
|
self.log("=" * 60)
|
|
self.log(f"ROUND {round_num}: {name}")
|
|
self.log(f" {len(methods)} contenders | {volume} prompt pairs | "
|
|
f"top {advance_count} advance")
|
|
self.log("=" * 60)
|
|
|
|
harmful, harmless = self._load_prompts(volume)
|
|
|
|
rnd = TourneyRound(
|
|
round_num=round_num,
|
|
name=name,
|
|
prompt_volume=volume,
|
|
)
|
|
|
|
# If resuming an interrupted round, restore already-completed
|
|
# contenders and only run the remaining methods.
|
|
methods_to_run = list(methods)
|
|
if is_interrupted_round and partial_contenders:
|
|
for c in partial_contenders:
|
|
rnd.contenders.append(c)
|
|
self.log(f" [restored] {c.method}: score={c.score:.4f}")
|
|
methods_to_run = list(resume_remaining)
|
|
self.log(f" {len(partial_contenders)} method(s) restored, "
|
|
f"{len(methods_to_run)} remaining")
|
|
|
|
total_in_round = len(rnd.contenders) + len(methods_to_run)
|
|
|
|
for i, method in enumerate(methods_to_run, len(rnd.contenders) + 1):
|
|
self.log(f"\n[{i}/{total_in_round}] Running: {method}")
|
|
yield (
|
|
f"**Round {round_num} ({name}):** running `{method}` [{i}/{total_in_round}]",
|
|
result,
|
|
)
|
|
|
|
save_dir = str(self.output_dir / f"r{round_num}_{method}")
|
|
|
|
try:
|
|
contender = self._run_one_method(
|
|
method, harmful, harmless, save_dir, verify_sz,
|
|
gpu_wrapper,
|
|
)
|
|
except Exception as exc:
|
|
if self._is_quota_error(exc):
|
|
# Save checkpoint so the tournament can resume later.
|
|
# Include the failed method in remaining so it retries.
|
|
still_remaining = methods_to_run[methods_to_run.index(method):]
|
|
_save_checkpoint(
|
|
output_dir=self.output_dir,
|
|
result=result,
|
|
current_round_num=round_num,
|
|
current_round_name=name,
|
|
current_round_volume=volume,
|
|
current_round_advance=advance_count,
|
|
current_round_verify=verify_sz,
|
|
completed_methods=list(rnd.contenders),
|
|
remaining_methods=still_remaining,
|
|
alive=alive,
|
|
model_name=self.model_name,
|
|
dataset_key=self.dataset_key,
|
|
quantization=self.quantization,
|
|
methods=self.methods,
|
|
)
|
|
self.log("\nGPU SESSION INTERRUPTED — checkpoint saved")
|
|
self.log(f" Reason: {exc}")
|
|
self.log(f" Completed: {len(rnd.contenders)} methods in round {round_num}")
|
|
self.log(f" Remaining: {len(still_remaining)} methods")
|
|
self.log(" Click Run again to resume automatically.")
|
|
raise
|
|
|
|
rnd.contenders.append(contender)
|
|
self.log(
|
|
f" {method}: score={contender.score:.4f} "
|
|
f"(refusal={contender.metrics.get('refusal_rate', '?')}, "
|
|
f"coherence={contender.metrics.get('coherence', '?')}) "
|
|
f"[{contender.time_s:.0f}s]"
|
|
)
|
|
|
|
# Rank, advance, eliminate
|
|
ranked = sorted(rnd.contenders, key=lambda c: c.score, reverse=True)
|
|
rnd.advanced_to = [c.method for c in ranked[:advance_count]]
|
|
rnd.eliminated = [c.method for c in ranked[advance_count:]]
|
|
for c in ranked[advance_count:]:
|
|
c.round_eliminated = round_num
|
|
|
|
self.log(f"\n{'─' * 40}")
|
|
self.log(f"Round {round_num} results:")
|
|
for idx, c in enumerate(ranked, 1):
|
|
status = "ADVANCE" if c.method in rnd.advanced_to else "OUT"
|
|
self.log(f" {idx}. {c.method}: {c.score:.4f} [{status}]")
|
|
|
|
# Clean up eliminated checkpoints
|
|
for c in ranked[advance_count:]:
|
|
if c.output_dir and Path(c.output_dir).exists():
|
|
shutil.rmtree(c.output_dir, ignore_errors=True)
|
|
|
|
self._on_round(rnd)
|
|
result.rounds.append(rnd)
|
|
alive = list(rnd.advanced_to)
|
|
|
|
# Schedule next round dynamically
|
|
if round_num == 1 and len(alive) > 1:
|
|
r2_advance = max(2, math.ceil(len(alive) / 2))
|
|
rounds_schedule.append((2, "Semifinals", alive, 128, r2_advance, 30))
|
|
elif round_num == 2 and len(alive) > 1:
|
|
r3_name = "Championship" if len(alive) == 2 else "Finals"
|
|
rounds_schedule.append((3, r3_name, alive, 256, 1, 50))
|
|
|
|
# ── Determine winner ──────────────────────────────────────────
|
|
last_round = result.rounds[-1]
|
|
ranked = sorted(last_round.contenders, key=lambda c: c.score, reverse=True)
|
|
winner = ranked[0] if ranked and not ranked[0].error else None
|
|
result.winner = winner
|
|
result.total_time_s = time.time() - t_start
|
|
|
|
# Clean up non-winner finalist dirs to free disk
|
|
for c in ranked[1:]:
|
|
if c.output_dir and Path(c.output_dir).exists():
|
|
shutil.rmtree(c.output_dir, ignore_errors=True)
|
|
|
|
self.log("")
|
|
self.log("=" * 60)
|
|
if winner:
|
|
self.log(f"CHAMPION: {winner.method} (score: {winner.score:.4f})")
|
|
else:
|
|
n_errors = sum(1 for c in ranked if c.error)
|
|
self.log(f"NO WINNER — {n_errors}/{len(ranked)} methods errored")
|
|
self.log(f"Total tournament time: {result.total_time_s / 60:.1f} minutes")
|
|
self.log("=" * 60)
|
|
|
|
# Save results
|
|
results_path = self.output_dir / "tourney_results.json"
|
|
results_path.write_text(json.dumps(result.to_dict(), indent=2))
|
|
self.log(f"Results saved to {results_path}")
|
|
|
|
bracket_path = self.output_dir / "tourney_bracket.md"
|
|
bracket_path.write_text(render_bracket(result))
|
|
self.log(f"Bracket saved to {bracket_path}")
|
|
|
|
# Clean up checkpoint file on successful completion
|
|
ckpt_path = self.output_dir / CHECKPOINT_FILENAME
|
|
if ckpt_path.exists():
|
|
ckpt_path.unlink()
|
|
|
|
# Push winner
|
|
if winner and winner.output_dir and (self.hub_org or self.hub_repo):
|
|
self._push_winner(result)
|
|
|
|
# Final yield with completed result
|
|
yield ("Tournament complete", result)
|
|
|
|
def _push_winner(self, result: TourneyResult):
|
|
"""Push the winning model to HuggingFace Hub."""
|
|
winner = result.winner
|
|
if not winner or not winner.output_dir:
|
|
return
|
|
|
|
try:
|
|
from huggingface_hub import HfApi
|
|
|
|
short_model = self.model_name.split("/")[-1] if "/" in self.model_name else self.model_name
|
|
if self.hub_repo:
|
|
repo_id = self.hub_repo
|
|
else:
|
|
repo_id = f"{self.hub_org}/{short_model}-OBLITERATED"
|
|
result.hub_repo = repo_id
|
|
|
|
self.log(f"\nPushing winner to Hub: {repo_id}")
|
|
|
|
_token = os.environ.get("HF_PUSH_TOKEN") or os.environ.get("HF_TOKEN") or None
|
|
api = HfApi(token=_token) if _token else HfApi()
|
|
api.create_repo(repo_id, exist_ok=True)
|
|
|
|
# Write model card
|
|
model_card = generate_model_card(result)
|
|
card_path = Path(winner.output_dir) / "README.md"
|
|
card_path.write_text(model_card)
|
|
|
|
# Write tourney results alongside model
|
|
results_dest = Path(winner.output_dir) / "tourney_results.json"
|
|
results_dest.write_text(json.dumps(result.to_dict(), indent=2))
|
|
|
|
api.upload_folder(
|
|
folder_path=winner.output_dir,
|
|
repo_id=repo_id,
|
|
commit_message=(
|
|
f"OBLITERATUS tourney: {winner.method} wins "
|
|
f"(score {winner.score:.4f}) on {self.model_name}"
|
|
),
|
|
)
|
|
self.log(f"Pushed to https://huggingface.co/{repo_id}")
|
|
|
|
except Exception as e:
|
|
self.log(f"Hub push failed: {e}")
|