lint + precommit

Signed-off-by: jwilber <jwilber@nvidia.com>

Author: jwilber

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/run.py

@@ -55,7 +55,7 @@ def _torchrun_prefix(nproc: int) -> list:
     return [sys.executable]
 
 
-def run_extract(cfg: DictConfig, cache_dir: Path) -> None:
+def run_extract(cfg: DictConfig, cache_dir: Path) -> None:  # noqa: D103
     cmd = [
         *_torchrun_prefix(cfg.nproc),
         str(SCRIPTS_DIR / "extract.py"),
@@ -82,7 +82,7 @@ def run_extract(cfg: DictConfig, cache_dir: Path) -> None:
     _run(cmd, f"STEP 1: Extract activations from {cfg.model_path}")
 
 
-def run_train(cfg: DictConfig, cache_dir: Path, output_dir: Path) -> None:
+def run_train(cfg: DictConfig, cache_dir: Path, output_dir: Path) -> None:  # noqa: D103
     checkpoint_dir = output_dir / "checkpoints"
     t = cfg.train
 
@@ -143,7 +143,7 @@ def run_train(cfg: DictConfig, cache_dir: Path, output_dir: Path) -> None:
     _run(cmd, "STEP 2: Train SAE")
 
 
-def run_eval(cfg: DictConfig, output_dir: Path) -> None:
+def run_eval(cfg: DictConfig, output_dir: Path) -> None:  # noqa: D103
     checkpoint = output_dir / "checkpoints" / "checkpoint_final.pt"
     eval_dir = output_dir / "eval"
 
@@ -176,7 +176,7 @@ def run_eval(cfg: DictConfig, output_dir: Path) -> None:
 
 
 @hydra.main(version_base=None, config_path="run_configs", config_name="config")
-def main(cfg: DictConfig) -> None:
+def main(cfg: DictConfig) -> None:  # noqa: D103
     os.chdir(hydra.utils.get_original_cwd())
 
     print(OmegaConf.to_yaml(cfg))

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/scripts/analyze.py

@@ -45,11 +45,11 @@
 _CODONFM_TE_DIR = _REPO_ROOT / "recipes" / "codonfm_ptl_te"
 sys.path.insert(0, str(_CODONFM_TE_DIR))
 
-from codonfm_sae.data import read_codon_csv
-from sae.architectures import TopKSAE
-from sae.utils import get_device, set_seed
-from src.data.preprocess.codon_sequence import process_item
-from src.inference.encodon import EncodonInference
+from codonfm_sae.data import read_codon_csv  # noqa: E402
+from sae.architectures import TopKSAE  # noqa: E402
+from sae.utils import get_device, set_seed  # noqa: E402
+from src.data.preprocess.codon_sequence import process_item  # noqa: E402
+from src.inference.encodon import EncodonInference  # noqa: E402
 
 
 # ── Standard codon usage table (human, per 1000 codons) ──────────────
@@ -189,7 +189,7 @@
 }
 
 
-def parse_args():
+def parse_args():  # noqa: D103
     p = argparse.ArgumentParser(description="Analyze CodonFM SAE features")
     p.add_argument("--checkpoint", type=str, required=True)
     p.add_argument("--top-k", type=int, default=None, help="Override top-k (default: read from checkpoint)")
@@ -232,7 +232,7 @@ def parse_args():
     return p.parse_args()
 
 
-def load_sae(checkpoint_path: str, top_k_override: int | None = None) -> TopKSAE:
+def load_sae(checkpoint_path: str, top_k_override: int | None = None) -> TopKSAE:  # noqa: D103
     ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
     state_dict = ckpt["model_state_dict"]
     if any(k.startswith("module.") for k in state_dict):
@@ -533,7 +533,7 @@ def stream_annotations_and_topk(
 # ── 3. Auto-interpretation ───────────────────────────────────────────
 
 
-def get_llm_client(provider: str, model: str = None):
+def get_llm_client(provider: str, model: str | None = None):
     """Create LLM client based on provider."""
     from sae.autointerp import (
         AnthropicClient,
@@ -834,7 +834,7 @@ def build_feature_labels(
 # ── Main ─────────────────────────────────────────────────────────────
 
 
-def main():
+def main():  # noqa: D103
     args = parse_args()
     set_seed(args.seed)
     device = args.device or get_device()

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/scripts/eval.py

@@ -41,14 +41,14 @@
 _CODONFM_TE_DIR = _REPO_ROOT / "recipes" / "codonfm_ptl_te"
 sys.path.insert(0, str(_CODONFM_TE_DIR))
 
-from codonfm_sae.data import read_codon_csv
-from codonfm_sae.eval import evaluate_codonfm_loss_recovered
-from sae.architectures import TopKSAE
-from sae.utils import get_device, set_seed
-from src.inference.encodon import EncodonInference
+from codonfm_sae.data import read_codon_csv  # noqa: E402
+from codonfm_sae.eval import evaluate_codonfm_loss_recovered  # noqa: E402
+from sae.architectures import TopKSAE  # noqa: E402
+from sae.utils import get_device, set_seed  # noqa: E402
+from src.inference.encodon import EncodonInference  # noqa: E402
 
 
-def parse_args():
+def parse_args():  # noqa: D103
     p = argparse.ArgumentParser(description="Evaluate CodonFM SAE")
 
     # Checkpoint
@@ -110,7 +110,7 @@ def load_sae_from_checkpoint(checkpoint_path: str, top_k_override: int | None =
     return sae
 
 
-def main():
+def main():  # noqa: D103
     args = parse_args()
     set_seed(args.seed)
     device = args.device or get_device()

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/scripts/eval_swissprot_f1.py

@@ -51,10 +51,10 @@
 _CODONFM_TE_DIR = _REPO_ROOT / "recipes" / "codonfm_ptl_te"
 sys.path.insert(0, str(_CODONFM_TE_DIR))
 
-from sae.architectures import TopKSAE
-from sae.utils import get_device, set_seed
-from src.data.preprocess.codon_sequence import process_item
-from src.inference.encodon import EncodonInference
+from sae.architectures import TopKSAE  # noqa: E402
+from sae.utils import get_device, set_seed  # noqa: E402
+from src.data.preprocess.codon_sequence import process_item  # noqa: E402
+from src.inference.encodon import EncodonInference  # noqa: E402
 
 
 # ── Annotation parsing (adapted from esm2_sae) ─────────────────────────

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/scripts/extract.py

@@ -48,13 +48,13 @@
 _CODONFM_TE_DIR = _REPO_ROOT / "recipes" / "codonfm_ptl_te"
 sys.path.insert(0, str(_CODONFM_TE_DIR))
 
-from codonfm_sae.data import read_codon_csv
-from sae.activation_store import ActivationStore, ActivationStoreConfig
-from src.data.preprocess.codon_sequence import process_item
-from src.inference.encodon import EncodonInference
+from codonfm_sae.data import read_codon_csv  # noqa: E402
+from sae.activation_store import ActivationStore, ActivationStoreConfig  # noqa: E402
+from src.data.preprocess.codon_sequence import process_item  # noqa: E402
+from src.inference.encodon import EncodonInference  # noqa: E402
 
 
-def parse_args():
+def parse_args():  # noqa: D103
     p = argparse.ArgumentParser(description="Extract CodonFM layer activations")
     p.add_argument(
         "--csv-path", type=str, required=True, help="Path to CSV with DNA sequences (auto-detects 'seq'/'cds' column)"
@@ -138,7 +138,7 @@ def _merge_rank_stores(cache_path: Path, world_size: int, metadata: dict) -> Non
         print(f"Merged {world_size} rank stores: {total_samples:,} tokens, {shard_idx} shards")
 
 
-def main():
+def main():  # noqa: D103
     args = parse_args()
     torch.manual_seed(args.seed)
 

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/scripts/launch_dashboard.py

@@ -60,7 +60,7 @@ def _filter_and_copy_parquet(src: Path, dst: Path, live_ids: set):
     return n_before, len(df)
 
 
-def main():
+def main():  # noqa: D103
     p = argparse.ArgumentParser(description="Launch codon SAE dashboard")
     p.add_argument(
         "--data-dir",

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/scripts/train.py

@@ -45,7 +45,7 @@
 from sae.utils import get_device, set_seed
 
 
-def parse_args():
+def parse_args():  # noqa: D103
     p = argparse.ArgumentParser(
         description="Train SAE from cached CodonFM activations",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -110,7 +110,7 @@ def parse_args():
     return p.parse_args()
 
 
-def build_sae(args, input_dim: int) -> torch.nn.Module:
+def build_sae(args, input_dim: int) -> torch.nn.Module:  # noqa: D103
     hidden_dim = input_dim * args.expansion_factor
 
     if args.model_type == "topk":
@@ -133,7 +133,7 @@ def build_sae(args, input_dim: int) -> torch.nn.Module:
         raise ValueError(f"Unknown model type: {args.model_type}")
 
 
-def build_training_config(args, device: str) -> TrainingConfig:
+def build_training_config(args, device: str) -> TrainingConfig:  # noqa: D103
     return TrainingConfig(
         lr=args.lr,
         n_epochs=args.n_epochs,
@@ -150,7 +150,7 @@ def build_training_config(args, device: str) -> TrainingConfig:
     )
 
 
-def build_wandb_config(args) -> WandbConfig:
+def build_wandb_config(args) -> WandbConfig:  # noqa: D103
     return WandbConfig(
         enabled=args.wandb_enabled,
         project=args.wandb_project,
@@ -161,11 +161,11 @@ def build_wandb_config(args) -> WandbConfig:
     )
 
 
-def build_parallel_config(args) -> ParallelConfig:
+def build_parallel_config(args) -> ParallelConfig:  # noqa: D103
     return ParallelConfig(dp_size=args.dp_size)
 
 
-def main():
+def main():  # noqa: D103
     args = parse_args()
 
     set_seed(args.seed)

bionemo-recipes/interpretability/sparse_autoencoders/recipes/codonfm/src/codonfm_sae/data/types.py

@@ -26,5 +26,5 @@ class CodonRecord:
     metadata: Dict[str, Any] = field(default_factory=dict)
 
     @property
-    def num_codons(self) -> int:
+    def num_codons(self) -> int:  # noqa: D102
         return len(self.sequence) // 3

bionemo-recipes/interpretability/sparse_autoencoders/recipes/esm2/scripts/launch_dashboard.py

@@ -59,7 +59,7 @@ def _filter_and_copy_parquet(src: Path, dst: Path, live_ids: set):
     return n_before, len(df)
 
 
-def main():
+def main():  # noqa: D103
     p = argparse.ArgumentParser(description="Launch ESM2 SAE dashboard")
     p.add_argument(
         "--data-dir",