{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7382b6c",
   "metadata": {},
   "source": [
    "# Jacobian lens — walkthrough\n",
    "\n",
    "Load a model, load a pre-fitted Jacobian lens from the Hub, apply it to a prompt, and render the interactive slice visualisation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5771294e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jlens\n",
    "\n",
    "jlens.configure_logging()\n",
    "\n",
    "MODEL_NAME = \"Qwen/Qwen3.5-4B\"\n",
    "# MODEL_NAME = \"Qwen/Qwen3.6-27B\"\n",
    "\n",
    "LENS_REPO = \"neuronpedia/jacobian-lens\"\n",
    "LENS_REVISION = \"qwen-n1000\"\n",
    "LENS_FILE = {\n",
    "    \"Qwen/Qwen3.5-4B\": \"qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt\",\n",
    "    \"Qwen/Qwen3.6-27B\": \"qwen3.6-27b/jlens/Salesforce-wikitext/Qwen3.6-27B_jacobian_lens_n1000.pt\",\n",
    "}[MODEL_NAME]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "975fa03c",
   "metadata": {},
   "source": [
    "## 1. Load the model\n",
    "\n",
    "`jlens.from_hf` wraps an already-loaded HuggingFace model into `LensModel` interface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a25cd5cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformers\n",
    "\n",
    "hf_model = transformers.AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_NAME, dtype=torch.bfloat16\n",
    ").cuda()\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "model = jlens.from_hf(hf_model, tokenizer)\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba6e94e0",
   "metadata": {},
   "source": [
    "## 2. Load a pre-fitted lens\n",
    "\n",
    "`JacobianLens.from_pretrained` pulls a `.pt` from the Hub (or a local path). The lens holds one `[d_model, d_model]` matrix per layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f74b29c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "lens = jlens.JacobianLens.from_pretrained(\n",
    "    LENS_REPO, filename=LENS_FILE, revision=LENS_REVISION\n",
    ")\n",
    "lens"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da54c19c",
   "metadata": {},
   "source": [
    "## 3. Apply: J-lens vs logit lens\n",
    "\n",
    "`lens.apply(model, prompt, positions=...)` runs one forward pass, transports each layer's residual into the final-layer basis with `J_l`, and decodes through the model's own unembedding. `use_jacobian=False` skips the transport — that's the vanilla logit lens.\n",
    "\n",
    "Below: a two-hop factual question, read out at the boot token. The J-lens surfaces interpretable tokens at layers where the logit lens is still noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2679d83e",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"Fact: The currency used in the country shaped like a boot is\"\n",
    "layers = [\n",
    "    model.n_layers // 4,\n",
    "    model.n_layers // 2,\n",
    "    model.n_layers // 4 * 3,\n",
    "    model.n_layers - 2,\n",
    "]\n",
    "\n",
    "jlens_logits, model_logits, _ = lens.apply(model, prompt, layers=layers, positions=[-2])\n",
    "logit_lens, _, _ = lens.apply(\n",
    "    model, prompt, layers=layers, positions=[-2], use_jacobian=False\n",
    ")\n",
    "\n",
    "\n",
    "def top5(logits):\n",
    "    return [tokenizer.decode([t]) for t in logits.topk(5).indices]\n",
    "\n",
    "\n",
    "for layer in layers:\n",
    "    print(f\"L{layer:>3} logit-lens: {top5(logit_lens[layer][0])}\")\n",
    "    print(f\"L{layer:>3} J-lens:     {top5(jlens_logits[layer][0])}\")\n",
    "print(f\"model:           {top5(model_logits[0])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90026704",
   "metadata": {},
   "source": [
    "## 4. Render a slice page (inline)\n",
    "\n",
    "`compute_slice` + `build_page` produce an interactive position × layer view of the lens's token ranks (the `?` in the corner explains the controls). `mode=\"embed\"` inlines everything so the page is self-contained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c84b2d42",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import json\n",
    "\n",
    "from jlens.examples import EXAMPLES, resolve_prompt\n",
    "from jlens.vis import build_page, compute_slice, notebook_iframe\n",
    "\n",
    "# English gloss for Qwen's Chinese/Japanese/Korean vocab tokens (machine-\n",
    "# generated, best-effort), shown next to\n",
    "# the token in the page (alt_token=).\n",
    "gloss = {\n",
    "    int(k): v for k, v in json.load(gzip.open(\"assets/qwen_gloss.json.gz\")).items()\n",
    "}\n",
    "\n",
    "example = next(e for e in EXAMPLES if e.slug == \"multihop\")\n",
    "prompt = resolve_prompt(example, tokenizer)\n",
    "\n",
    "slice_data = compute_slice(\n",
    "    model,\n",
    "    lens,\n",
    "    prompt,\n",
    "    layer_stride=2,\n",
    "    # Empirically on Qwen, the interesting word tokens trail punctuation and\n",
    "    # single-character tokens in the raw top-K; mask to word-like tokens only.\n",
    "    mask_display=True,\n",
    ")\n",
    "page, _, _ = build_page(\n",
    "    slice_data,\n",
    "    prompt,\n",
    "    title=example.section,\n",
    "    description=example.description,\n",
    "    alt_token=gloss,\n",
    ")\n",
    "notebook_iframe(page)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15e9fd00",
   "metadata": {},
   "source": [
    "## 5. Render a slice page (served)\n",
    "\n",
    "For longer prompts prefer `mode=\"fetch\"`: `build_page` writes the data as sidecar files to `out_dir` and the page fetches rank files lazily on pin, so it stays small regardless of how many tokens are tracked."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64318be2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import threading\n",
    "from functools import partial\n",
    "from http.server import HTTPServer, SimpleHTTPRequestHandler\n",
    "from pathlib import Path\n",
    "\n",
    "example = next(e for e in EXAMPLES if e.slug == \"ascii-face\")\n",
    "prompt = resolve_prompt(example, tokenizer)\n",
    "\n",
    "slice_data = compute_slice(model, lens, prompt, mask_display=True)\n",
    "out_dir = Path(\"slices\") / example.slug\n",
    "page, _, _ = build_page(\n",
    "    slice_data,\n",
    "    prompt,\n",
    "    title=example.section,\n",
    "    description=example.description,\n",
    "    alt_token=gloss,\n",
    "    mode=\"fetch\",\n",
    "    out_dir=out_dir,\n",
    ")\n",
    "(out_dir / \"index.html\").write_text(page)\n",
    "\n",
    "if \"_jlens_httpd\" not in globals():\n",
    "    _handler = partial(SimpleHTTPRequestHandler, directory=os.path.abspath(\"slices\"))\n",
    "    _jlens_httpd = HTTPServer((\"127.0.0.1\", 0), _handler)\n",
    "    threading.Thread(target=_jlens_httpd.serve_forever, daemon=True).start()\n",
    "print(f\"-> http://localhost:{_jlens_httpd.server_address[1]}/{example.slug}/\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "604e6c25",
   "metadata": {},
   "source": [
    "### More to explore\n",
    "\n",
    "A few more prompts are bundled in `jlens.examples.EXAMPLES` — change the `slug` above and see what surfaces, or try a prompt of your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d98765a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for e in EXAMPLES:\n",
    "    print(f\"{e.slug:>24}  {e.section}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ef45078",
   "metadata": {},
   "source": [
    "## 6. Fitting\n",
    "\n",
    "`fit(model, prompts)` computes `J_l` over the supplied prompts. 100 prompts is enough for a usable lens; the released lenses use 1000. `dim_batch` is the memory knob — each prompt does `ceil(d_model / dim_batch)` backward passes on a retained graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fb27b941602401d91542211134fc71a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jlens.examples import load_wikitext_prompts\n",
    "\n",
    "prompts = load_wikitext_prompts(n_prompts=100)\n",
    "lens = jlens.fit(\n",
    "    model, prompts, dim_batch=32, max_seq_len=128, checkpoint_path=\"ckpt.pt\"\n",
    ")\n",
    "lens.save(\"jacobian_lens.pt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv (3.13.14.final.0)",
   "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.13.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}