load_polysemantic_benchmark.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from poly_bench.io import load_from_hf\n",
    "from poly_bench.poly_hl_model import PolyModelDataset\n",
    "\n",
    "n_samples = 10_000\n",
    "n_ctx = 15"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Case 0 + 1 + 2 + 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4722d8ad21b4726946288508a2c6243",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/162k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ll_model, hl_model, dset_classes = load_from_hf(model_name=\"cases_0+1+2+3\", benchmarks_file='cases.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:14: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.inputs = t.tensor(inputs).to(t.int)\n",
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:15: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.targets = t.tensor(targets).to(t.float32)\n",
      "/Users/evananders/far_cluster/polysemantic-benchmark/poly_bench/utils.py:19: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  self.markers = t.tensor(markers).to(t.int)\n"
     ]
    }
   ],
   "source": [
    "dsets = [dsetcase(N_samples=n_samples, n_ctx=n_ctx, seed=42) for dsetcase in dset_classes]\n",
    "dataset = PolyModelDataset(dsets, n_ctx=hl_model.cfg.n_ctx).get_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/ry/qny1f95136l2lpppg_d78n6c0000gq/T/ipykernel_50893/22537286.py:1: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  inputs = torch.tensor(dataset.inputs)\n",
      "/var/folders/ry/qny1f95136l2lpppg_d78n6c0000gq/T/ipykernel_50893/22537286.py:2: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  labels = torch.tensor(dataset.targets)\n"
     ]
    }
   ],
   "source": [
    "inputs = torch.tensor(dataset.inputs)\n",
    "labels = torch.tensor(dataset.targets)\n",
    "ll_output = ll_model(inputs)\n",
    "hl_output = hl_model((inputs, None, None))\n",
    "\n",
    "assert torch.allclose(hl_output, labels.to(hl_output.device))\n",
    "assert (torch.argmax(hl_output, dim=-1) -  torch.argmax(ll_output, dim=-1) != 0).sum() < 1e-4 * n_samples * n_ctx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "circuits_bench",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}