{
 "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. Run All (`Runtime > Run all` or `Ctrl+F9`)\n",
    "4. 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/obliteratus-project/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_mem / 1024**3:.1f} GB\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "config-header"
   },
   "source": [
    "## 2. Configure\n",
    "\n",
    "Edit the cell below to set your target model and abliteration method."
   ]
  },
  {
   "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": [
    "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.key, \"\")\n",
    "    print(f\"\\n{'='*60}\")\n",
    "    print(f\"{icon}  STAGE: {stage.key.upper()} — {stage.description}\")\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.get('output_dir', OUTPUT_DIR)}\")\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\n",
    "#@markdown Set your HF repo name. You'll need to be logged in (`huggingface-cli login`).\n",
    "HF_REPO = \"your-username/model-name-abliterated\" #@param {type: \"string\"}\n",
    "\n",
    "from huggingface_hub import HfApi\n",
    "api = HfApi()\n",
    "\n",
    "# Login if needed\n",
    "from huggingface_hub import notebook_login\n",
    "notebook_login()\n",
    "\n",
    "# Upload\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}\")"
   ]
  },
  {
   "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))"
   ]
  }
 ]
}