OBLITERATUS/notebooks/abliterate.ipynb

{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "header"
   },
   "source": [
    "# OBLITERATUS — One-Click Abliteration\n",
    "\n",
    "**SOTA refusal removal** running on free Colab GPU. SVD multi-direction extraction, norm-preserving projection, iterative refinement.\n",
    "\n",
    "Based on: Arditi et al. (2024) | Gabliteration (arXiv:2512.18901) | grimjim norm-preserving biprojection (2025)\n",
    "\n",
    "---\n",
    "\n",
    "**How to use:**\n",
    "1. Make sure GPU runtime is enabled: `Runtime > Change runtime type > T4 GPU`\n",
    "2. Set your model and method in the config cell below\n",
    "3. For gated models, complete the Hugging Face access setup in section 2 before running.\n",
    "4. Run All (`Runtime > Run all` or `Ctrl+F9`)\n",
    "5. Download the abliterated model from the output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "setup-header"
   },
   "source": [
    "## 1. Install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "install"
   },
   "outputs": [],
   "source": "!pip install -q git+https://github.com/elder-plinius/OBLITERATUS.git\n!pip install -q accelerate bitsandbytes\n\nimport torch\nprint(f\"PyTorch {torch.__version__}\")\nprint(f\"CUDA available: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "config-header"
   },
   "source": [
    "## 2. Configure and prepare model access\n",
    "\n",
    "Edit the cell below to set your target model and abliteration method.\n",
    "\n",
    "**The default Llama model is gated.** Open its [model page](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct), accept the license/request access, and wait for approval on your Hugging Face account. A token alone does not grant access.\n",
    "\n",
    "Create a **read** token for that account (or a fine-grained token permitted to read the selected model). In Colab, open **Secrets** (the key icon), add a secret named `HF_TOKEN`, and enable **Notebook access**. Never paste a token into a code cell or saved output. Outside Colab, an existing `hf auth login` session or an `HF_TOKEN` environment variable also works.\n",
    "\n",
    "Ungated public models such as `Qwen/Qwen2.5-7B-Instruct` or `openai-community/gpt2` need no login. Step 3 checks access to the selected model's small `config.json` before constructing the pipeline or downloading weights; if access cannot be confirmed, it stops with setup instructions. After fixing access or changing the model, rerun the configuration and step 3. The check runs again even when step 3 is executed directly.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "config"
   },
   "outputs": [],
   "source": [
    "#@title Abliteration Config { run: \"auto\" }\n",
    "\n",
    "#@markdown ### Target Model\n",
    "#@markdown Pick a model from the dropdown or paste a custom HuggingFace ID.\n",
    "MODEL = \"meta-llama/Llama-3.1-8B-Instruct\" #@param [\"meta-llama/Llama-3.1-8B-Instruct\", \"Qwen/Qwen2.5-7B-Instruct\", \"mistralai/Mistral-7B-Instruct-v0.3\", \"google/gemma-2-9b-it\", \"microsoft/Phi-3.5-mini-instruct\", \"THUDM/glm-4-9b-chat\", \"NousResearch/Hermes-3-Llama-3.1-8B\", \"cognitivecomputations/dolphin-2.9.4-llama3.1-8b\", \"TinyLlama/TinyLlama-1.1B-Chat-v1.0\", \"openai-community/gpt2\"] {allow-input: true}\n",
    "\n",
    "#@markdown ### Method\n",
    "METHOD = \"advanced\" #@param [\"basic\", \"advanced\", \"aggressive\"]\n",
    "\n",
    "#@markdown ### Advanced Overrides (leave 0 to use method defaults)\n",
    "N_DIRECTIONS = 0 #@param {type: \"integer\"}\n",
    "REGULARIZATION = 0.0 #@param {type: \"number\"}\n",
    "REFINEMENT_PASSES = 0 #@param {type: \"integer\"}\n",
    "\n",
    "#@markdown ### Output\n",
    "OUTPUT_DIR = \"abliterated\" #@param {type: \"string\"}\n",
    "\n",
    "print(f\"Model: {MODEL}\")\n",
    "print(f\"Method: {METHOD}\")\n",
    "print(f\"Output: {OUTPUT_DIR}/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "run-header"
   },
   "source": [
    "## 3. Run Abliteration Pipeline\n",
    "\n",
    "This runs all 6 stages: SUMMON → PROBE → DISTILL → EXCISE → VERIFY → REBIRTH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "run-pipeline"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from huggingface_hub import get_token, hf_hub_download\n",
    "from obliteratus.credential_sources import resolve_secret\n",
    "\n",
    "\n",
    "def check_model_access(model_name):\n",
    "    # Resolve the same configured sources as the loader, then Hub/Colab login.\n",
    "    # Only the small config is fetched; force a live check even if it is cached.\n",
    "    try:\n",
    "        token = resolve_secret(\"HF_TOKEN\") or get_token()\n",
    "        hf_hub_download(\n",
    "            repo_id=model_name, filename=\"config.json\",\n",
    "            token=token or False, force_download=True,\n",
    "        )\n",
    "    except Exception:\n",
    "        # Hub/credential exceptions may contain sensitive request details.\n",
    "        raise RuntimeError(\n",
    "            \"Model access could not be confirmed. Check the model ID and your \"\n",
    "            \"connection. For gated/private models, accept the license/request \"\n",
    "            \"access on the model page and wait for approval; use a read token \"\n",
    "            \"from that account. In Colab, add HF_TOKEN in Secrets and enable \"\n",
    "            \"Notebook access (outside Colab use hf auth login or HF_TOKEN). \"\n",
    "            \"Then rerun this cell, or choose an ungated public model. \"\n",
    "            \"No model weights have been loaded.\"\n",
    "        ) from None\n",
    "    if token:\n",
    "        # Runtime only: keep the loader on the credential that passed the check.\n",
    "        os.environ[\"HF_TOKEN\"] = token\n",
    "\n",
    "\n",
    "check_model_access(MODEL)\n",
    "\n",
    "from obliteratus.abliterate import AbliterationPipeline\n",
    "\n",
    "# Build kwargs, only pass overrides if non-zero\n",
    "kwargs = dict(\n",
    "    model_name=MODEL,\n",
    "    output_dir=OUTPUT_DIR,\n",
    "    device=\"auto\",\n",
    "    dtype=\"float16\",\n",
    "    method=METHOD,\n",
    ")\n",
    "if N_DIRECTIONS > 0:\n",
    "    kwargs[\"n_directions\"] = N_DIRECTIONS\n",
    "if REGULARIZATION > 0:\n",
    "    kwargs[\"regularization\"] = REGULARIZATION\n",
    "if REFINEMENT_PASSES > 0:\n",
    "    kwargs[\"refinement_passes\"] = REFINEMENT_PASSES\n",
    "\n",
    "# Progress callback\n",
    "def on_stage(stage):\n",
    "    icons = {\"summon\": \"\\u26a1\", \"probe\": \"\\u2692\", \"distill\": \"\\u269b\",\n",
    "             \"excise\": \"\\u2702\", \"verify\": \"\\u2713\", \"rebirth\": \"\\u2606\"}\n",
    "    icon = icons.get(stage.stage, \"\")\n",
    "    print(f\"\\n{'='*60}\")\n",
    "    print(f\"{icon}  STAGE: {stage.stage.upper()} — {stage.message}\")\n",
    "    print(f\"{'='*60}\")\n",
    "\n",
    "def on_log(msg):\n",
    "    print(f\"  {msg}\")\n",
    "\n",
    "kwargs[\"on_stage\"] = on_stage\n",
    "kwargs[\"on_log\"] = on_log\n",
    "\n",
    "pipeline = AbliterationPipeline(**kwargs)\n",
    "result = pipeline.run()\n",
    "\n",
    "print(f\"\\n{'='*60}\")\n",
    "print(f\"ABLITERATION COMPLETE\")\n",
    "print(f\"Output: {result}\")\n",
    "print(f\"{'='*60}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "download-header"
   },
   "source": [
    "## 4. Download the Abliterated Model\n",
    "\n",
    "Run the cell below to zip and download, or upload directly to HuggingFace Hub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "download"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "# Find the output directory\n",
    "out_path = Path(OUTPUT_DIR)\n",
    "subdirs = [d for d in out_path.iterdir() if d.is_dir()] if out_path.exists() else []\n",
    "model_dir = subdirs[0] if subdirs else out_path\n",
    "\n",
    "print(f\"Model saved at: {model_dir}\")\n",
    "print(f\"Contents:\")\n",
    "for f in sorted(model_dir.rglob(\"*\")):\n",
    "    if f.is_file():\n",
    "        size_mb = f.stat().st_size / 1024**2\n",
    "        print(f\"  {f.relative_to(model_dir)}  ({size_mb:.1f} MB)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zip-download"
   },
   "outputs": [],
   "source": [
    "#@title Option A: Download as ZIP\n",
    "import shutil\n",
    "from google.colab import files\n",
    "\n",
    "zip_name = model_dir.name\n",
    "shutil.make_archive(zip_name, 'zip', model_dir)\n",
    "print(f\"Downloading {zip_name}.zip ...\")\n",
    "files.download(f\"{zip_name}.zip\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "push-hub"
   },
   "outputs": [],
   "source": [
    "#@title Option B: Push to HuggingFace Hub (opt-in)\n",
    "#@markdown Enable only when ready to upload. Publishing requires a write token with permission for the destination repo; read access to the source model is separate.\n",
    "UPLOAD_TO_HUB = False #@param {type: \"boolean\"}\n",
    "HF_REPO = \"your-username/model-name-abliterated\" #@param {type: \"string\"}\n",
    "\n",
    "if UPLOAD_TO_HUB:\n",
    "    from huggingface_hub import HfApi, get_token\n",
    "    from obliteratus.credential_sources import resolve_secret\n",
    "\n",
    "    if HF_REPO == \"your-username/model-name-abliterated\":\n",
    "        raise ValueError(\"Set HF_REPO to your destination repository before uploading.\")\n",
    "    upload_token = resolve_secret(\"HF_TOKEN\") or get_token()\n",
    "    if not upload_token:\n",
    "        raise RuntimeError(\"Set a write token in HF_TOKEN Secrets before uploading.\")\n",
    "    api = HfApi(token=upload_token)\n",
    "    api.create_repo(HF_REPO, exist_ok=True)\n",
    "    api.upload_folder(\n",
    "        folder_path=str(model_dir),\n",
    "        repo_id=HF_REPO,\n",
    "        repo_type=\"model\",\n",
    "    )\n",
    "    print(f\"\\nUploaded to: https://huggingface.co/{HF_REPO}\")\n",
    "else:\n",
    "    print(\"Hub upload skipped. Enable UPLOAD_TO_HUB to publish explicitly.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "test-header"
   },
   "source": [
    "## 5. Test the Abliterated Model\n",
    "\n",
    "Quick sanity check — try some prompts that the original model would refuse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "test-model"
   },
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "print(f\"Loading abliterated model from {model_dir}...\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(str(model_dir))\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    str(model_dir),\n",
    "    torch_dtype=torch.float16,\n",
    "    device_map=\"auto\",\n",
    ")\n",
    "\n",
    "def generate(prompt, max_new_tokens=100):\n",
    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
    "    with torch.no_grad():\n",
    "        outputs = model.generate(\n",
    "            **inputs,\n",
    "            max_new_tokens=max_new_tokens,\n",
    "            do_sample=True,\n",
    "            temperature=0.7,\n",
    "            top_p=0.9,\n",
    "        )\n",
    "    return tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "\n",
    "# Test prompts\n",
    "test_prompts = [\n",
    "    \"The capital of France is\",\n",
    "    \"Explain how a neural network works:\",\n",
    "    \"Write a poem about the ocean:\",\n",
    "]\n",
    "\n",
    "for prompt in test_prompts:\n",
    "    print(f\"\\n{'='*60}\")\n",
    "    print(f\"PROMPT: {prompt}\")\n",
    "    print(f\"{'='*60}\")\n",
    "    print(generate(prompt))"
   ]
  }
 ]
}