Docs generation (#60)

* finished setting up initial docs generation * added hot reload command * added github docs action * changed build command * ensured generated directory exists * made sure static directory exists * only used cuda if available * commented out blocks * commented out more blocks * added block again * added two more blocks * added two more blocks * added another block * added more blocks * removed another block * added cuda check again * removed block again * commented out last block * commented out remaining block * cleaned up markdown * setup mkdocs * setup commands * installed mkdocs * installed poethepoet * added init file for new module * fixed doc string typing * removed doc typing exceptions * regenerated lock file * added init files
ai-safety-foundation · Nov 16, 2023 · 214285f · 214285f
1 parent 0059c3d
commit 214285f
Show file tree

Hide file tree

Showing 26 changed files with 1,008 additions and 397 deletions.
diff --git a/.github/workflows/gh-pages.yml b/.github/workflows/gh-pages.yml
@@ -0,0 +1,60 @@
+name: Docs
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - '*'
+
+permissions:
+  contents: write
+
+jobs:
+  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
+    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 poe
+        run: pip install poethepoet
+      - name: Install mkdocs
+        run: pip install mkdocs
+      - name: Install dependencies
+        run: poetry install --with docs
+      - name: Build Docs
+        run: poe make-docs
+      - 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: |
+            *.*.*/
diff --git a/.gitignore b/.gitignore
@@ -52,6 +52,10 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 
+# Docs stuff
+docs/content/reference
+docs/generated
+
 # Translations
 *.mo
 *.pot

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -18,13 +18,13 @@ repos:
     hooks:
       - id: ruff_lint
         name: Ruff Lint
-        entry: poetry run ruff check sparse_autoencoder
+        entry: poetry run ruff check sparse_autoencoder --fix
         language: system
         types: [python]
         require_serial: true
       - id: ruff_format
         name: Ruff Format
-        entry: poetry run ruff format sparse_autoencoder --check
+        entry: poetry run ruff format sparse_autoencoder --check --diff
         language: system
         types: [python]
         require_serial: true

diff --git a/benchmarks.ipynb b/benchmarks.ipynb
@@ -238,22 +238,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "num_items: int = 1_000_000\n",
-    "num_neurons: int = 2048\n",
+    "# num_items: int = 1_000_000\n",
+    "# num_neurons: int = 2048\n",
     "\n",
     "\n",
-    "def create_dummy_activations(\n",
-    "    n_items: int, n_neurons: int\n",
-    ") -> list[Float[Tensor, \"batch neurons\"]]:\n",
-    "    \"\"\"Create Dummy Activations for Benchmarks.\"\"\"\n",
-    "    batch_size = 1_000\n",
-    "    n_batches = int(n_items // batch_size)\n",
-    "    activations = [torch.rand(batch_size, n_neurons) for _ in range(n_batches)]\n",
-    "    return activations\n",
+    "# def create_dummy_activations(\n",
+    "#     n_items: int, n_neurons: int\n",
+    "# ) -> list[Float[Tensor, \"batch neurons\"]]:\n",
+    "#     \"\"\"Create Dummy Activations for Benchmarks.\"\"\"\n",
+    "#     batch_size = 1_000\n",
+    "#     n_batches = int(n_items // batch_size)\n",
+    "#     activations = [torch.rand(batch_size, n_neurons) for _ in range(n_batches)]\n",
+    "#     return activations\n",
     "\n",
     "\n",
-    "dummy_activations = create_dummy_activations(num_items, num_neurons)\n",
-    "dummy_activations[0].shape"
+    "# dummy_activations = create_dummy_activations(num_items, num_neurons)\n",
+    "# dummy_activations[0].shape"
    ]
   },
   {
@@ -262,19 +262,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "benchmarks_to_run = {\n",
-    "    \"GPU Tensor\": TensorActivationStore(\n",
-    "        max_items=num_items, num_neurons=num_neurons, device=torch.device(\"cuda\")\n",
-    "    ),\n",
-    "    \"CPU Tensor\": TensorActivationStore(\n",
-    "        max_items=num_items, num_neurons=num_neurons, device=torch.device(\"cpu\")\n",
-    "    ),\n",
-    "    \"CPU List, No Multiprocessing\": ListActivationStore(),\n",
-    "    \"CPU List, Multiprocessing (multiple GPUs)\": ListActivationStore(\n",
-    "        multiprocessing_enabled=True\n",
-    "    ),\n",
-    "    \"Disk\": DiskActivationStore(empty_dir=True, max_cache_size=100_000),\n",
-    "}"
+    "# benchmarks_to_run = {\n",
+    "#     \"GPU Tensor\": TensorActivationStore(\n",
+    "#         max_items=num_items, num_neurons=num_neurons, device=torch.device(\"cuda\")\n",
+    "#     ),\n",
+    "#     \"CPU Tensor\": TensorActivationStore(\n",
+    "#         max_items=num_items, num_neurons=num_neurons, device=torch.device(\"cpu\")\n",
+    "#     ),\n",
+    "#     \"CPU List, No Multiprocessing\": ListActivationStore(),\n",
+    "#     \"CPU List, Multiprocessing (multiple GPUs)\": ListActivationStore(\n",
+    "#         multiprocessing_enabled=True\n",
+    "#     ),\n",
+    "#     \"Disk\": DiskActivationStore(empty_dir=True, max_cache_size=100_000),\n",
+    "# }"
    ]
   },
   {
@@ -298,21 +298,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "results = {}\n",
+    "# results = {}\n",
     "\n",
-    "for name, store in tqdm.tqdm(benchmarks_to_run.items()):\n",
-    "    store.empty()\n",
-    "    start_time = time.time()\n",
-    "    for batch in dummy_activations:\n",
-    "        store.extend(batch)\n",
-    "    if hasattr(store, \"wait_for_writes_to_complete\"):\n",
-    "        store.wait_for_writes_to_complete()\n",
-    "    end_time = time.time()\n",
-    "    results[name] = end_time - start_time\n",
+    "# for name, store in tqdm.tqdm(benchmarks_to_run.items()):\n",
+    "#     store.empty()\n",
+    "#     start_time = time.time()\n",
+    "#     for batch in dummy_activations:\n",
+    "#         store.extend(batch)\n",
+    "#     if hasattr(store, \"wait_for_writes_to_complete\"):\n",
+    "#         store.wait_for_writes_to_complete()\n",
+    "#     end_time = time.time()\n",
+    "#     results[name] = end_time - start_time\n",
     "\n",
-    "df = pd.DataFrame(results, index=[\"Time (s)\"]).T\n",
-    "df[\"Time 10B (h estimate)\"] = df[\"Time (s)\"] * 10**10 / num_items / 3600\n",
-    "df"
+    "# df = pd.DataFrame(results, index=[\"Time (s)\"]).T\n",
+    "# df[\"Time 10B (h estimate)\"] = df[\"Time (s)\"] * 10**10 / num_items / 3600\n",
+    "# df"
    ]
   },
   {

diff --git a/docs/content/SUMMARY.md b/docs/content/SUMMARY.md
@@ -0,0 +1,5 @@
+* [Home](index.md)
+* [Getting Started](getting_started.md)
+* [Contributing](contributing.md)
+* [Citation](citation.md)
+* [Reference](reference/)
diff --git a/docs/content/citation.md b/docs/content/citation.md
@@ -0,0 +1,4 @@
+
+# Citation
+
+Please cite this library as:
diff --git a/docs/content/contributing.md b/docs/content/contributing.md
@@ -0,0 +1,140 @@
+# Contributing
+
+## Setup
+
+### DevContainer
+
+For a one-click setup of your development environment, this project includes a
+[DevContainer](https://containers.dev/). It can be used locally with [VS
+Code](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers) or
+with [GitHub Codespaces](https://github.com/features/codespaces).
+
+### Manual Setup
+
+This project uses [Poetry](https://python-poetry.org/docs/#installation) for package management.
+Install as follows (this will also setup your virtual environment):
+
+```bash
+poetry config virtualenvs.in-project true
+poetry install --with dev,docs,jupyter
+```
+
+## Testing
+
+If adding a feature, please add unit tests for it. If you need a model, please use one of the ones
+that are cached by GitHub Actions (so that it runs quickly on the CD). These are `gpt2`,
+`attn-only-1l`, `attn-only-2l`, `attn-only-3l`, `attn-only-4l`, `tiny-stories-1M`. Note `gpt2` is
+quite slow (as we only have CPU actions) so the smaller models like `attn-only-1l` and
+`tiny-stories-1M` are preferred if possible.
+
+### Running the tests
+
+- Unit tests only via `make unit-test`
+- Acceptance tests only via `make acceptance-test`
+- Docstring tests only via `make docstring-test`
+
+## Formatting
+
+This project uses `pycln`, `isort` and `black` for formatting, pull requests are checked in github
+actions.
+
+- Format all files via `make format`
+- Only check the formatting via `make check-format`
+
+## Documentation
+
+Please make sure to add thorough documentation for any features you add. You should do this directly
+in the docstring, and this will then automatically generate the API docs when merged into `main`.
+They will also be automatically checked with [pytest](https://docs.pytest.org/) (via
+[doctest](https://docs.python.org/3/library/doctest.html)).
+
+If you want to view your documentation changes, run `pytest run docs-hot-reload`. This will give you
+hot-reloading docs (they change in real time as you edit docstrings).
+
+### Docstring Style Guide
+
+We follow the Google Python Docstring Style for writing docstrings, with some added features from
+reStructuredText (reST).
+
+#### Sections and Order
+
+You should follow this order:
+
+```python
+"""Title In Title Case.
+
+A description of what the function/class does, including as much detail as is necessary to fully understand it.
+
+Warning:
+
+Any warnings to the user (e.g. common pitfalls).
+
+Examples:
+
+Include any examples here. They will be checked with doctest.
+
+  >>> print(1 + 2)
+  3
+
+Args:
+    param_without_type_signature:
+        Each description should be indented once more.
+    param_2:
+        Another example parameter.
+
+Returns:
+    Returns description without type signature.
+
+Raises:
+    Information about the error it may raise (if any).
+"""
+```
+
+#### Supported Sphinx Properties
+
+##### References to Other Functions/Classes
+
+You can reference other parts of the codebase using
+[cross-referencing](https://www.sphinx-doc.org/en/master/usage/domains/python.html#cross-referencing-python-objects)
+(noting that you can omit the full path if it is in the same file).
+
+```reStructuredText
+:mod:transformer_lens # Function or module
+
+:const:`transformer_lens.loading_from_pretrained.OFFICIAL_MODEL_NAMES`
+
+:class:`transformer_lens.HookedTransformer`
+
+:meth:`transformer_lens.HookedTransformer.from_pretrained`
+
+:attr:`transformer_lens.HookedTransformer.cfg`
+```
+
+##### Maths
+
+You can use LaTeX, but note that as you're placing this in python strings the backwards slash (`\`)
+must be repeated (i.e. `\\`). You can write LaTeX inline, or in "display mode".
+
+```reStructuredText
+:math:`(a + b)^2 = a^2 + 2ab + b^2`
+```
+
+```reStructuredText
+.. math::
+   :nowrap:
+
+   \\begin{eqnarray}
+      y    & = & ax^2 + bx + c \\
+      f(x) & = & x^2 + 2xy + y^2
+   \\end{eqnarray}
+```
+
+#### Markup
+
+- Italics - `*text*`
+- Bold - `**text**`
+- Code - ` ``code`` `
+- List items - `*item`
+- Numbered items - `1. Item`
+- Quotes - indent one level
+- External links = ``` `Link text <https://domain.invalid/>` ```
diff --git a/docs/content/getting_started.md b/docs/content/getting_started.md
@@ -0,0 +1 @@
+# Getting Started
diff --git a/docs/content/index.md b/docs/content/index.md
@@ -0,0 +1,16 @@
+# Home
+
+[![PyPI](https://img.shields.io/pypi/v/sparse_autoencoder?color=blue)](https://pypi.org/project/transformer-lens/)
+![PyPI -
+License](https://img.shields.io/pypi/l/sparse_autoencoder?color=blue) [![Checks](https://github.com/alan-cooney/sparse_autoencoder/actions/workflows/checks.yml/badge.svg)](https://github.com/alan-cooney/sparse_autoencoder/actions/workflows/checks.yml)
+[![Release](https://github.com/alan-cooney/sparse_autoencoder/actions/workflows/release.yml/badge.svg)](https://github.com/alan-cooney/sparse_autoencoder/actions/workflows/release.yml)
+
+A sparse autoencoder for mechanistic interpretability research.
+
+```shell
+pip install sparse_autoencoder
+```
+
+## Demo
+
+Check out the demo notebook for a guide to using this library.