flake8 and black

.flake8

@@ -5,4 +5,4 @@ max-complexity = 25
 extend-select = E9, F63, F7, F82
 show-source = true
 statistics = true
-exclude = ./sae_training/geom_median/
+exclude = ./sae_training/geom_median/, ./wandb/*, ./research/wandb/*

.pre-commit-config.yaml

@@ -7,19 +7,19 @@ repos:
     -   id: check-added-large-files
         args: [--maxkb=250000]
 -   repo: /~https://github.com/psf/black
-    rev: 23.3.0
+    rev: 24.2.0
     hooks:
     -   id: black
 -   repo: /~https://github.com/PyCQA/flake8
     rev: 6.0.0
     hooks:
     -   id: flake8
+        args: ['--config=.flake8']
         additional_dependencies: [
             'flake8-blind-except',
-            'flake8-docstrings',
+            # 'flake8-docstrings',
             'flake8-bugbear',
             'flake8-comprehensions',
-            'flake8-docstrings',
             'flake8-implicit-str-concat',
             'pydocstyle>=5.0.0',
         ]

pyproject.toml

@@ -18,6 +18,7 @@ ipykernel = "^6.29.2"
 matplotlib = "^3.8.3"
 matplotlib-inline = "^0.1.6"
 eindex = {git = "/~https://github.com/callummcdougall/eindex.git"}
+datasets = "^2.17.1"
 
 
 [tool.poetry.group.dev.dependencies]
@@ -33,4 +34,4 @@ profile = "black"
 
 [build-system]
 requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+build-backend = "poetry.core.masonry.api"

sae_analysis/dashboard_runner.py

@@ -17,10 +17,10 @@
 import plotly
 import plotly.express as px
 import torch
-import wandb
 from torch.nn.functional import cosine_similarity
 from tqdm import tqdm
 
+import wandb
 from sae_analysis.visualizer.data_fns import get_feature_data
 from sae_training.utils import LMSparseAutoencoderSessionloader
 

sae_training/config.py

@@ -3,6 +3,7 @@
 from typing import Optional
 
 import torch
+
 import wandb
 
 
@@ -21,9 +22,9 @@ class RunnerConfig(ABC):
     is_dataset_tokenized: bool = True
     context_size: int = 128
     use_cached_activations: bool = False
-    cached_activations_path: Optional[
-        str
-    ] = None  # Defaults to "activations/{dataset}/{model}/{full_hook_name}_{hook_point_head_index}"
+    cached_activations_path: Optional[str] = (
+        None  # Defaults to "activations/{dataset}/{model}/{full_hook_name}_{hook_point_head_index}"
+    )
 
     # SAE Parameters
     d_in: int = 512
@@ -61,7 +62,9 @@ class LanguageModelSAERunnerConfig(RunnerConfig):
     l1_coefficient: float = 1e-3
     lp_norm: float = 1
     lr: float = 3e-4
-    lr_scheduler_name: str = "constantwithwarmup"  # constant, constantwithwarmup, linearwarmupdecay, cosineannealing, cosineannealingwarmup
+    lr_scheduler_name: str = (
+        "constantwithwarmup"  # constant, constantwithwarmup, linearwarmupdecay, cosineannealing, cosineannealingwarmup
+    )
     lr_warm_up_steps: int = 500
     train_batch_size: int = 4096
 

sae_training/evals.py

@@ -2,11 +2,11 @@
 
 import pandas as pd
 import torch
-import wandb
 from tqdm import tqdm
 from transformer_lens import HookedTransformer
 from transformer_lens.utils import get_act_name
 
+import wandb
 from sae_training.activations_store import ActivationsStore
 from sae_training.sparse_autoencoder import SparseAutoencoder
 

sae_training/sae_group.py

@@ -1,10 +1,12 @@
+import dataclasses
 import gzip
 import os
 import pickle
-import dataclasses
-from sae_training.sparse_autoencoder import SparseAutoencoder
+
 import torch
 
+from sae_training.sparse_autoencoder import SparseAutoencoder
+
 
 class SAEGroup:
     def __init__(self, cfg):

sae_training/toy_model_runner.py

@@ -2,8 +2,8 @@
 
 import einops
 import torch
-import wandb
 
+import wandb
 from sae_training.sparse_autoencoder import SparseAutoencoder
 from sae_training.toy_models import Config as ToyConfig
 from sae_training.toy_models import Model as ToyModel

sae_training/train_sae_on_language_model.py

@@ -1,14 +1,14 @@
 import torch
-import wandb
 from torch.optim import Adam
 from tqdm import tqdm
 from transformer_lens import HookedTransformer
 
+import wandb
 from sae_training.activations_store import ActivationsStore
 from sae_training.evals import run_evals
+from sae_training.geom_median.src.geom_median.torch import compute_geometric_median
 from sae_training.optim import get_scheduler
 from sae_training.sae_group import SAEGroup
-from sae_training.geom_median.src.geom_median.torch import compute_geometric_median
 
 
 def train_sae_on_language_model(

sae_training/train_sae_on_toy_model.py

@@ -1,8 +1,8 @@
 import torch
-import wandb
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+import wandb
 from sae_training.sparse_autoencoder import SparseAutoencoder
 
 

scripts/generate_dashboards.py

@@ -16,10 +16,10 @@
 import plotly
 import plotly.express as px
 import torch
-import wandb
 from torch.nn.functional import cosine_similarity
 from tqdm import tqdm
 
+import wandb
 from sae_analysis.visualizer.data_fns import get_feature_data
 from sae_training.utils import LMSparseAutoencoderSessionloader