Fix Exploratory Analysis Demo (#429)

Make it work with the latest transformer_lens version, add testing & deploy to the docs pages.

.github/workflows/checks.yml

@@ -46,13 +46,19 @@ jobs:
       - uses: actions/checkout@v3
       - name: Install Poetry
         uses: snok/install-poetry@v1
-        with:
-          version: 1.4.0
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
           cache: "poetry"
+      - name: Cache Models used with Tests
+        uses: actions/cache@v3
+        with:
+          path: |
+            ~/.cache/huggingface/hub/models--gpt2
+            ~/.cache/huggingface/hub/models--NeelNanda--Attn_Only*
+            ~/.cache/huggingface/hub/models--roneneldan--TinyStories-1M*
+          key: ${{ runner.os }}-huggingface-models
       - name: Install dependencies
         run: |
           poetry lock --check
@@ -69,3 +75,32 @@ jobs:
       #   run: poetry run mypy transformer_lens
       - name: Build check
         run: poetry build
+  notebook-checks:
+    name: Notebook Checks
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.11"
+          cache: "poetry"
+      - name: Re-use HuggingFace models cache
+        uses: actions/cache/restore@v3
+        with:
+          path: ~/.cache/huggingface/hub
+          key: ${{ runner.os }}-huggingface-models
+      - name: Install dependencies
+        run: |
+          poetry lock --check
+          poetry install --with dev,jupyter
+      - name: Install pandoc
+        uses: awalsh128/cache-apt-pkgs-action@latest
+        with:
+          packages: pandoc
+          version: 1.0
+      - name: Check Notebook Output Consistency
+        # Note: currently only checks notebooks we have specifically setup for this
+        run: make notebook-test

.github/workflows/gh-pages.yml

@@ -24,6 +24,11 @@ jobs:
         uses: actions/setup-python@v2
         with:
           python-version: "3.11"
+      - name: Install pandoc
+        uses: awalsh128/cache-apt-pkgs-action@latest
+        with:
+          packages: pandoc
+          version: 1.0
       - name: Install dependencies
         run: poetry install --with docs
       - name: Build Docs

.vscode/extensions.json

@@ -19,5 +19,6 @@
         "stkb.rewrap",
         "streetsidesoftware.code-spell-checker",
         "yzhang.markdown-all-in-one",
+        "streetsidesoftware.code-spell-checker-british-english"
     ]
 }

.vscode/settings.json

@@ -17,12 +17,23 @@
     "pylint.importStrategy": "fromEnvironment",
     "notebook.formatOnCellExecution": true,
     "notebook.formatOnSave.enabled": true,
+    "cSpell.language": "en,en-GB",
     "cSpell.words": [
+        "arange",
+        "argmax",
         "autoregressive",
+        "caxis",
         "checkpointed",
+        "circuitsvis",
+        "Colab",
+        "colour",
+        "Cooney's",
+        "disconfirm",
         "docstrings",
         "doctest",
         "doctree",
+        "dtype",
+        "einops",
         "Elhage",
         "evals",
         "firstpage",
@@ -32,11 +43,16 @@
         "gptj",
         "howpublished",
         "huggingface",
+        "imshow",
+        "interp",
         "interpretability",
         "isort",
+        "jaxtyping",
         "keepdim",
         "lastpage",
         "layernorm",
+        "logits",
+        "logsumexp",
         "maxdepth",
         "mingpt",
         "Nanda",
@@ -45,16 +61,24 @@
         "Nitpicky",
         "Olah",
         "pagename",
+        "pretrained",
         "probs",
         "pycln",
         "pytest",
+        "resid",
         "Softmax",
         "solu",
         "templatedir",
         "templatename",
         "toctree",
+        "topk",
         "transformerlens",
         "Unembed",
-        "unembedding"
+        "unembedding",
+        "unigram",
+        "visualisation",
+        "xaxis",
+        "yaxis"
     ],
+    "cSpell.caseSensitive": false,
 }