{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "3ca18c70",
      "metadata": {},
      "source": [
        "# FeatureLens offline study\n",
        "\n",
        "This notebook runs the full FeatureLens empirical study on a CUDA runtime while persisting experiment artifacts to Google Drive. It is designed to be resumable after Colab disconnects.\n",
        "\n",
        "**Before running:** choose a GPU runtime in Colab, then execute the cells from top to bottom.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dd012fdd",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 1. Verify that Colab actually assigned a GPU.\n",
        "import subprocess, sys\n",
        "\n",
        "subprocess.run([\"nvidia-smi\"], check=True)\n",
        "\n",
        "try:\n",
        "    import torch\n",
        "    assert torch.cuda.is_available(), \"CUDA is not available. Change the Colab runtime to a GPU and reconnect.\"\n",
        "    props = torch.cuda.get_device_properties(0)\n",
        "    gpu_name = torch.cuda.get_device_name(0)\n",
        "    gpu_vram_gb = props.total_memory / 1024**3\n",
        "    print(f\"\\nGPU: {gpu_name} | VRAM: {gpu_vram_gb:.1f} GB\")\n",
        "except Exception as exc:\n",
        "    raise RuntimeError(\"A CUDA GPU runtime is required for the model stages.\") from exc\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bdac04b1",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 2. Mount Google Drive so completed experiment stages survive a runtime reset.\n",
        "from google.colab import drive\n",
        "drive.mount(\"/content/drive\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "0446d7cd",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 3. Configuration \u2014 edit REPO_URL before running this cell.\n",
        "from pathlib import Path\n",
        "\n",
        "REPO_URL = \"PASTE_YOUR_GIT_REPO_URL_HERE\"\n",
        "BRANCH = \"main\"\n",
        "DRIVE_RUN_NAME = \"FeatureLens_offline_v016_full\"\n",
        "\n",
        "REPO_DIR = Path(\"/content/FeatureLens\")\n",
        "DRIVE_ROOT = Path(\"/content/drive/MyDrive\") / DRIVE_RUN_NAME\n",
        "DRIVE_ARTIFACTS = DRIVE_ROOT / \"artifacts\"\n",
        "LOG_PATH = DRIVE_ROOT / \"offline_study.log\"\n",
        "\n",
        "if REPO_URL.startswith(\"PASTE_\"):\n",
        "    raise ValueError(\"Set REPO_URL to your FeatureLens Git repository URL first.\")\n",
        "\n",
        "DRIVE_ARTIFACTS.mkdir(parents=True, exist_ok=True)\n",
        "print(\"Persistent run directory:\", DRIVE_ROOT)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c305468e",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 4. Clone or refresh the FeatureLens source.\n",
        "import shutil, subprocess\n",
        "\n",
        "if not REPO_DIR.exists():\n",
        "    subprocess.run([\"git\", \"clone\", \"--branch\", BRANCH, \"--single-branch\", REPO_URL, str(REPO_DIR)], check=True)\n",
        "else:\n",
        "    subprocess.run([\"git\", \"-C\", str(REPO_DIR), \"fetch\", \"origin\", BRANCH], check=True)\n",
        "    subprocess.run([\"git\", \"-C\", str(REPO_DIR), \"checkout\", BRANCH], check=True)\n",
        "    subprocess.run([\"git\", \"-C\", str(REPO_DIR), \"pull\", \"--ff-only\", \"origin\", BRANCH], check=True)\n",
        "\n",
        "print(subprocess.check_output([\"git\", \"-C\", str(REPO_DIR), \"rev-parse\", \"--short\", \"HEAD\"], text=True).strip())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e34c567a",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 5. Install the project environment. This can take a few minutes on a fresh runtime.\n",
        "import subprocess, sys\n",
        "subprocess.run([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", \"-r\", str(REPO_DIR / \"requirements.txt\")], check=True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "501273c0",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 6. Re-check CUDA after dependency installation and choose a conservative activation batch.\n",
        "import os, torch\n",
        "\n",
        "assert torch.cuda.is_available(), \"CUDA disappeared after dependency setup.\"\n",
        "gpu_name = torch.cuda.get_device_name(0)\n",
        "gpu_vram_gb = torch.cuda.get_device_properties(0).total_memory / 1024**3\n",
        "ACTIVATION_BATCH_SIZE = 16 if gpu_vram_gb >= 20 else 8\n",
        "ACTIVATION_MAX_LENGTH = 192\n",
        "\n",
        "# Keep model/SAE downloads on Colab's local disk for speed.\n",
        "os.environ[\"HF_HOME\"] = \"/content/hf_cache\"\n",
        "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
        "\n",
        "print(f\"GPU: {gpu_name} ({gpu_vram_gb:.1f} GB)\")\n",
        "print(f\"Activation batch size: {ACTIVATION_BATCH_SIZE}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "41a3a0e8",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 7. Link FeatureLens artifacts to Google Drive.\n",
        "# Existing small repo artifacts (for example README.md) are copied once; the local directory is then replaced by a symlink.\n",
        "import shutil\n",
        "\n",
        "local_artifacts = REPO_DIR / \"artifacts\"\n",
        "if local_artifacts.is_symlink():\n",
        "    local_artifacts.unlink()\n",
        "elif local_artifacts.exists():\n",
        "    shutil.copytree(local_artifacts, DRIVE_ARTIFACTS, dirs_exist_ok=True)\n",
        "    shutil.rmtree(local_artifacts)\n",
        "\n",
        "local_artifacts.symlink_to(DRIVE_ARTIFACTS, target_is_directory=True)\n",
        "print(\"artifacts ->\", local_artifacts.resolve())\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dbfda9e4",
      "metadata": {},
      "source": [
        "## Run / resume the study\n",
        "\n",
        "The command below is safe to rerun. Completed stages are skipped. The causal and feature-set stages also checkpoint completed tasks, so a disconnect during either stage does not discard earlier tasks from that stage.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3f0bfd41",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 8. Run the full pipeline with live output and a persistent log.\n",
        "import subprocess, sys, time\n",
        "\n",
        "command = [\n",
        "    sys.executable, \"-m\", \"experiments.run_all\",\n",
        "    \"--resume\",\n",
        "    \"--activation-batch-size\", str(ACTIVATION_BATCH_SIZE),\n",
        "    \"--activation-max-length\", str(ACTIVATION_MAX_LENGTH),\n",
        "]\n",
        "\n",
        "print(\"$\", \" \".join(command))\n",
        "print(\"Log:\", LOG_PATH)\n",
        "start = time.time()\n",
        "\n",
        "with LOG_PATH.open(\"a\", encoding=\"utf-8\") as log:\n",
        "    log.write(\"\\n\\n=== FeatureLens run ===\\n\")\n",
        "    log.write(\"$ \" + \" \".join(command) + \"\\n\")\n",
        "    process = subprocess.Popen(\n",
        "        command, cwd=REPO_DIR, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1\n",
        "    )\n",
        "    assert process.stdout is not None\n",
        "    for line in process.stdout:\n",
        "        print(line, end=\"\")\n",
        "        log.write(line)\n",
        "        log.flush()\n",
        "    return_code = process.wait()\n",
        "\n",
        "if return_code != 0:\n",
        "    raise RuntimeError(\n",
        "        f\"Pipeline exited with code {return_code}. Fix the error, then rerun this cell; --resume will keep completed work.\"\n",
        "    )\n",
        "\n",
        "print(f\"\\nCompleted in {(time.time() - start) / 60:.1f} minutes.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "645d7ee5",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 9. Validate the measured artifact set.\n",
        "import subprocess, sys\n",
        "subprocess.run([sys.executable, \"-m\", \"scripts.validate_artifacts\"], cwd=REPO_DIR, check=True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9e823e4d",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 10. Inspect the study summary and report.\n",
        "from pathlib import Path\n",
        "import json, pandas as pd\n",
        "from IPython.display import display, Markdown\n",
        "\n",
        "summary_path = DRIVE_ARTIFACTS / \"study_summary.json\"\n",
        "study_table_path = DRIVE_ARTIFACTS / \"study_feature_summary.csv\"\n",
        "report_path = DRIVE_ARTIFACTS / \"report.md\"\n",
        "\n",
        "summary = json.loads(summary_path.read_text(encoding=\"utf-8\"))\n",
        "display(summary)\n",
        "display(pd.read_csv(study_table_path))\n",
        "display(Markdown(report_path.read_text(encoding=\"utf-8\")))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "62c96dc6",
      "metadata": {},
      "outputs": [],
      "source": [
        "# 11. Create a small publishable artifact bundle (activation caches and checkpoint markers are excluded).\n",
        "import zipfile\n",
        "\n",
        "PUBLISH_ZIP = DRIVE_ROOT / \"FeatureLens_offline_results.zip\"\n",
        "\n",
        "with zipfile.ZipFile(PUBLISH_ZIP, \"w\", compression=zipfile.ZIP_DEFLATED) as zf:\n",
        "    for path in sorted(DRIVE_ARTIFACTS.rglob(\"*\")):\n",
        "        if not path.is_file():\n",
        "            continue\n",
        "        rel = path.relative_to(DRIVE_ARTIFACTS)\n",
        "        if rel.parts and rel.parts[0] == \"activations\":\n",
        "            continue\n",
        "        if path.name.endswith(\".complete\") or path.name.endswith(\".tmp\"):\n",
        "            continue\n",
        "        zf.write(path, arcname=str(Path(\"artifacts\") / rel))\n",
        "\n",
        "print(\"Publishable bundle:\", PUBLISH_ZIP)\n",
        "print(f\"Size: {PUBLISH_ZIP.stat().st_size / 1024**2:.2f} MiB\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a6634bfe",
      "metadata": {},
      "source": [
        "## After Colab\n",
        "\n",
        "Download `FeatureLens_offline_results.zip` from the Drive run folder. Extract it over your local FeatureLens repository so the files land under `artifacts/`, run the normal release checks locally, inspect the measured report, and only then commit the small study artifacts. Do **not** commit `artifacts/activations/`.\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "FeatureLens Offline Study",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}