updated loading in llama 2 demo to use transformer bridge

.github/workflows/checks.yml

@@ -154,7 +154,7 @@ jobs:
           # - "Head_Detector_Demo"
           # - "Interactive_Neuroscope"
           # - "LLaMA"
-          # - "LLaMA2_GPU_Quantized"
+          - "LLaMA2_GPU_Quantized"
           - "Main_Demo"
           # - "No_Position_Experiment"
           - "Othello_GPT"

demos/LLaMA2_GPU_Quantized.ipynb

@@ -212,7 +212,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {
     "id": "P8zS3MPkCUsR"
    },
@@ -232,7 +232,7 @@
     "from transformer_lens.hook_points import (\n",
     "    HookPoint,\n",
     ")  # Hooking utilities\n",
-    "from transformer_lens import HookedTransformer\n",
+    "from transformer_lens.model_bridge import TransformerBridge\n",
     "\n",
     "torch.set_grad_enabled(False)\n",
     "\n",
@@ -291,7 +291,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {
     "id": "RdJ0AuW_CUsS"
    },
@@ -303,7 +303,8 @@
     "    tokenizer = LlamaTokenizer.from_pretrained(MODEL_PATH)\n",
     "    hf_model = LlamaForCausalLM.from_pretrained(MODEL_PATH, low_cpu_mem_usage=True)\n",
     "\n",
-    "    model = HookedTransformer.from_pretrained(\"llama-7b\", hf_model=hf_model, device=\"cpu\", fold_ln=False, center_writing_weights=False, center_unembed=False, tokenizer=tokenizer)\n",
+    "    model = TransformerBridge.boot_transformers(\"llama-7b\", hf_model=hf_model, device=\"cpu\", fold_ln=False, center_writing_weights=False, center_unembed=False, tokenizer=tokenizer)\n",
+    "    model.enable_compatibility_mode()\n",
     "\n",
     "    model = model.to(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
     "    model.generate(\"The capital of Germany is\", max_new_tokens=20, temperature=0)"
@@ -406,7 +407,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
@@ -730,14 +731,15 @@
     "\n",
     "tokenizer = AutoTokenizer.from_pretrained(LLAMA_2_7B_CHAT_PATH)\n",
     "\n",
-    "model = HookedTransformer.from_pretrained(LLAMA_2_7B_CHAT_PATH,\n",
+    "model = TransformerBridge.boot_transformers(LLAMA_2_7B_CHAT_PATH,\n",
     "                                             hf_model=hf_model,\n",
     "                                             dtype=inference_dtype,\n",
     "                                             fold_ln=False,\n",
     "                                             fold_value_biases=False,\n",
     "                                             center_writing_weights=False,\n",
     "                                             center_unembed=False,\n",
     "                                             tokenizer=tokenizer)\n",
+    "model.enable_compatibility_mode()\n",
     "\n",
     "model.generate(\"The capital of Germany is\", max_new_tokens=2, temperature=0)\n",
     "\n"