mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-07-10 22:33:50 +02:00
Stella Biderman
877 lines
32 KiB
Python
877 lines
32 KiB
Python
"""CLI entry point for Obliteratus — Master Ablation Suite."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
|
|
_BANNER = r"""
|
|
[bold red]
|
|
░▒█▀▀▀█ ░▒█▀▀▄ ░▒█░░░ ▀█▀ ▀▀█▀▀ ░▒█▀▀▀ ░▒█▀▀█ ▒█▀▀█ ▀▀█▀▀ ░▒█░░▒█ ░▒█▀▀▀█
|
|
░▒█░░▒█ ░▒█▀▀▄ ░▒█░░░ ░█░ ░░█░░ ░▒█▀▀▀ ░▒█▄▄▀ ▒█▄▄█ ░░█░░ ░▒█░░▒█ ░░▀▀▀▄▄
|
|
░▒█▄▄▄█ ░▒█▄▄▀ ░▒█▄▄█ ▄█▄ ░░▀░░ ░▒█▄▄▄ ░▒█░▒█ ▒█░▒█ ░░▀░░ ░░▒█▄▄█ ░▒█▄▄▄█
|
|
[/bold red]
|
|
[dim] ════════════════════════════════════════════════════════════════════[/dim]
|
|
[bold white] MASTER ABLATION SUITE[/bold white] [dim]//[/dim] [bold red]Break the chains. Free the mind.[/bold red]
|
|
[dim] ════════════════════════════════════════════════════════════════════[/dim]
|
|
"""
|
|
|
|
|
|
def _add_gpu_args(parser):
|
|
"""Add --gpus flag for multi-GPU control."""
|
|
gpu_group = parser.add_argument_group("GPU selection")
|
|
gpu_group.add_argument(
|
|
"--gpus", type=str, default=None, metavar="IDS",
|
|
help=(
|
|
"Comma-separated GPU IDs to use (e.g. '0,1,2,3' or 'all'). "
|
|
"Sets CUDA_VISIBLE_DEVICES. By default uses all available GPUs. "
|
|
"Models are automatically split across selected GPUs via accelerate."
|
|
),
|
|
)
|
|
|
|
|
|
def _add_remote_args(parser):
|
|
"""Add --remote execution flags to a subcommand parser."""
|
|
remote_group = parser.add_argument_group("remote execution")
|
|
remote_group.add_argument(
|
|
"--remote", type=str, default=None, metavar="[USER@]HOST",
|
|
help="Run on a remote GPU node via SSH (e.g. root@gpu-node or just gpu-node)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--ssh-key", type=str, default=None,
|
|
help="Path to SSH private key (default: use SSH agent or ~/.ssh/id_rsa)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--ssh-port", type=int, default=22,
|
|
help="SSH port on remote host (default: 22)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--remote-dir", type=str, default="/tmp/obliteratus_run",
|
|
help="Working directory on the remote machine (default: /tmp/obliteratus_run)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--remote-python", type=str, default="python3",
|
|
help="Python binary on the remote machine (default: python3)",
|
|
)
|
|
remote_group.add_argument(
|
|
"--no-sync", action="store_true", default=False,
|
|
help="Don't copy results back to local machine after remote run",
|
|
)
|
|
|
|
|
|
def _apply_gpu_selection(args):
|
|
"""Set CUDA_VISIBLE_DEVICES based on --gpus flag (for local runs only)."""
|
|
import os
|
|
|
|
gpus = getattr(args, "gpus", None)
|
|
if gpus is None or getattr(args, "remote", None):
|
|
return # skip for remote runs (handled by remote runner)
|
|
|
|
if gpus.lower() == "all":
|
|
return # use all GPUs (default behavior)
|
|
|
|
# Validate: should be comma-separated integers
|
|
try:
|
|
gpu_ids = [int(g.strip()) for g in gpus.split(",")]
|
|
except ValueError:
|
|
console.print(f"[red]Invalid --gpus value: {gpus!r}. Expected comma-separated integers or 'all'.[/]")
|
|
raise SystemExit(1)
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
|
|
console.print(f"[dim]Using GPUs: {gpu_ids} (CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']})[/dim]")
|
|
|
|
|
|
def main(argv: list[str] | None = None):
|
|
console.print(_BANNER)
|
|
parser = argparse.ArgumentParser(
|
|
prog="obliteratus",
|
|
description="Master Ablation Suite for HuggingFace transformers",
|
|
)
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
|
|
# --- run ---
|
|
run_parser = subparsers.add_parser("run", help="Run an ablation from a YAML config")
|
|
run_parser.add_argument("config", type=str, help="Path to YAML config file")
|
|
run_parser.add_argument("--output-dir", type=str, default=None, help="Override output dir")
|
|
run_parser.add_argument(
|
|
"--preset",
|
|
type=str,
|
|
default=None,
|
|
help="Apply a preset (e.g. quick, full, attention, jailbreak, guardrail)",
|
|
)
|
|
_add_gpu_args(run_parser)
|
|
_add_remote_args(run_parser)
|
|
|
|
# --- info ---
|
|
info_parser = subparsers.add_parser("info", help="Print model architecture info")
|
|
info_parser.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
info_parser.add_argument("--task", type=str, default="causal_lm", choices=["causal_lm", "classification"])
|
|
info_parser.add_argument("--device", type=str, default="cpu")
|
|
info_parser.add_argument("--dtype", type=str, default="float32")
|
|
|
|
# --- interactive ---
|
|
subparsers.add_parser(
|
|
"interactive",
|
|
help="Guided setup — pick hardware, model, and preset interactively",
|
|
)
|
|
|
|
# --- models ---
|
|
models_parser = subparsers.add_parser("models", help="Browse curated models by compute tier")
|
|
models_parser.add_argument(
|
|
"--tier",
|
|
type=str,
|
|
default=None,
|
|
choices=["tiny", "small", "medium", "large", "frontier"],
|
|
help="Filter by compute tier",
|
|
)
|
|
|
|
# --- presets ---
|
|
subparsers.add_parser("presets", help="Browse ablation presets (quick, full, jailbreak, etc.)")
|
|
|
|
# --- strategies ---
|
|
subparsers.add_parser("strategies", help="List available ablation strategies")
|
|
|
|
# --- ui ---
|
|
ui_parser = subparsers.add_parser(
|
|
"ui",
|
|
help="Launch the Gradio web UI locally (same UI as the HuggingFace Space)",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--port", type=int, default=7860, help="Server port (default: 7860)",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--host", type=str, default="0.0.0.0", help="Server host (default: 0.0.0.0)",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--share", action="store_true", help="Create a public Gradio share link",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--no-browser", action="store_true", help="Don't auto-open browser on launch",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--auth", type=str, default=None,
|
|
help="Basic auth as user:pass",
|
|
)
|
|
ui_parser.add_argument(
|
|
"--quiet", action="store_true", help="Suppress the startup banner",
|
|
)
|
|
|
|
# --- obliterate (primary) + abliterate (backward-compat alias) ---
|
|
def _add_obliterate_args(p):
|
|
p.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
p.add_argument("--output-dir", type=str, default=None, help="Where to save the obliterated model")
|
|
p.add_argument("--device", type=str, default="auto")
|
|
p.add_argument("--dtype", type=str, default="float16")
|
|
p.add_argument(
|
|
"--method", type=str, default="advanced",
|
|
choices=[
|
|
"basic", "advanced", "aggressive", "spectral_cascade",
|
|
"informed", "surgical", "optimized", "inverted", "nuclear",
|
|
],
|
|
help="Liberation method (default: advanced)",
|
|
)
|
|
p.add_argument("--n-directions", type=int, default=None, help="Override: number of refusal directions to extract")
|
|
p.add_argument(
|
|
"--direction-method", type=str, default=None,
|
|
choices=["diff_means", "svd", "leace"],
|
|
help="Direction extraction method: diff_means (simple, robust), svd (multi-direction), leace (optimal erasure)",
|
|
)
|
|
p.add_argument("--regularization", type=float, default=None, help="Override: fraction to preserve (0.0-1.0)")
|
|
p.add_argument("--refinement-passes", type=int, default=None, help="Override: number of iterative passes")
|
|
p.add_argument(
|
|
"--quantization", type=str, default=None, choices=["4bit", "8bit"],
|
|
help="Load model with quantization (4bit or 8bit). Requires bitsandbytes.",
|
|
)
|
|
p.add_argument(
|
|
"--large-model", action="store_true", default=False,
|
|
help="Enable conservative defaults for 120B+ models (fewer directions, 1 pass, lower SAE expansion).",
|
|
)
|
|
p.add_argument(
|
|
"--verify-sample-size", type=int, default=None,
|
|
help="Number of harmful prompts to test for refusal rate (default: 30). "
|
|
"Increase for tighter confidence intervals (e.g. 100 for ~1%% resolution).",
|
|
)
|
|
p.add_argument(
|
|
"--contribute", action="store_true", default=False,
|
|
help="Save a community contribution record after the run completes.",
|
|
)
|
|
p.add_argument(
|
|
"--contribute-notes", type=str, default="",
|
|
help="Optional notes to include with the community contribution.",
|
|
)
|
|
|
|
abl_parser = subparsers.add_parser(
|
|
"obliterate",
|
|
help="One-click: remove refusal directions from a model (SOTA multi-technique)",
|
|
)
|
|
_add_obliterate_args(abl_parser)
|
|
_add_gpu_args(abl_parser)
|
|
_add_remote_args(abl_parser)
|
|
# Backward-compat alias (hidden from help)
|
|
abl_alias = subparsers.add_parser("abliterate", help=argparse.SUPPRESS)
|
|
_add_obliterate_args(abl_alias)
|
|
_add_gpu_args(abl_alias)
|
|
_add_remote_args(abl_alias)
|
|
|
|
# --- report ---
|
|
report_parser = subparsers.add_parser("report", help="Regenerate report from saved results")
|
|
report_parser.add_argument("results_json", type=str, help="Path to results.json")
|
|
report_parser.add_argument("--output-dir", type=str, default=None)
|
|
|
|
# --- aggregate ---
|
|
aggregate_parser = subparsers.add_parser("aggregate", help="Aggregate community contribution results")
|
|
aggregate_parser.add_argument(
|
|
"--dir", type=str, default="community_results",
|
|
help="Directory containing contribution JSON files",
|
|
)
|
|
|
|
# --- tourney ---
|
|
tourney_parser = subparsers.add_parser(
|
|
"tourney",
|
|
help="March Madness tournament — pit all methods against each other, push winner to Hub",
|
|
)
|
|
tourney_parser.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
tourney_parser.add_argument("--hub-org", type=str, default=None, help="HF org to push winner (e.g. my-org)")
|
|
tourney_parser.add_argument("--hub-repo", type=str, default=None, help="Full HF repo ID (overrides --hub-org)")
|
|
tourney_parser.add_argument("--device", type=str, default="auto")
|
|
tourney_parser.add_argument("--dtype", type=str, default="float16")
|
|
tourney_parser.add_argument("--dataset", type=str, default="builtin", help="Dataset source (default: builtin)")
|
|
tourney_parser.add_argument(
|
|
"--quantization", type=str, default=None, choices=["4bit", "8bit"],
|
|
help="Load model with quantization",
|
|
)
|
|
tourney_parser.add_argument("--output-dir", type=str, default="/tmp/obliteratus_tourney")
|
|
tourney_parser.add_argument(
|
|
"--methods", type=str, nargs="+", default=None,
|
|
help="Override: only run these methods (space-separated)",
|
|
)
|
|
_add_gpu_args(tourney_parser)
|
|
_add_remote_args(tourney_parser)
|
|
|
|
# --- recommend ---
|
|
recommend_parser = subparsers.add_parser(
|
|
"recommend",
|
|
help="Show telemetry-driven best method + hyperparams for a model",
|
|
)
|
|
recommend_parser.add_argument("model", type=str, help="HuggingFace model name/path")
|
|
recommend_parser.add_argument("--device", type=str, default="cpu")
|
|
recommend_parser.add_argument("--dtype", type=str, default="float32")
|
|
recommend_parser.add_argument(
|
|
"--insights", action="store_true", default=False,
|
|
help="Also show global cross-architecture insights",
|
|
)
|
|
|
|
args = parser.parse_args(argv)
|
|
|
|
# Apply GPU selection early (before any CUDA init)
|
|
_apply_gpu_selection(args)
|
|
|
|
if args.command == "run":
|
|
if getattr(args, "remote", None):
|
|
_cmd_remote_run(args)
|
|
else:
|
|
_cmd_run(args)
|
|
elif args.command == "interactive":
|
|
_cmd_interactive()
|
|
elif args.command == "models":
|
|
_cmd_models(args)
|
|
elif args.command == "presets":
|
|
_cmd_presets()
|
|
elif args.command == "info":
|
|
_cmd_info(args)
|
|
elif args.command == "strategies":
|
|
_cmd_strategies()
|
|
elif args.command == "report":
|
|
_cmd_report(args)
|
|
elif args.command == "aggregate":
|
|
_cmd_aggregate(args)
|
|
elif args.command == "ui":
|
|
_cmd_ui(args)
|
|
elif args.command == "recommend":
|
|
_cmd_recommend(args)
|
|
elif args.command == "tourney":
|
|
if getattr(args, "remote", None):
|
|
_cmd_remote_tourney(args)
|
|
else:
|
|
_cmd_tourney(args)
|
|
elif args.command in ("obliterate", "abliterate"):
|
|
if getattr(args, "remote", None):
|
|
_cmd_remote_abliterate(args)
|
|
else:
|
|
_cmd_abliterate(args)
|
|
|
|
|
|
def _cmd_ui(args):
|
|
from obliteratus.local_ui import launch_local_ui
|
|
|
|
auth = tuple(args.auth.split(":", 1)) if args.auth else None
|
|
launch_local_ui(
|
|
host=args.host,
|
|
port=args.port,
|
|
share=args.share,
|
|
open_browser=not args.no_browser,
|
|
auth=auth,
|
|
quiet=args.quiet,
|
|
)
|
|
|
|
|
|
def _cmd_interactive():
|
|
from obliteratus.interactive import run_interactive
|
|
run_interactive()
|
|
|
|
|
|
def _cmd_models(args):
|
|
from rich.table import Table
|
|
|
|
from obliteratus.presets import get_presets_by_tier, list_all_presets
|
|
|
|
presets = get_presets_by_tier(args.tier) if args.tier else list_all_presets()
|
|
|
|
table = Table(title="Model Library — Curated Targets")
|
|
table.add_column("Model", style="green")
|
|
table.add_column("HuggingFace ID", style="cyan")
|
|
table.add_column("Params", justify="right")
|
|
table.add_column("Tier", style="yellow")
|
|
table.add_column("Dtype")
|
|
table.add_column("Quant")
|
|
table.add_column("Description")
|
|
|
|
for p in presets:
|
|
table.add_row(
|
|
p.name,
|
|
p.hf_id,
|
|
p.params,
|
|
p.tier.upper(),
|
|
p.recommended_dtype,
|
|
p.recommended_quantization or "—",
|
|
p.description,
|
|
)
|
|
|
|
console.print(table)
|
|
console.print(
|
|
"\n[dim]Tiers: TINY = CPU/laptop | SMALL = 4-8GB | "
|
|
"MEDIUM = 8-16GB | LARGE = 24GB+ | FRONTIER = multi-GPU/cloud[/dim]"
|
|
)
|
|
|
|
|
|
def _cmd_presets():
|
|
from rich.table import Table
|
|
|
|
from obliteratus.study_presets import list_study_presets
|
|
|
|
presets = list_study_presets()
|
|
|
|
table = Table(title="Ablation Presets")
|
|
table.add_column("Key", style="cyan", min_width=12)
|
|
table.add_column("Name", style="green")
|
|
table.add_column("Strategies", style="yellow")
|
|
table.add_column("Samples", justify="right")
|
|
table.add_column("Description", max_width=55)
|
|
|
|
for p in presets:
|
|
strats = ", ".join(s["name"] for s in p.strategies)
|
|
table.add_row(p.key, p.name, strats, str(p.max_samples), p.description)
|
|
|
|
console.print(table)
|
|
console.print(
|
|
"\n[dim]Usage: obliteratus run config.yaml --preset quick\n"
|
|
" or: set preset: quick in your YAML file[/dim]"
|
|
)
|
|
|
|
|
|
def _cmd_run(args):
|
|
from obliteratus.config import StudyConfig
|
|
from obliteratus.runner import run_study
|
|
|
|
config = StudyConfig.from_yaml(args.config)
|
|
# If --preset flag given, inject it so from_dict picks it up
|
|
if args.preset:
|
|
import yaml
|
|
|
|
raw = yaml.safe_load(Path(args.config).read_text())
|
|
raw["preset"] = args.preset
|
|
config = StudyConfig.from_dict(raw)
|
|
if args.output_dir:
|
|
config.output_dir = args.output_dir
|
|
|
|
# If YAML has a remote: section, dispatch to remote runner
|
|
if config.remote is not None:
|
|
from obliteratus.remote import RemoteConfig as _RC, RemoteRunner
|
|
|
|
rc = _RC(
|
|
host=config.remote.host,
|
|
user=config.remote.user,
|
|
port=config.remote.port,
|
|
ssh_key=config.remote.ssh_key,
|
|
remote_dir=config.remote.remote_dir,
|
|
python=config.remote.python,
|
|
sync_results=config.remote.sync_results,
|
|
gpus=config.remote.gpus,
|
|
)
|
|
runner = RemoteRunner(rc)
|
|
result_path = runner.run_config(
|
|
local_config_path=args.config,
|
|
local_output_dir=config.output_dir,
|
|
preset=args.preset,
|
|
)
|
|
if result_path:
|
|
console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
|
|
else:
|
|
console.print("[red]Remote run failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
return
|
|
|
|
run_study(config)
|
|
|
|
|
|
def _cmd_info(args):
|
|
from obliteratus.models.loader import load_model
|
|
|
|
console.print(f"[bold cyan]Loading model:[/bold cyan] {args.model}")
|
|
handle = load_model(
|
|
model_name=args.model,
|
|
task=args.task,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
)
|
|
summary = handle.summary()
|
|
for key, val in summary.items():
|
|
if isinstance(val, int) and val > 1000:
|
|
console.print(f" {key}: {val:,}")
|
|
else:
|
|
console.print(f" {key}: {val}")
|
|
|
|
|
|
def _cmd_strategies():
|
|
from obliteratus.strategies import STRATEGY_REGISTRY
|
|
|
|
console.print("[bold]Available ablation strategies:[/bold]\n")
|
|
for name, cls in sorted(STRATEGY_REGISTRY.items()):
|
|
doc = (cls.__doc__ or "").strip().split("\n")[0]
|
|
console.print(f" [cyan]{name}[/cyan] — {doc}")
|
|
|
|
|
|
def _cmd_report(args):
|
|
from obliteratus.reporting.report import AblationReport, AblationResult
|
|
|
|
path = Path(args.results_json)
|
|
data = json.loads(path.read_text())
|
|
|
|
report = AblationReport(model_name=data["model_name"])
|
|
report.add_baseline(data["baseline_metrics"])
|
|
for r in data["results"]:
|
|
report.add_result(
|
|
AblationResult(
|
|
strategy=r["strategy"],
|
|
component=r["component"],
|
|
description=r["description"],
|
|
metrics=r["metrics"],
|
|
metadata=r.get("metadata"),
|
|
)
|
|
)
|
|
|
|
report.print_summary()
|
|
|
|
output_dir = Path(args.output_dir) if args.output_dir else path.parent
|
|
metric_name = list(data["baseline_metrics"].keys())[0]
|
|
try:
|
|
report.plot_impact(metric=metric_name, output_path=output_dir / "impact.png")
|
|
report.plot_heatmap(output_path=output_dir / "heatmap.png")
|
|
console.print(f"\nPlots saved to {output_dir}/")
|
|
except Exception as e:
|
|
console.print(f"[yellow]Could not generate plots: {e}[/yellow]")
|
|
|
|
|
|
def _cmd_aggregate(args):
|
|
from obliteratus.community import aggregate_results, load_contributions
|
|
|
|
contrib_dir = args.dir
|
|
records = load_contributions(contrib_dir)
|
|
if not records:
|
|
console.print(f"[yellow]No contributions found in {contrib_dir}[/yellow]")
|
|
return
|
|
|
|
aggregated = aggregate_results(records)
|
|
|
|
from rich.table import Table
|
|
|
|
table = Table(title="Aggregated Community Results")
|
|
table.add_column("Model", style="green")
|
|
table.add_column("Method", style="cyan")
|
|
table.add_column("Runs", justify="right")
|
|
table.add_column("Mean Refusal", justify="right")
|
|
table.add_column("Mean Perplexity", justify="right")
|
|
|
|
for model_name, methods in sorted(aggregated.items()):
|
|
for method_name, stats in sorted(methods.items()):
|
|
refusal = stats.get("refusal_rate", {}).get("mean", "N/A")
|
|
ppl = stats.get("perplexity", {}).get("mean", "N/A")
|
|
if isinstance(refusal, float):
|
|
refusal = f"{refusal:.4f}"
|
|
if isinstance(ppl, float):
|
|
ppl = f"{ppl:.2f}"
|
|
table.add_row(
|
|
model_name.split("/")[-1] if "/" in model_name else model_name,
|
|
method_name,
|
|
str(stats["n_runs"]),
|
|
str(refusal),
|
|
str(ppl),
|
|
)
|
|
|
|
console.print(table)
|
|
|
|
|
|
def _cmd_recommend(args):
|
|
from rich.markdown import Markdown
|
|
from rich.panel import Panel
|
|
|
|
from obliteratus.architecture_profiles import detect_architecture, enhance_profile_with_telemetry
|
|
from obliteratus.adaptive_defaults import format_recommendation, get_global_insights
|
|
|
|
model_name = args.model
|
|
console.print(f"\nAnalyzing [bold]{model_name}[/]...")
|
|
|
|
# Detect architecture
|
|
try:
|
|
from transformers import AutoConfig
|
|
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
|
|
num_layers = getattr(config, "num_hidden_layers", 0)
|
|
hidden_size = getattr(config, "hidden_size", 0)
|
|
except Exception:
|
|
config = None
|
|
num_layers = 0
|
|
hidden_size = 0
|
|
|
|
profile = detect_architecture(model_name, config, num_layers, hidden_size)
|
|
profile, rec = enhance_profile_with_telemetry(profile)
|
|
|
|
console.print(Panel(
|
|
f"[bold]{profile.profile_label}[/]\n"
|
|
f"Architecture: {profile.arch_class.value} | Reasoning: {profile.reasoning_class.value}\n"
|
|
f"Params: ~{profile.total_params_b:.1f}B | Layers: {profile.num_layers} | "
|
|
f"Hidden: {profile.hidden_size}",
|
|
title="Architecture Profile",
|
|
border_style="cyan",
|
|
))
|
|
|
|
if rec:
|
|
md = format_recommendation(rec)
|
|
console.print(Markdown(md))
|
|
else:
|
|
console.print("\n[yellow]Could not fetch telemetry — using research-grounded defaults.[/]")
|
|
|
|
console.print(f"\n[bold green]Research default method:[/] {profile.recommended_method}")
|
|
if profile.method_overrides:
|
|
console.print("[bold green]Overrides:[/]")
|
|
for k, v in sorted(profile.method_overrides.items()):
|
|
console.print(f" {k}: {v}")
|
|
|
|
if args.insights:
|
|
console.print("\n")
|
|
console.rule("[bold magenta]Global Telemetry Insights")
|
|
insights = get_global_insights()
|
|
console.print(f"Total records analyzed: {insights['total_records']}")
|
|
if insights["overall_best_methods"]:
|
|
console.print("\n[bold]Overall method ranking (all architectures):[/]")
|
|
for entry in insights["overall_best_methods"][:10]:
|
|
console.print(
|
|
f" {entry['method']}: {entry['mean_score']:.4f} "
|
|
f"({entry['n_runs']} runs)"
|
|
)
|
|
if insights["architecture_breakdown"]:
|
|
console.print("\n[bold]Per-architecture breakdown:[/]")
|
|
for label, info in insights["architecture_breakdown"].items():
|
|
console.print(
|
|
f" {label}: best={info['best_method']} "
|
|
f"({info['best_score']:.4f}), "
|
|
f"{info['n_methods_tested']} methods tested, "
|
|
f"{info['total_runs']} runs"
|
|
)
|
|
|
|
|
|
def _cmd_tourney(args):
|
|
from obliteratus.tourney import TourneyRunner, render_bracket
|
|
|
|
def on_log(msg):
|
|
console.print(msg)
|
|
|
|
def on_round(rnd):
|
|
console.print()
|
|
console.rule(f"[bold green]Round {rnd.round_num} complete — "
|
|
f"{len(rnd.advanced_to)} advance, {len(rnd.eliminated)} eliminated")
|
|
|
|
runner = TourneyRunner(
|
|
model_name=args.model,
|
|
hub_org=args.hub_org,
|
|
hub_repo=args.hub_repo,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
dataset_key=args.dataset,
|
|
quantization=args.quantization,
|
|
methods=args.methods,
|
|
output_dir=args.output_dir,
|
|
on_log=on_log,
|
|
on_round=on_round,
|
|
)
|
|
|
|
result = runner.run()
|
|
|
|
if result.winner:
|
|
console.print()
|
|
console.rule("[bold magenta]TOURNAMENT CHAMPION", style="magenta")
|
|
console.print(f" [bold]{result.winner.method}[/] — score {result.winner.score:.4f}")
|
|
console.print(f" Refusal rate: {result.winner.metrics.get('refusal_rate', '?')}")
|
|
console.print(f" Coherence: {result.winner.metrics.get('coherence', '?')}")
|
|
if result.hub_repo:
|
|
console.print(f" Pushed to: [link=https://huggingface.co/{result.hub_repo}]{result.hub_repo}[/link]")
|
|
console.print(f"\n Full bracket: {args.output_dir}/tourney_bracket.md")
|
|
|
|
|
|
def _cmd_abliterate(args):
|
|
from rich.live import Live
|
|
from rich.panel import Panel
|
|
from rich.table import Table
|
|
from rich.text import Text
|
|
|
|
from obliteratus.abliterate import METHODS, STAGES, AbliterationPipeline
|
|
|
|
model_name = args.model
|
|
output_dir = args.output_dir or f"abliterated/{model_name.replace('/', '_')}"
|
|
method = args.method
|
|
method_label = METHODS.get(method, {}).get("label", method)
|
|
|
|
# Stage state tracking
|
|
stage_status = {s.key: "waiting" for s in STAGES}
|
|
stage_msgs = {s.key: "" for s in STAGES}
|
|
log_lines: list[str] = []
|
|
|
|
def make_display():
|
|
table = Table(show_header=False, expand=True, border_style="green")
|
|
table.add_column("", width=6)
|
|
table.add_column("Stage", min_width=10)
|
|
table.add_column("Status", min_width=50)
|
|
for i, s in enumerate(STAGES):
|
|
st = stage_status[s.key]
|
|
if st == "done":
|
|
icon = "[bold green]✓[/]"
|
|
bar = "[green]" + "█" * 20 + "[/]"
|
|
elif st == "running":
|
|
icon = "[bold yellow]⚡[/]"
|
|
bar = "[yellow]" + "▓" * 10 + "░" * 10 + "[/]"
|
|
else:
|
|
icon = "[dim]○[/]"
|
|
bar = "[dim]" + "░" * 20 + "[/]"
|
|
msg = stage_msgs.get(s.key, "")
|
|
table.add_row(
|
|
f"[cyan][{i + 1}/6][/]",
|
|
f"{icon} [bold]{s.name}[/]",
|
|
f"{bar} {msg}",
|
|
)
|
|
|
|
header = Text.from_markup(
|
|
f"[bold green]OBLITERATUS — ABLITERATION PIPELINE[/]\n"
|
|
f"[dim]Target:[/] [cyan]{model_name}[/] → [cyan]{output_dir}[/]\n"
|
|
f"[dim]Method:[/] [magenta]{method_label}[/]"
|
|
)
|
|
|
|
# Last 12 log lines
|
|
recent = log_lines[-12:] if log_lines else ["Initializing..."]
|
|
log_text = "\n".join(f"[dim]>[/] {line}" for line in recent)
|
|
|
|
return Panel(
|
|
f"{header}\n\n{table}\n\n[dim]─── LOG ───[/]\n{log_text}",
|
|
border_style="green",
|
|
title="[bold green]⚗ ABLITERATE ⚗[/]",
|
|
)
|
|
|
|
def on_stage(result):
|
|
stage_status[result.stage] = result.status
|
|
stage_msgs[result.stage] = result.message
|
|
if live:
|
|
live.update(make_display())
|
|
|
|
def on_log(msg):
|
|
log_lines.append(msg)
|
|
if live:
|
|
live.update(make_display())
|
|
|
|
live = None
|
|
pipeline = AbliterationPipeline(
|
|
model_name=model_name,
|
|
output_dir=output_dir,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
method=method,
|
|
n_directions=args.n_directions,
|
|
direction_method=getattr(args, "direction_method", None),
|
|
regularization=args.regularization,
|
|
refinement_passes=args.refinement_passes,
|
|
quantization=args.quantization,
|
|
large_model_mode=getattr(args, "large_model", False),
|
|
verify_sample_size=getattr(args, "verify_sample_size", None),
|
|
on_stage=on_stage,
|
|
on_log=on_log,
|
|
)
|
|
|
|
with Live(make_display(), console=console, refresh_per_second=4) as live_ctx:
|
|
live = live_ctx
|
|
try:
|
|
result_path = pipeline.run()
|
|
live.update(make_display())
|
|
except Exception as e:
|
|
log_lines.append(f"[red]ERROR: {e}[/]")
|
|
live.update(make_display())
|
|
raise
|
|
|
|
# ── Telemetry: send pipeline report to community leaderboard ──
|
|
try:
|
|
from obliteratus.telemetry import maybe_send_pipeline_report
|
|
maybe_send_pipeline_report(pipeline)
|
|
except Exception:
|
|
pass # Telemetry is best-effort
|
|
|
|
# ── Community contribution (--contribute flag) ──
|
|
contrib_path = None
|
|
if getattr(args, "contribute", False):
|
|
try:
|
|
from obliteratus.community import save_contribution
|
|
contrib_path = save_contribution(
|
|
pipeline,
|
|
model_name=model_name,
|
|
notes=getattr(args, "contribute_notes", ""),
|
|
)
|
|
except Exception as e:
|
|
console.print(f"[yellow]Could not save contribution: {e}[/yellow]")
|
|
|
|
console.print()
|
|
contrib_line = ""
|
|
if contrib_path:
|
|
contrib_line = f"\n Contribution: [cyan]{contrib_path}[/]"
|
|
console.print(
|
|
Panel(
|
|
f"[bold green]Abliteration complete![/]\n\n"
|
|
f" Model saved to: [cyan]{result_path}[/]\n"
|
|
f" Metadata: [cyan]{result_path}/abliteration_metadata.json[/]"
|
|
f"{contrib_line}\n\n"
|
|
f" [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
|
|
border_style="green",
|
|
title="[bold green]✓ REBIRTH COMPLETE[/]",
|
|
)
|
|
)
|
|
|
|
|
|
def _make_remote_runner(args):
|
|
"""Create a RemoteRunner from CLI --remote flags."""
|
|
from obliteratus.remote import RemoteConfig, RemoteRunner
|
|
|
|
rc = RemoteConfig.from_cli_args(
|
|
args.remote,
|
|
port=args.ssh_port,
|
|
ssh_key=args.ssh_key,
|
|
remote_dir=args.remote_dir,
|
|
python=args.remote_python,
|
|
sync_results=not args.no_sync,
|
|
gpus=getattr(args, "gpus", None),
|
|
)
|
|
return RemoteRunner(rc)
|
|
|
|
|
|
def _cmd_remote_abliterate(args):
|
|
from rich.panel import Panel
|
|
|
|
runner = _make_remote_runner(args)
|
|
|
|
kwargs = {}
|
|
if args.method:
|
|
kwargs["method"] = args.method
|
|
if args.device:
|
|
kwargs["device"] = args.device
|
|
if args.dtype:
|
|
kwargs["dtype"] = args.dtype
|
|
if args.quantization:
|
|
kwargs["quantization"] = args.quantization
|
|
if args.n_directions is not None:
|
|
kwargs["n_directions"] = args.n_directions
|
|
if getattr(args, "direction_method", None):
|
|
kwargs["direction_method"] = args.direction_method
|
|
if args.regularization is not None:
|
|
kwargs["regularization"] = args.regularization
|
|
if args.refinement_passes is not None:
|
|
kwargs["refinement_passes"] = args.refinement_passes
|
|
if getattr(args, "large_model", False):
|
|
kwargs["large_model"] = True
|
|
if getattr(args, "verify_sample_size", None) is not None:
|
|
kwargs["verify_sample_size"] = args.verify_sample_size
|
|
|
|
result_path = runner.run_obliterate(
|
|
model=args.model,
|
|
local_output_dir=args.output_dir,
|
|
**kwargs,
|
|
)
|
|
|
|
if result_path:
|
|
console.print(
|
|
Panel(
|
|
f"[bold green]Remote abliteration complete![/]\n\n"
|
|
f" Results at: [cyan]{result_path}[/]\n\n"
|
|
f" [dim]Load with:[/] AutoModelForCausalLM.from_pretrained('{result_path}')",
|
|
border_style="green",
|
|
title="[bold green]REBIRTH COMPLETE (remote)[/]",
|
|
)
|
|
)
|
|
else:
|
|
console.print("[red]Remote abliteration failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
|
|
|
|
def _cmd_remote_run(args):
|
|
runner = _make_remote_runner(args)
|
|
result_path = runner.run_config(
|
|
local_config_path=args.config,
|
|
local_output_dir=args.output_dir,
|
|
preset=args.preset,
|
|
)
|
|
if result_path:
|
|
console.print(f"\n[bold green]Remote run complete.[/] Results at: [cyan]{result_path}[/]")
|
|
else:
|
|
console.print("[red]Remote run failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
|
|
|
|
def _cmd_remote_tourney(args):
|
|
from rich.panel import Panel
|
|
|
|
runner = _make_remote_runner(args)
|
|
result_path = runner.run_tourney(
|
|
model=args.model,
|
|
local_output_dir=args.output_dir,
|
|
device=args.device,
|
|
dtype=args.dtype,
|
|
quantization=args.quantization,
|
|
methods=args.methods,
|
|
hub_org=args.hub_org,
|
|
hub_repo=args.hub_repo,
|
|
dataset=args.dataset,
|
|
)
|
|
if result_path:
|
|
console.print(
|
|
Panel(
|
|
f"[bold green]Remote tournament complete![/]\n\n"
|
|
f" Results at: [cyan]{result_path}[/]",
|
|
border_style="green",
|
|
title="[bold green]TOURNAMENT COMPLETE (remote)[/]",
|
|
)
|
|
)
|
|
else:
|
|
console.print("[red]Remote tournament failed. Check logs above.[/]")
|
|
raise SystemExit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|