Refactor components

.github/workflows/checks.yml

@@ -5,24 +5,24 @@ on:
     branches:
       - main
     paths:
-      - '**' # Include all files by default
-      - '!.devcontainer/**'
-      - '!.vscode/**'
-      - '!.git*'
-      - '!*.md'
-      - '!.github/**'
-      - '.github/workflows/checks.yml' # Still include current workflow
+      - "**" # Include all files by default
+      - "!.devcontainer/**"
+      - "!.vscode/**"
+      - "!.git*"
+      - "!*.md"
+      - "!.github/**"
+      - ".github/workflows/checks.yml" # Still include current workflow
   pull_request:
     branches:
       - main
     paths:
-      - '**'
-      - '!.devcontainer/**'
-      - '!.vscode/**'
-      - '!.git*'
-      - '!*.md'
-      - '!.github/**'
-      - '.github/workflows/checks.yml'
+      - "**"
+      - "!.devcontainer/**"
+      - "!.vscode/**"
+      - "!.git*"
+      - "!*.md"
+      - "!.github/**"
+      - ".github/workflows/checks.yml"
   # Allow this workflow to be called from other workflows
   workflow_call:
     inputs:
@@ -36,16 +36,15 @@ permissions:
   contents: write
 
 jobs:
-  checks:
-    name: Checks
+  compatibility-checks:
+    name: Compatibility Checks
     runs-on: ubuntu-latest
     strategy:
       matrix:
         python-version:
           - "3.8"
           - "3.9"
           - "3.10"
-          - "3.11"
     steps:
       - uses: actions/checkout@v3
       - name: Install Poetry
@@ -67,20 +66,20 @@ jobs:
         run: |
           poetry lock --check
           poetry install --with dev
-      - name: Check format
-        run: make check-format
-      - name: Unit test
+      - name: Unit Test
         run: make unit-test
-      - name: Docstring test
-        run: make docstring-test
-      # - name: Type check
-      #   run: poetry run mypy transformer_lens
+      - name: Acceptance Test
+        run: make acceptance-test
       - name: Build check
         run: poetry build
-
-  # Acceptance tests are run in parallel with unit checks.
-  acceptance-tests:
-    name: Acceptance Tests
+      - name: Upload Coverage Report Artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: documentation
+          path: htmlcov
+
+  code-checks:
+    name: Code Checks
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
@@ -103,8 +102,21 @@ jobs:
         run: |
           poetry lock --check
           poetry install --with dev
-      - name: Acceptance test
-        run: make acceptance-test
+      - name: Check format
+        run: make check-format
+      - name: Docstring test
+        run: make docstring-test
+      - name: Type check
+        run: poetry run mypy .
+      - name: Test Suite with Coverage Report
+        run: make coverage-report-test
+      - name: Build check
+        run: poetry build
+      - name: Upload Coverage Report Artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: test-coverage
+          path: htmlcov
 
   notebook-checks:
     name: Notebook Checks
@@ -135,3 +147,61 @@ jobs:
       - name: Check Notebook Output Consistency
         # Note: currently only checks notebooks we have specifically setup for this
         run: make notebook-test
+
+
+  build-docs:
+    # When running on a PR, this just checks we can build the docs without errors
+    # When running on merge to main, it builds the docs and then another job deploys them
+    name: ${{ github.event_name == 'pull_request' && 'Check Build Docs' || 'Build Docs' }}
+    runs-on: ubuntu-latest
+    needs: code-checks
+    steps:
+      - uses: actions/checkout@v4
+      - 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: Install pandoc
+        uses: awalsh128/cache-apt-pkgs-action@latest
+        with:
+          packages: pandoc
+          version: 1.0
+      - name: Install dependencies
+        run: poetry install --with docs
+      - name: Download Test Coverage Artifact
+        uses: actions/download-artifact@v3
+        with:
+          name: test-coverage
+          path: docs/source/coverage
+      - name: Build Docs
+        run: HF_TOKEN="$HF_TOKEN" poetry run build-docs
+        env:
+          HF_TOKEN: "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"
+      - name: Upload Docs Artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: documentation
+          path: docs/build
+
+  deploy-docs:
+    name: Deploy Docs
+    runs-on: ubuntu-latest
+    # Only run if merging a PR into main
+    if: github.event_name == 'push' && github.ref == 'refs/heads/main'
+    needs: build-docs
+    steps:
+      - uses: actions/checkout@v4
+      - name: Download Docs Artifact
+        uses: actions/download-artifact@v3
+        with:
+          name: documentation
+          path: docs/build
+      - name: Upload to GitHub Pages
+        uses: JamesIves/github-pages-deploy-action@v4
+        with:
+          folder: docs/build
+          clean-exclude: |
+            *.*.*/

.gitignore

@@ -18,3 +18,4 @@ docs/build
 .Ds_Store
 .pylintrc
 docs/source/generated
+**.orig

.vscode/cspell.json

@@ -1,41 +1,53 @@
 {
     "language": "en,en-GB",
     "words": [
+        "accum",
         "adrià",
         "aengus",
+        "allclose",
         "alonso",
         "arange",
         "argmax",
+        "argmaxy",
         "autodiff",
         "autoregressive",
         "barez",
-        "Beartype",
+        "beartype",
         "belrose",
         "bertsimas",
         "biderman",
         "bilal",
+        "bincount",
         "caxis",
         "checkpointed",
         "chughtai",
         "circuitsvis",
-        "Codespaces",
+        "Codeparrot",
+        "codespaces",
         "colab",
         "collectstart",
         "colour",
         "conmy",
         "cooney",
+        "crfm",
+        "cumsum",
+        "datapoint",
+        "dictmodel",
         "dimitris",
         "disconfirm",
         "dmitrii",
         "docstrings",
         "doctest",
         "doctree",
         "dtype",
+        "dtypes",
         "einops",
         "elhage",
+        "endoftext",
         "eqnarray",
         "esben",
         "evals",
+        "explictly",
         "fazl",
         "firstpage",
         "fspath",
@@ -49,53 +61,77 @@
         "howpublished",
         "huggingface",
         "icml",
+        "idxs",
         "imshow",
         "interp",
         "interpretability",
         "ioannis",
         "ipynb",
+        "isin",
         "isort",
         "janiak",
+        "Janky",
         "jaxtyping",
         "jett",
+        "kaiming",
         "keepdim",
         "kissane",
         "konstas",
         "kran",
         "lastpage",
         "layernorm",
+        "ldim",
         "lieberum",
         "logits",
         "logsumexp",
         "mavor",
         "maxdepth",
         "mingpt",
         "nanda",
+        "ndarray",
+        "ndim",
         "neel",
         "neox",
         "nitpicky",
+        "occurences",
         "olah",
+        "openwebtext",
+        "overcomplete",
+        "Overriden",
         "pagename",
         "pauly",
         "pretrained",
         "probs",
+        "producting",
         "pycln",
         "pypi",
         "pytest",
+        "randn",
+        "rdim",
+        "relu",
         "resid",
+        "rprint",
         "rtml",
+        "rtol",
+        "shortformer",
         "softmax",
+        "softmaxing",
         "solu",
+        "stas",
         "templatedir",
         "templatename",
         "toctree",
         "topk",
+        "tqdm",
         "transformerlens",
+        "tril",
+        "triu",
         "troitskii",
         "unembed",
         "unembedded",
         "unembedding",
         "unigram",
+        "unsqueeze",
         "virtualenvs",
         "visualisation",
         "xaxis",

.vscode/extensions.json

@@ -17,8 +17,9 @@
         "ms-toolsai.jupyter",
         "richie5um2.vscode-sort-json",
         "stkb.rewrap",
+        "streetsidesoftware.code-spell-checker-british-english",
         "streetsidesoftware.code-spell-checker",
-        "yzhang.markdown-all-in-one",
-        "streetsidesoftware.code-spell-checker-british-english"
+        "tamasfe.even-better-toml",
+        "yzhang.markdown-all-in-one"
     ]
 }