"""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", "som", "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", "som"], help="Direction extraction method: diff_means, svd, leace, or som", ) 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( "--min-layer-fraction", type=float, default=None, help="Optional layer floor as fraction of depth; e.g. 0.75 keeps only the final quarter.", ) p.add_argument( "--max-layer-fraction", type=float, default=None, help="Optional cap on selected layers as fraction of depth; e.g. 0.25 caps edits to one quarter.", ) p.add_argument( "--harmless-pc-count", type=int, default=None, help="Subtract this many top harmless activation PCs from extracted refusal directions.", ) p.add_argument( "--shield-concept-count", type=int, default=None, help="Use this many built-in capability/style shield atoms for direction cleaning.", ) p.add_argument( "--shield-ridge", type=float, default=None, help="Ridge strength for shield atom residualization (default: 0.05).", ) p.add_argument( "--shield-residualize", action="store_true", default=None, help="Directly residualize refusal directions against shield atoms.", ) p.add_argument( "--shield-layer-penalty", type=float, default=None, help="Penalize layer selection by shield-atom overlap instead of editing directions.", ) p.add_argument( "--projection-target", type=str, default=None, choices=["all", "attention", "ffn", "output"], help="Projection target modules: all, attention, ffn, or output.", ) p.add_argument( "--projection-row-fraction", type=float, default=None, help="Only project the strongest fraction of output rows/columns per weight matrix.", ) 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( "--dataset", type=str, default="builtin", help="Prompt dataset source for contrastive extraction when using residue mining (default: builtin).", ) p.add_argument( "--residue-file", action="append", default=[], help="Refusal-audit/residue JSON to upweight as hard negatives. Can be passed multiple times.", ) p.add_argument( "--residue-weight", type=int, default=5, help="How many extra copies to add per mined residue prompt (default: 5).", ) p.add_argument( "--residue-max", type=int, default=None, help="Optional cap on unique residue prompts loaded from residue files.", ) 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) # --- self-improve --- si_parser = subparsers.add_parser( "self-improve", help="Recursive hard-negative mining: turn refusal audits into weighted residue runs", ) si_parser.add_argument("model", type=str, help="Source model or previous artifact to improve") si_parser.add_argument("--audit", action="append", required=True, help="Refusal-audit JSON/summary. Can be repeated.") si_parser.add_argument("--output-dir", type=str, required=True, help="Output directory for the next candidate") si_parser.add_argument("--residue-out", type=str, default=None, help="Where to write canonical residue JSON") si_parser.add_argument("--dataset", type=str, default="builtin") si_parser.add_argument("--residue-weight", type=int, default=None) si_parser.add_argument("--residue-max", type=int, default=None) si_parser.add_argument("--params-b", type=float, default=None, help="Override detected total parameter count in billions for planning/defaults") si_parser.add_argument("--no-param-auto-scale", action="store_true", default=False, help="Do not apply size-aware defaults when knobs are omitted") si_parser.add_argument("--method", type=str, default="advanced", choices=["basic", "advanced", "aggressive", "spectral_cascade", "informed", "surgical", "optimized", "som", "inverted", "nuclear"]) si_parser.add_argument("--direction-method", type=str, default="diff_means", choices=["diff_means", "svd", "leace", "som"]) si_parser.add_argument("--n-directions", type=int, default=None) si_parser.add_argument("--regularization", type=float, default=None) si_parser.add_argument("--refinement-passes", type=int, default=None) si_parser.add_argument("--min-layer-fraction", type=float, default=None) si_parser.add_argument("--max-layer-fraction", type=float, default=None) si_parser.add_argument("--harmless-pc-count", type=int, default=None) si_parser.add_argument("--shield-concept-count", type=int, default=None) si_parser.add_argument("--shield-ridge", type=float, default=None) si_parser.add_argument("--shield-residualize", action="store_true", default=None) si_parser.add_argument("--shield-layer-penalty", type=float, default=None) si_parser.add_argument("--projection-target", type=str, default=None, choices=["all", "attention", "ffn", "output"]) si_parser.add_argument("--projection-row-fraction", type=float, default=None) si_parser.add_argument("--device", type=str, default="auto") si_parser.add_argument("--dtype", type=str, default="float16") si_parser.add_argument("--verify-sample-size", type=int, default=None) si_parser.add_argument("--dry-run", action="store_true", default=False, help="Only write residue/plan; do not run surgery") _add_gpu_args(si_parser) # --- 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", ) # --- gpu-calc --- calc_parser = subparsers.add_parser( "gpu-calc", help="Estimate minimum GPUs needed for a model", ) calc_parser.add_argument( "model", type=str, nargs="?", default=None, help="HuggingFace model name/path (auto-fetches param counts)", ) calc_parser.add_argument( "--params", type=float, default=None, metavar="B", help="Total parameters in billions (overrides auto-detection)", ) calc_parser.add_argument( "--active-params", type=float, default=None, metavar="B", help="Active parameters in billions (for MoE models; defaults to --params)", ) calc_parser.add_argument( "--dtype", type=str, default="bfloat16", choices=["float32", "float16", "bfloat16", "int8", "int4"], help="Data type for model weights (default: bfloat16)", ) calc_parser.add_argument( "--gpu-mem", type=float, default=80.0, metavar="GB", help="VRAM per GPU in GB (default: 80 for A100-80GB)", ) args = parser.parse_args(argv) # Apply GPU selection early (before any CUDA init) _apply_gpu_selection(args) if args.command == "gpu-calc": _cmd_gpu_calc(args) return elif 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 == "self-improve": _cmd_self_improve(args) elif args.command in ("obliterate", "abliterate"): if getattr(args, "remote", None): _cmd_remote_abliterate(args) else: _cmd_abliterate(args) def _cmd_self_improve(args): """Run one recursive hard-negative OBLITERATUS iteration.""" from rich.panel import Panel from rich.table import Table from obliteratus.hard_negative import ( build_weighted_prompt_pairs, load_residue_file, save_residue_file, ) from obliteratus.model_profile import ModelProfile, default_self_improve_params, profile_model profile = profile_model(args.model, dtype=args.dtype) if args.params_b is not None: profile = ModelProfile( **{**profile.to_json(), "total_params": int(args.params_b * 1e9), "total_params_b": args.params_b, "active_params_b": args.params_b, "source": "cli_override"} ) size_defaults = default_self_improve_params(profile) if not args.no_param_auto_scale else { "n_directions": 3, "regularization": 0.30, "refinement_passes": 1, "residue_weight": 5, "verify_sample_size": 30, "note": "parameter auto-scale disabled; using generic defaults for omitted knobs.", } n_directions = args.n_directions if args.n_directions is not None else size_defaults["n_directions"] regularization = args.regularization if args.regularization is not None else size_defaults["regularization"] refinement_passes = args.refinement_passes if args.refinement_passes is not None else size_defaults["refinement_passes"] residue_weight = args.residue_weight if args.residue_weight is not None else size_defaults["residue_weight"] verify_sample_size = args.verify_sample_size if args.verify_sample_size is not None else size_defaults["verify_sample_size"] residue_examples = [] for audit in args.audit: residue_examples.extend(load_residue_file(audit)) residue_path = Path(args.residue_out or Path(args.output_dir, "mined_residue.json")) save_residue_file(residue_examples, residue_path) harmful, harmless, meta = build_weighted_prompt_pairs( base_dataset=args.dataset, residue_files=[residue_path], residue_weight=residue_weight, max_residue=args.residue_max, ) meta["model_profile"] = profile.to_json() meta["size_aware_defaults"] = size_defaults meta["effective_settings"] = { "n_directions": n_directions, "regularization": regularization, "refinement_passes": refinement_passes, "residue_weight": residue_weight, "verify_sample_size": verify_sample_size, "projection_row_fraction": args.projection_row_fraction, } plan_path = Path(args.output_dir, "self_improve_plan.json") plan_path.parent.mkdir(parents=True, exist_ok=True) plan_path.write_text(json.dumps(meta, indent=2)) table = Table(title="Recursive hard-negative mining plan") table.add_column("Field", style="cyan") table.add_column("Value", style="green") table.add_row("source model", args.model) table.add_row("params", f"{profile.total_params_b:.3f}B" if profile.total_params_b is not None else "unknown") table.add_row("profile source", profile.source) table.add_row("layers / hidden", f"{profile.num_layers or '?'} / {profile.hidden_size or '?'}") table.add_row("size default note", size_defaults["note"]) table.add_row("output", args.output_dir) table.add_row("residue file", str(residue_path)) table.add_row("plan file", str(plan_path)) table.add_row("unique residue prompts", str(meta["residue_examples"])) table.add_row("weighted residue pairs", str(meta["residue_added_pairs"])) table.add_row("total prompt pairs", str(meta["total_pairs"])) table.add_row("method", args.method) table.add_row("directions", str(n_directions)) table.add_row("regularization", str(regularization)) table.add_row("refinement passes", str(refinement_passes)) table.add_row("verify sample size", str(verify_sample_size)) console.print(table) if args.dry_run: console.print( Panel( "[bold yellow]Dry run only.[/]\n\n" f"Plan JSON: {plan_path}\n" "Run without --dry-run to launch surgery, or pass this residue file to:\n" f" obliteratus obliterate {args.model} --residue-file {residue_path} --output-dir {args.output_dir}", border_style="yellow", title="SELF-IMPROVE PLAN WRITTEN", ) ) return from obliteratus.abliterate import AbliterationPipeline log_lines: list[str] = [] def on_log(msg): log_lines.append(msg) console.print(f"[dim]>[/] {msg}") pipeline = AbliterationPipeline( model_name=args.model, output_dir=args.output_dir, device=args.device, dtype=args.dtype, method=args.method, n_directions=n_directions, direction_method=args.direction_method, regularization=regularization, refinement_passes=refinement_passes, min_layer_fraction=args.min_layer_fraction, max_layer_fraction=args.max_layer_fraction, harmless_pc_count=args.harmless_pc_count, shield_concept_count=args.shield_concept_count, shield_ridge=args.shield_ridge, shield_residualize=args.shield_residualize, shield_layer_penalty=args.shield_layer_penalty, projection_target=args.projection_target, projection_row_fraction=args.projection_row_fraction, verify_sample_size=verify_sample_size, harmful_prompts=harmful, harmless_prompts=harmless, on_log=on_log, ) result_path = pipeline.run() Path(result_path, "hard_negative_residue.json").write_text(json.dumps(meta, indent=2)) console.print( Panel( f"[bold green]Self-improvement iteration complete![/]\n\n" f" Model saved to: [cyan]{result_path}[/]\n" f" Residue metadata: [cyan]{result_path}/hard_negative_residue.json[/]\n" f" Next loop: audit refusals from this artifact and feed them back with --audit.", border_style="green", title="RECURSIVE REBIRTH COMPLETE", ) ) 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 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 prompt_kwargs = {} residue_meta = None if getattr(args, "residue_file", None): from obliteratus.hard_negative import build_weighted_prompt_pairs harmful, harmless, residue_meta = build_weighted_prompt_pairs( base_dataset=getattr(args, "dataset", "builtin"), residue_files=args.residue_file, residue_weight=getattr(args, "residue_weight", 5), max_residue=getattr(args, "residue_max", None), ) prompt_kwargs = {"harmful_prompts": harmful, "harmless_prompts": harmless} log_lines.append( "Loaded hard-negative residue: " f"{residue_meta['residue_examples']} unique prompts, " f"+{residue_meta['residue_added_pairs']} weighted pairs " f"({residue_meta['total_pairs']} total)." ) 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, min_layer_fraction=getattr(args, "min_layer_fraction", None), max_layer_fraction=getattr(args, "max_layer_fraction", None), harmless_pc_count=getattr(args, "harmless_pc_count", None), shield_concept_count=getattr(args, "shield_concept_count", None), shield_ridge=getattr(args, "shield_ridge", None), shield_residualize=getattr(args, "shield_residualize", None), shield_layer_penalty=getattr(args, "shield_layer_penalty", None), projection_target=getattr(args, "projection_target", None), projection_row_fraction=getattr(args, "projection_row_fraction", None), 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, **prompt_kwargs, ) with Live(make_display(), console=console, refresh_per_second=4) as live_ctx: live = live_ctx try: result_path = pipeline.run() if residue_meta: import json Path(result_path, "hard_negative_residue.json").write_text( json.dumps(residue_meta, indent=2) ) 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 _cmd_gpu_calc(args): import math from rich.panel import Panel from rich.table import Table BYTES_PER_PARAM = { "float32": 4, "float16": 2, "bfloat16": 2, "int8": 1, "int4": 0.5, } # Resolve param counts total_params_b = args.params active_params_b = args.active_params if total_params_b is None: if args.model is None: console.print("[red]Provide either a model name or --params.[/]") raise SystemExit(1) console.print(f"Profiling [cyan]{args.model}[/]...") try: from obliteratus.model_profile import profile_model profile = profile_model(args.model, dtype=args.dtype) total_params_b = profile.total_params_b if active_params_b is None: active_params_b = profile.active_params_b except Exception as e: console.print(f"[red]Could not profile model: {e}[/]") raise SystemExit(1) if total_params_b is None: console.print("[red]Cannot determine model size. Use --params.[/]") raise SystemExit(1) if active_params_b is None: active_params_b = total_params_b bpp = BYTES_PER_PARAM[args.dtype] gpu_mem_gb = args.gpu_mem # Model weight memory (use base-10 GB to match HF/nvidia conventions) weight_gb = total_params_b * bpp # Activation overhead during forward passes (PROBE/VERIFY). # Scales with active params, not total. Empirical from benchmarks: # - DeepSeek-70B (149GB): failed at 160GB (2 GPUs), OK at 240GB (3 GPUs) # - GPT-OSS-120B (234GB): failed at 240GB (3 GPUs), OK at 320GB (4 GPUs) # This implies ~15-35% overhead. We use 20% as a reasonable middle ground. active_weight_gb = active_params_b * bpp activation_overhead_gb = active_weight_gb * 0.20 # CUDA context + fragmentation overhead: ~1.5 GB per GPU (fixed cost) cuda_overhead_per_gpu = 1.5 # Total memory needed (before splitting across GPUs) total_needed_gb = weight_gb + activation_overhead_gb # Find minimum GPUs: we need total_needed / (gpu_mem - cuda_overhead) GPUs usable_per_gpu = gpu_mem_gb - cuda_overhead_per_gpu if usable_per_gpu <= 0: console.print("[red]GPU memory too small after CUDA overhead.[/]") raise SystemExit(1) min_gpus = math.ceil(total_needed_gb / usable_per_gpu) min_gpus = max(min_gpus, 1) # Show results for a range of GPU counts is_moe = active_params_b < total_params_b * 0.99 table = Table(title="GPU Configurations", show_edge=True) table.add_column("GPUs", justify="right", style="cyan") table.add_column("VRAM/GPU", justify="right") table.add_column("Total VRAM", justify="right") table.add_column("Headroom", justify="right") table.add_column("Verdict", min_width=20) # Show from min_gpus-1 (to show why it fails) up to 8 low = max(1, min_gpus - 1) high = max(min_gpus + 3, 8) for n in range(low, high + 1): total_vram = n * gpu_mem_gb usable_vram = n * usable_per_gpu headroom = usable_vram - total_needed_gb headroom_pct = headroom / total_needed_gb * 100 vram_per = total_needed_gb / n if headroom < 0: verdict = "[red]INSUFFICIENT[/]" elif headroom_pct < 15: verdict = "[yellow]TIGHT — may fail[/]" elif n == min_gpus: verdict = "[bold green]MINIMUM (recommended)[/]" else: verdict = "[green]OK[/] [dim](more GPUs = slower)[/]" table.add_row( str(n), f"{vram_per:.1f} GB", f"{total_vram:.0f} GB", f"{headroom:+.1f} GB ({headroom_pct:+.0f}%)", verdict, ) model_label = args.model or f"{total_params_b:.1f}B params" moe_line = "" if is_moe: moe_line = f"\n Active params: [cyan]{active_params_b:.1f}B[/] ({active_params_b/total_params_b*100:.0f}% of total — MoE)" console.print(Panel( f" Model: [cyan]{model_label}[/]\n" f" Total params: [cyan]{total_params_b:.1f}B[/]" f"{moe_line}\n" f" Dtype: [cyan]{args.dtype}[/] ({bpp} bytes/param)\n" f" Weight memory: [cyan]{weight_gb:.1f} GB[/]\n" f" Activation est: [cyan]{activation_overhead_gb:.1f} GB[/]\n" f" Total needed: [bold]{total_needed_gb:.1f} GB[/]\n" f" GPU VRAM: [cyan]{gpu_mem_gb:.0f} GB[/] per device", title="[bold]GPU Calculator[/]", border_style="cyan", )) console.print(table) console.print( f"\n [bold green]Minimum GPUs: {min_gpus}[/]" f" ({min_gpus} x {gpu_mem_gb:.0f} GB = {min_gpus * gpu_mem_gb:.0f} GB)\n" ) console.print( "[dim]Note: fewer GPUs = faster (pipeline parallel has cross-device overhead).\n" "Estimates are conservative. Actual memory may vary with sequence length\n" "and model architecture. See 'obliteratus obliterate --help' for runtime options.[/]\n" ) def _estimate_total_params_b(config) -> float: """Estimate total parameter count in billions from a HuggingFace config.""" # Some configs have explicit param counts for attr in ("num_parameters", "n_params"): val = getattr(config, attr, None) if val and val > 1000: return val / 1e9 # Estimate from architecture dimensions h = getattr(config, "hidden_size", 0) L = getattr(config, "num_hidden_layers", 0) V = getattr(config, "vocab_size", 0) i = getattr(config, "intermediate_size", h * 4) if h == 0 or L == 0: console.print("[red]Cannot determine model size from config. Use --params.[/]") raise SystemExit(1) n_heads = getattr(config, "num_attention_heads", None) or (h // 128) head_dim = getattr(config, "head_dim", None) or (h // n_heads if n_heads else 128) kv_heads = getattr(config, "num_key_value_heads", None) or n_heads # Attention: Q + K + V projections + output projection attn_params = h * (n_heads * head_dim) + h * (kv_heads * head_dim) * 2 + (n_heads * head_dim) * h # FFN (MoE or dense) n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1 # MoE models often have a separate intermediate size for expert FFNs moe_i = getattr(config, "moe_intermediate_size", i) # gate + up + down projections per expert ffn_per_expert = h * moe_i * 3 ffn_params = ffn_per_expert * n_experts # Some architectures (Qwen, DeepSeek) also have a shared/dense FFN per layer if n_experts > 1 and hasattr(config, "moe_intermediate_size"): # The dense FFN uses the main intermediate_size ffn_params += h * i * 3 # Router if n_experts > 1: ffn_params += h * n_experts # Per-layer: attention + FFN + layernorms layer_params = attn_params + ffn_params + h * 4 # 2 layernorms, 2 params each # Embedding + LM head embed_params = V * h * 2 # input + output embeddings (may be tied but counts for memory) total = L * layer_params + embed_params return total / 1e9 def _estimate_active_params_b(config, total_params_b: float) -> float: """For MoE models, estimate active parameters per forward pass.""" n_experts = getattr(config, "num_local_experts", getattr(config, "num_experts", 1)) or 1 if n_experts <= 1: return total_params_b top_k = getattr(config, "num_experts_per_tok", getattr(config, "top_k", 2)) or 2 h = getattr(config, "hidden_size", 0) i = getattr(config, "intermediate_size", h * 4) moe_i = getattr(config, "moe_intermediate_size", i) L = getattr(config, "num_hidden_layers", 0) # FFN per expert (uses moe_intermediate_size if available) ffn_per_expert = h * moe_i * 3 # Active FFN = top_k experts instead of all n_experts ffn_all = ffn_per_expert * n_experts * L ffn_active = ffn_per_expert * top_k * L # Non-FFN params (includes any shared/dense FFN) non_ffn = total_params_b * 1e9 - ffn_all active = non_ffn + ffn_active return max(active / 1e9, 0.1) 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, "min_layer_fraction", None) is not None: kwargs["min_layer_fraction"] = args.min_layer_fraction if getattr(args, "max_layer_fraction", None) is not None: kwargs["max_layer_fraction"] = args.max_layer_fraction if getattr(args, "harmless_pc_count", None) is not None: kwargs["harmless_pc_count"] = args.harmless_pc_count if getattr(args, "shield_concept_count", None) is not None: kwargs["shield_concept_count"] = args.shield_concept_count if getattr(args, "shield_ridge", None) is not None: kwargs["shield_ridge"] = args.shield_ridge if getattr(args, "shield_residualize", None) is not None: kwargs["shield_residualize"] = args.shield_residualize if getattr(args, "shield_layer_penalty", None) is not None: kwargs["shield_layer_penalty"] = args.shield_layer_penalty if getattr(args, "projection_target", None) is not None: kwargs["projection_target"] = args.projection_target if getattr(args, "projection_row_fraction", None) is not None: kwargs["projection_row_fraction"] = args.projection_row_fraction 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()