"""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 tempfile import time from dataclasses import dataclass, field from datetime import datetime from pathlib import Path from typing import Any, Callable from obliteratus.credential_sources import resolve_first from obliteratus.tourney_contracts import parse_checkpoint_document # --------------------------------------------------------------------------- # 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: return parse_checkpoint_document(path.read_text()) except (OSError, UnicodeError): 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), direction_method=c_data.get("direction_method", ""), spectral_cert=c_data.get("spectral_cert", ""), )) 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 = """ """ # ── Header ─────────────────────────────────────────────────────────── header_parts = [ '
', '

OBLITERATUS TOURNEY

', f'
{html_mod.escape(model_short)}
', ] 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('
') header_parts.append( f'🏆 ' f'{html_mod.escape(w.method)}' ) dir_m = w.direction_method or "—" cert = w.spectral_cert or "—" header_parts.append( f'
' f'Score: {w.score:.4f}  |  Refusal: {rr}  |  ' f'Coherence: {co}  |  Dir: {html_mod.escape(dir_m)}  |  Cert: {html_mod.escape(cert)}' f'
' ) header_parts.append("
") else: header_parts.append('
No winner determined
') if result.total_time_s: header_parts.append( f'
{result.total_time_s / 60:.1f} min total
' ) header_parts.append("
") # ── Bracket columns ────────────────────────────────────────────────── bracket_parts = ['
'] n_rounds = len(result.rounds) for ri, rnd in enumerate(result.rounds): if ri > 0: bracket_parts.append('
▶
') bracket_parts.append('
') bracket_parts.append( f'
{html_mod.escape(rnd.name)}' f'
{rnd.prompt_volume} pairs
' ) 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 = 'ERR' elif is_final and rank == 1 and result.winner and not result.winner.error: css_cls = "champion" badge = '★ CHAMP' elif c.method in (rnd.advanced_to or []): css_cls = "advanced" badge = 'ADV' else: css_cls = "eliminated" badge = 'OUT' # Score color if c.error: score_html = 'ERR' elif c.score >= 0.7: score_html = f'{c.score:.4f}' elif c.score >= 0.4: score_html = f'{c.score:.4f}' else: score_html = f'{c.score:.4f}' # 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'dir {html_mod.escape(dm)}' ) rr = m.get("refusal_rate") if rr is not None: metric_spans.append( f'ref {rr:.0%}' ) co = m.get("coherence") if co is not None: metric_spans.append( f'coh {co:.3f}' ) 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'cert {html_mod.escape(sc)}' ) kl = m.get("kl_divergence") if kl is not None: metric_spans.append( f'kl {kl:.4f}' ) pp = m.get("perplexity") if pp is not None: metric_spans.append( f'ppl {pp:.1f}' ) metrics_html = "".join(metric_spans) bracket_parts.append(f'
') bracket_parts.append( f'
' f'#{rank}' f'{html_mod.escape(c.method)}' f'{score_html}' f'{badge}' f'
' ) if metrics_html: bracket_parts.append(f'
{metrics_html}
') bracket_parts.append("
") bracket_parts.append("
") bracket_parts.append("
") return css + '
' + "\n".join(header_parts + bracket_parts) + "
" 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 = os.path.join(tempfile.gettempdir(), "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 = resolve_first("HF_PUSH_TOKEN", "HF_TOKEN") 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}")