updating loading in head detector demo to use transformer bridge

Author: degenfabian

.github/workflows/checks.yml

@@ -151,7 +151,7 @@ jobs:
           - "BERT"
           - "Exploratory_Analysis_Demo"
           # - "Grokking_Demo"
-          # - "Head_Detector_Demo"
+          - "Head_Detector_Demo"
           - "Interactive_Neuroscope"
           # - "LLaMA"
           # - "LLaMA2_GPU_Quantized"

demos/Head_Detector_Demo.ipynb

@@ -347,7 +347,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 3,
+      "execution_count": null,
       "metadata": {
         "id": "ScWILAgIGt5O"
       },
@@ -359,7 +359,8 @@
         "from tqdm import tqdm\n",
         "\n",
         "import transformer_lens\n",
-        "from transformer_lens import HookedTransformer, ActivationCache\n",
+        "from transformer_lens import ActivationCache\n",
+        "from transformer_lens.model_bridge import TransformerBridge\n",
         "from neel_plotly import line, imshow, scatter"
       ]
     },
@@ -479,7 +480,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": null,
       "metadata": {
         "id": "5ikyL8-S7u2Z"
       },
@@ -493,7 +494,6 @@
         "import numpy as np\n",
         "import torch\n",
         "\n",
-        "from transformer_lens import HookedTransformer, ActivationCache\n",
         "# from transformer_lens.utils import is_lower_triangular, is_square\n",
         "\n",
         "HeadName = Literal[\"previous_token_head\", \"duplicate_token_head\", \"induction_head\"]\n",
@@ -515,7 +515,7 @@
         "\n",
         "\n",
         "def detect_head(\n",
-        "    model: HookedTransformer,\n",
+        "    model: TransformerBridge,\n",
         "    seq: Union[str, List[str]],\n",
         "    detection_pattern: Union[torch.Tensor, HeadName],\n",
         "    heads: Optional[Union[List[LayerHeadTuple], LayerToHead]] = None,\n",
@@ -566,14 +566,15 @@
         "    --------\n",
         "    .. code-block:: python\n",
         "\n",
-        "        >>> from transformer_lens import HookedTransformer,  utils\n",
+        "        >>> from transformer_lens import utils\n",
+        "        >>> from transformer_lens.model_bridge import TransformerBridge\n",
         "        >>> from transformer_lens.head_detector import detect_head\n",
         "        >>> import plotly.express as px\n",
         "\n",
         "        >>> def imshow(tensor, renderer=None, xaxis=\"\", yaxis=\"\", **kwargs):\n",
         "        >>>     px.imshow(utils.to_numpy(tensor), color_continuous_midpoint=0.0, color_continuous_scale=\"RdBu\", labels={\"x\":xaxis, \"y\":yaxis}, **kwargs).show(renderer)\n",
         "\n",
-        "        >>> model = HookedTransformer.from_pretrained(\"gpt2-small\")\n",
+        "        >>> model = TransformerBridge.boot_transformers(\"gpt2\")\n",
         "        >>> sequence = \"This is a test sequence. This is a test sequence.\"\n",
         "\n",
         "        >>> attention_score = detect_head(model, sequence, \"previous_token_head\")\n",
@@ -777,7 +778,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": null,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/"
@@ -802,7 +803,7 @@
         }
       ],
       "source": [
-        "model = HookedTransformer.from_pretrained(\"gpt2-small\", device=device)"
+        "model = TransformerBridge.boot_transformers(\"gpt2\", device=device)"
       ]
     },
     {