update loading in bert demo to use transformer bridge

.github/workflows/checks.yml

@@ -225,7 +225,7 @@ jobs:
       matrix:
         notebook:
           # - "Activation_Patching_in_TL_Demo"
-          # - "Attribution_Patching_Demo"
+          - "Attribution_Patching_Demo"
           - "ARENA_Content"
           - "BERT"
           - "Exploratory_Analysis_Demo"

demos/BERT.ipynb

@@ -148,7 +148,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -157,7 +157,7 @@
     "\n",
     "from transformers import AutoTokenizer\n",
     "\n",
-    "from transformer_lens import HookedEncoder, BertNextSentencePrediction"
+    "from transformer_lens.model_bridge import TransformerBridge"
    ]
   },
   {
@@ -192,7 +192,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -214,7 +214,8 @@
    ],
    "source": [
     "# NBVAL_IGNORE_OUTPUT\n",
-    "bert = HookedEncoder.from_pretrained(\"bert-base-cased\")\n",
+    "bert = TransformerBridge.boot_transformers(\"bert-base-cased\")\n",
+    "bert.enable_compatibility_mode()\n",
     "tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")"
    ]
   },
@@ -287,14 +288,13 @@
    "metadata": {},
    "source": [
     "## Next Sentence Prediction\n",
-    "To carry out Next Sentence Prediction, you have to use the class BertNextSentencePrediction, and pass a HookedEncoder in its constructor.  \n",
-    "Then, create a list with the two sentences you want to perform NSP on as elements and use that as input to the forward function.  \n",
+    "To carry out Next Sentence Prediction create a list with the two sentences you want to perform NSP on as elements and use that as input to the forward function.  \n",
     "The model will then predict the probability of the sentence at position 1 following (i.e. being the next sentence) to the sentence at position 0."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -308,13 +308,12 @@
     }
    ],
    "source": [
-    "nsp = BertNextSentencePrediction(bert)\n",
     "sentence_a = \"A man walked into a grocery store.\"\n",
     "sentence_b = \"He bought an apple.\"\n",
     "\n",
     "input = [sentence_a, sentence_b]\n",
     "\n",
-    "predictions = nsp(input, return_type=\"predictions\")\n",
+    "predictions = bert(input, return_type=\"predictions\")\n",
     "\n",
     "print(f\"Sentence A: {sentence_a}\")\n",
     "print(f\"Sentence B: {sentence_b}\")\n",