evo2 SAE: inference engine + steering + server/CLI, tests, Dockerfile

[polinabinder1]

Summary

The importable Evo2SAE inference engine + feature steering — the base of the serve stack — with tests and a runnable (layer-cached) Docker image. A single Evo2 inference engine is loaded once and serves both paths: encode reads the residual stream off a layer-L forward hook; generate drives the same model's decode with decode-only feature steering. This PR now also includes the FastAPI server + CLI (folded in from #1637, which merged into this branch); the dashboard (#1623) and steering eval (#1635) build on it.

Rebased onto the post-#1633 top-level layout (interpretability/sparse_autoencoders/).

Stack update (2026-06-29): #1637 (FastAPI server + CLI) was merged into this branch — its code (server.py, cli.py, their tests, launch_inference.sh) is now part of this PR. This PR's scope is therefore engine + steering + server + CLI.

Architecture: one model, both paths

Earlier iterations loaded two copies of Evo2 — a truncated post_process=False model for encode/highlight and the full inference engine for generate (~1.8× the weights). This collapses to a single engine (infer.setup_inference_engine, run eager with cuda_graph_impl="none" so the steering hook applies):

• load() builds the one engine and takes self.model = unwrap_model(comp.model) + comp.tokenizer from it.
• encode/highlight (_forward_hidden) runs a normal full-sequence forward and reads layer L off a forward hook — the engine model is post_process=True (it produces logits for generation), so output_embeddings can't be used; the hook captures the same [S, B, H] module output the steering clamp_hook reads, so encode and steer see identical activations by construction.
• generate steers on self.model.decoder.layers[L] — the same module encode reads.

Validated end-to-end on the 1B-8k-bf16 (21/21 tests, incl. a highlight↔steer interleaving test proving no state bleed between the shared model's encode forward and decode path). 7B fidelity is the remaining gate.

Note on the encode forward

The single engine model is post_process=True with no embeddings-only mode, so encode runs the full forward (all layers + the LM head) and reads layer L off the hook in passing — the trailing layers after L and the logits are computed and discarded. Cost is ~the layers past L (e.g. 6 of 32 on the 7B, 6 of 25 on the 1B; the vocab-512 head is cheap). Negligible for interactive encode; a ~20% overhead only for batch encode — and bulk extraction doesn't use this path anyway (scripts/extract.py reuses predict_evo2 --embedding-layer, which stops at L). An early-exit (raise from the hook to abort the forward) was left out deliberately: aborting mid-forward risks dirtying the engine's inference context (the no-bleed property) and would deadlock under tensor/pipeline parallel.

GPU & scaling

The engine is single-GPU. Everything keys off one self.device; serve is a single python -m evo2_sae.cli process (not torchrun), pinned via DEVICE / CUDA_VISIBLE_DEVICES. The 7B fits on one 80 GB H100, so model-parallel isn't needed for ≤7B.

• Model-parallel (TP/PP) is not supported as written. The SAE consumes the full [S,B,H] hidden off the layer-L hook (input_dim=4096); under tensor-parallel that hidden is sharded across ranks (SAE would see a fragment), and under pipeline-parallel layer L lives on a single rank (the hook fires nowhere else). Supporting it needs an all_gather of the layer output + rank-aware hook placement — real work, only relevant for a model too big for one GPU (e.g. 40B).
• Data-parallel serving scales trivially, no code change — run N independent serve processes, each pinned to its own GPU (CUDA_VISIBLE_DEVICES) + PORT, behind a load balancer. Each is a full, independent engine.
• Offline bulk extraction already runs multi-GPU — scripts/extract.py under torchrun --nproc_per_node N is data-parallel (each rank a full model writing its own shards); this is how the 7B activation set was built.

Contents

Engine + steering

• src/evo2_sae/core.py — Evo2SAE: load → encode / encode_batch / feature_tracks / generate (decode-only clamp via sae.steering) + input-sanitization guards (_sanitize_steering: feature-id range, clamp-magnitude cap, non-finite/top_k/temperature coercion). encode_batch is length-bucketed (work sorted by token length to minimize padding waste on mixed-length inputs; results un-sorted back to input order).
• Load-time SAE/model fit check — load() verifies the SAE's input_dim equals the model's hidden size (_model_hidden_size via config, or a 1-token forward) and raises a clear error on a mismatch ("wrong SAE/model pairing"), instead of a cryptic matmul failure on the first encode. Known gap: a wrong layer number with the same hidden size can't be caught here (the SAE checkpoint records no training layer) — it silently yields out-of-distribution features; /health surfaces the configured layer, and stamping the layer into the checkpoint at train time is a follow-up.
• sae/src/sae/steering.py — model-agnostic delta-clamp hook + steer().

Server + CLI (folded in from #1637)

• src/evo2_sae/server.py — FastAPI app (/health, /features, /annotate, /generate) under an /api prefix, with an optional static-dashboard mount (used by evo2 SAE recipe: feature-explorer dashboard (viz) #1623).
• src/evo2_sae/cli.py + scripts/launch_inference.sh — serve / encode / batch / generate entry points (required env vars, no hardcoded paths).

Build / run / CI

• .ci_build.sh (env | install | all) + .ci_test_env.sh — build the env by delegating to evo2_megatron's own build (no fork of the pinned megatron stack), then install sae + this recipe into that venv. The phase arg lets the Dockerfile cache the two steps separately.
• Dockerfile — thin, non-forking, layer-cached: the ~30-min mbridge megatron build is its own layer (depends only on recipes/evo2_megatron), and the SAE source + editable installs are a separate trailing layer — so editing engine/SAE code rebuilds only the cheap install layer, not megatron. (+ a per-Dockerfile .dockerignore.)
• tests/conftest.py — 1B-8k-bf16 fixture (bionemo_load → run_nemo2_to_mbridge) + a synthesized tiny SAE, GPU-memory-gated; honors EVO2_CKPT_DIR / SAE_CKPT_PATH for manual / 7B runs. The GPU tests are gated by @pytest.mark.skipif(not torch.cuda.is_available()), so they run on a GPU box and skip otherwise.

Dependency on bionemo.evo2

The engine reuses bionemo.evo2's model code (the mbridge recipes/evo2_megatron recipe), which isn't pip-installable. .ci_build.sh (and the Dockerfile) build it via evo2_megatron's own script; it's intentionally not in pyproject.toml, matching the codonfm/esm2 recipes (base model is environment-provided).

Interfaces: CLI vs serve (and why FastAPI)

The same Evo2SAE is exposed two ways via cli.py:

• CLI one-shot (encode / batch / generate) — load → do one thing → exit. Pays the model load per invocation; right for scripted/batch one-offs (batch: FASTA → parquet).
• serve — load once, serve many HTTP requests; right for interactive use (the dashboard). Carries the production machinery: single-GPU lock + bounded threadpool, 503-until-ready gate, body-size cap, and the CUDA-wedge supervisor.

The FastAPI app is built by build_app(engine) with the engine injected — prod passes a real Evo2SAE, the tests pass a torch-free FakeEngine and drive the same app via TestClient (no model, CPU-only). That DI is why test_server.py locks the full HTTP contract without a GPU.

Why FastAPI (over Flask / gRPC / Gradio-Streamlit / Triton-TorchServe): the need is a thin JSON API over one GPU-bound engine, served from the same origin as a React SPA, testable without a model, in one container. FastAPI gives async + a bounded threadpool (anyio, so concurrency caps cleanly to the single serialized GPU), Pydantic request validation, auto OpenAPI docs, and a StaticFiles mount for single-origin serving (/api + dashboard at /, no CORS) — all out of the box. The alternatives are either too thin (Flask/raw — re-roll all of that), UI-coupled (Gradio/Streamlit — can't cleanly serve a separate React SPA + JSON API), browser-hostile (gRPC), or heavyweight model-serving platforms (Triton/TorchServe/Ray Serve) built for autoscale/multi-model/dynamic-batching that this single-model, single-GPU, single-container tool doesn't need. Those become worth it only if deployment grows to many replicas / multi-model — at which point FastAPI sits behind them.

Error handling / failure modes

An escalation ladder — nothing leaves the server permanently broken; worst case is a brief 503 during a respawn.

• Bad input → 4xx, the engine never sees it. Oversized body (>16 MB, MAX_BODY_BYTES) → 413 in middleware; over-context sequence → 413; non-DNA / unknown organism / out-of-range or negative feature id / bad annotate mode / missing pick ids → 400; wrong JSON shape → 422 (Pydantic). Some of these are safety, not UX: an out-of-range feature id would trip a CUDA device-side assert and a negative id would silently index the wrong feature via torch negative-indexing — so both are rejected before reaching the engine.
• Not ready → 503. _require_ready() and /health return 503 while the model loads (so k8s/LB readiness probes shed the pod). A startup failure is caught and logged; /health simply stays not-ready — process up, no crash loop.
• Unrecoverable CUDA wedge → self-heal. A device-side assert poisons the CUDA context process-wide (uncleanable in-process). _is_unrecoverable_cuda detects it → marks the engine not-ready (→ 503) → and if EXIT_ON_CUDA_WEDGE=1, the worker exits and launch_inference.sh respawns it (backoff, ≤10 consecutive fails, SIGTERM/SIGINT forwarded for docker stop/k8s). Recovery is a clean restart, host-independent.
• Other uncaught exception → 500, server continues. Model calls are serialized behind a lock, so one failed request can't corrupt shared state for the others.

(The dashboard's /gene_embed adds a partial-failure path — skip + report invalid sequences, all-invalid → 400 — documented in #1623.)

How to run

# Build once from the repo root, then run with a GPU:
docker build -f interpretability/sparse_autoencoders/recipes/evo2/Dockerfile -t evo2-sae .
docker run --gpus all -it evo2-sae bash -lc "source .ci_test_env.sh && pytest tests/"

On build time / making it easier. The engine needs bionemo.evo2 (the mbridge
evo2_megatron recipe), which isn't pip-installable — so the first docker build
compiles the full megatron stack (megatron-bridge, causal-conv1d, …) and takes ~30 min.
After that, the build is layer-cached: editing engine/SAE code re-runs only the two
editable pip installs (seconds), not the megatron compile. Other shortcuts:

• Build once, reuse: push the built image to a registry; coworkers docker run it and never rebuild.
• Skip the compile: the Dockerfile's ARG BASE_IMAGE can point at a prebuilt evo2_megatron / bionemo image once one exists — the build then collapses to just the two pip installs.
• No container at all (dev): inside an existing megatron env, pip install -e sae/ && pip install -e recipes/evo2/ (what local validation does).

from evo2_sae import Evo2SAE
eng   = Evo2SAE(evo2_ckpt_dir, sae_ckpt_path, layer=19).load()    # 1B layer 19 (7B: 26)
codes = eng.encode("ATGGCC...")                                    # [S, n_features], sparse (TopK)
out   = eng.generate(prompt="ATGGCC...", features=[{"feature_id": 123, "strength": 200}])

Tests

There's no dedicated CI lane right now (deferred — it should later fold into the repo-wide recipe lane, which already runs .ci_build.sh + pytest). Run them manually:

cd interpretability/sparse_autoencoders/recipes/evo2
bash .ci_build.sh && source .ci_test_env.sh   # build + activate the megatron venv
pytest tests/

• CPU (no model): test_core.py (engine plumbing — top_features, _load_sae, generate guards, the SAE/model dim check, encode_batch length-bucketing order) + test_server.py + test_cli.py (FastAPI app + CLI via a mocked FakeEngine) + test_steering.py sanitize guards + sae/tests/test_steering.py (exact clamp math). Quick CPU-only run without the venv: PYTHONPATH=src:../../sae/src pytest tests/test_core.py.
• GPU: test_steering.py — bf16 encode, generation in-distribution, steering changes the continuation (+ compare_baseline), batched/empty-sequence encode, max-clamp stays finite, and highlight↔steer interleaving (encode bit-identical across a steered generate; baseline unaffected by history). Gated by @pytest.mark.skipif(not torch.cuda.is_available()) — runs on a GPU box (megatron venv); set EVO2_CKPT_DIR/SAE_CKPT_PATH for a specific model, else the fixtures build the 1B-8k-bf16 + a synthesized SAE.

Base of

#1623 (dashboard) and #1635 (steering eval). (#1637 — server + CLI — merged into this branch and is now part of this PR.)

---

Note: recipes/evo2/ is co-owned with the eval stack (#1629)

This PR owns the recipe's Dockerfile / .ci_build.sh / src/evo2_sae (engine + server + CLI) + tests/conftest.py; the eval stack (#1629, which absorbed #1636) adds eval/probing (labelers, probe harness) and its biopython/pyrodigal deps to the same recipes/evo2/. The branches are reconciled — the pyproject.toml dep union is already resolved on #1629. Merge order: this PR first, then #1629 / #1635 / #1623 rebase onto main.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7c2e469a-96df-4765-ba83-fda964a7c92b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR introduces sparse autoencoder (SAE) feature steering capabilities for the Evo2 foundation model, along with a complete inference recipe. It adds a reusable clamp_hook steering primitive that injects only the delta between clamped and original SAE reconstructions, applies it to a new Evo2SAE inference engine supporting encoding and steered generation, and refactors FASTA parsing into a shared utility.

Changes

Evo2 SAE Steering and Inference Recipe

Layer / File(s)	Summary
Build configuration and runtime dependencies bionemo-recipes/.../evo2/pyproject.toml	Added fastapi>=0.110, uvicorn>=0.29, pandas>=1.5 to project dependencies and enabled setuptools package discovery under src/.
SAE feature steering primitives bionemo-recipes/.../sae/src/sae/steering.py	New clamp_hook forward hook re-encodes activations through the SAE, clamps specified feature codes, and injects the delta between clamped and original decoded outputs. steer context manager registers/removes the hook reliably.
SAE steering unit tests bionemo-recipes/.../sae/tests/test_steering.py	Validates delta-clamp correctness (no-op leaves unchanged, real clamp matches analytic delta), tuple output isolation (only hidden state modified), and decode-only mode (skips prefill).
Evo2SAE package API and lazy loading bionemo-recipes/.../evo2/src/evo2_sae/__init__.py	Public API via __all__ constraint on Evo2SAE, clean_dna, DEFAULT_ORGANISM_TAGS with module-level __getattr__ for lazy core module loading.
Evo2SAE core inference engine bionemo-recipes/.../evo2/src/evo2_sae/core.py	Main Evo2SAE class loads truncated Evo2 model + SAE checkpoint, supports tokenization, single/batch encoding to SAE codes, feature extraction, top-k feature selection, and generation with optional decode-time SAE steering via delta injection.
Shared FASTA parsing utility bionemo-recipes/.../evo2/src/evo2_sae/fasta.py	Streaming read_fasta() reader transparently supports plain and gzip-compressed FASTA, yields (seq_id, sequence) tuples, and auto-generates sequential IDs for headerless records.
FASTA integration in chunk script bionemo-recipes/.../evo2/scripts/chunk_fasta.py	Updated to use shared read_fasta() from evo2_sae.fasta instead of local parse_fasta().
Evo2 SAE recipe integration tests bionemo-recipes/.../evo2/tests/test_steering.py	CPU test validates clamp-hook arithmetic. GPU tests verify encode produces finite positive codes, unsteered generation produces valid DNA (ACGTN), and steering changes continuation deterministically.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

• NVIDIA-BioNeMo/bionemo-framework#1621: Both PRs modify the Evo2 recipe's FASTA chunking pipeline by refactoring parse_fasta logic in scripts/chunk_fasta.py; this PR removes the local parser in favor of a shared evo2_sae.fasta.read_fasta utility.

Suggested labels

ciflow:all

Suggested reviewers

• jstjohn
• pstjohn
• jwilber
• trvachov

---

🐰 A SAE hook so clever, it clamps with a delta,
Evo2 now steers genes with precision so bright—
Encoding and clamping, a feature to frame,
Generation refined: DNA shaped just right! 🧬

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 52.83% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title captures the main engine and steering work, though it is broader than the actual diff.
Description check	✅ Passed	The PR description is detailed and covers scope, architecture, usage, and testing, though it does not fully follow the repository template sections.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch pbinder/evo2-sae-serve

---

_{Comment @coderabbitai help to get the list of available commands.}

[jwilber]

Have you tried this? any examples you can share/screenshots?

[polinabinder1]

@jwilber This only deals with the steering backend. The visualization is in PR 1623.

[polinabinder1]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

Actionable comments posted: 10

🧹 Nitpick comments (8)

bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py (1)

37-50: ⚖️ Poor tradeoff

Consider more portable default paths.

Similar to the shell script, the default checkpoint and annotation paths are hardcoded to /data/interp/evo2/... which won't exist for other users. While these can be overridden via CLI arguments or environment variables (making this less critical than the shell script issue), consider removing these hardcoded defaults or documenting the required setup clearly.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py`
around lines 37 - 50, Default file paths for CLI args (--sae-ckpt-path,
--feature-annotations and EVO2_CKPT_DIR env fallback) are hardcoded to
/data/interp/evo2/...; remove or replace these with portable defaults by making
the argparse defaults None (or point to a user/home-relative path) and rely on
environment variables (SAE_CKPT_PATH, FEATURE_ANNOTATIONS, EVO2_CKPT_DIR) or
explicit CLI input, and update the code that consumes these values (where these
args are referenced) to validate and raise a clear error if no path is provided;
target the add_argument calls for "--sae-ckpt-path", "--feature-annotations" and
the EVO2_CKPT_DIR default.

bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py (7)

51-55: ⚡ Quick win

Add field docstrings to FeatureClamp.

📝 Example enhancement

 class FeatureClamp(BaseModel):
     """A single SAE-feature steering clamp (feature id + target strength)."""
 
-    feature_id: int
-    strength: float = 1.0
+    feature_id: int
+    """SAE feature ID to clamp during generation."""
+    strength: float = 1.0
+    """Target activation strength for the feature."""

As per coding guidelines, use Google-style docstrings (pydocstyle convention) in Python code.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`
around lines 51 - 55, Add Google-style (pydocstyle) docstrings describing each
field on the FeatureClamp Pydantic model: update the class docstring for
FeatureClamp (subclassing BaseModel) to include an Args section documenting
feature_id (int) and strength (float) with concise descriptions and
units/semantics (e.g., feature index and target steering strength, default 1.0).
Keep the top-line summary intact and ensure the Args block follows Google style
so linters accept it.

Source: Coding guidelines

---

58-68: ⚡ Quick win

Add field docstrings to GenerateRequest.

📝 Example enhancement

 class GenerateRequest(BaseModel):
     """Request body for /generate (autoregressive generation + optional SAE-feature clamps)."""
 
-    prompt: str = ""
-    organism: str = "None (raw DNA)"
-    tag: Optional[str] = None
-    features: list[FeatureClamp] = []
-    n_tokens: int = 120
-    temperature: float = 1.0
-    top_k: int = 0
-    compare_baseline: bool = False
+    prompt: str = ""
+    """Initial DNA sequence to condition generation."""
+    organism: str = "None (raw DNA)"
+    """Organism identifier for phylogenetic tagging."""
+    tag: Optional[str] = None
+    """Custom phylogenetic tag (overrides organism lookup)."""
+    features: list[FeatureClamp] = []
+    """SAE feature clamps for steering generation."""
+    n_tokens: int = 120
+    """Number of tokens to generate."""
+    temperature: float = 1.0
+    """Sampling temperature (higher = more random)."""
+    top_k: int = 0
+    """Top-k sampling parameter (0 = disabled)."""
+    compare_baseline: bool = False
+    """Whether to generate an unsteered baseline for comparison."""

As per coding guidelines, use Google-style docstrings (pydocstyle convention) in Python code.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`
around lines 58 - 68, Add Google-style (pydocstyle) docstrings for the
GenerateRequest datamodel: add a class docstring describing the purpose of
GenerateRequest and include an Args section that documents each attribute
(prompt, organism, tag, features: list[FeatureClamp], n_tokens, temperature,
top_k, compare_baseline) with types and brief descriptions (e.g., prompt: input
sequence string; organism: organism context or "None (raw DNA)"; tag: optional
user tag; features: SAE FeatureClamp list used for clamping; n_tokens: number of
tokens to generate; temperature: sampling temperature; top_k: top-k sampling
value; compare_baseline: whether to compare to baseline). Ensure the formatting
follows Google-style pydocstyle conventions and place the docstring immediately
under the class GenerateRequest declaration.

Source: Coding guidelines

---

39-48: ⚡ Quick win

Add field docstrings to AnnotateRequest.

The class is missing Google-style field docstrings. Each field should document its purpose, especially fields like mode that have specific allowed values ("topk" | "pick").

📝 Example enhancement

 class AnnotateRequest(BaseModel):
     """Request body for /annotate (top-k feature scan or an explicit feature pick)."""
 
-    sequence: str
-    organism: str = "None (raw DNA)"
-    tag: Optional[str] = None
-    mode: str = "topk"  # "topk" | "pick"
-    k: int = 8
-    feature_ids: Optional[list[int]] = None
-    feature_id: Optional[int] = None
+    sequence: str
+    """DNA sequence to annotate."""
+    organism: str = "None (raw DNA)"
+    """Organism identifier for phylogenetic tagging."""
+    tag: Optional[str] = None
+    """Custom phylogenetic tag (overrides organism lookup)."""
+    mode: str = "topk"
+    """Feature selection mode: 'topk' (top-k scan) or 'pick' (explicit features)."""
+    k: int = 8
+    """Number of top features to return when mode='topk'."""
+    feature_ids: Optional[list[int]] = None
+    """Explicit feature IDs when mode='pick'."""
+    feature_id: Optional[int] = None
+    """Single feature ID when mode='pick' (alternative to feature_ids)."""

As per coding guidelines, use Google-style docstrings (pydocstyle convention) in Python code.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`
around lines 39 - 48, The AnnotateRequest Pydantic model lacks Google-style
field docstrings; update the class docstring for AnnotateRequest to include a
Google-style "Attributes:" section that documents each field (sequence,
organism, tag, mode, k, feature_ids, feature_id), describing purpose,
types/constraints and allowed values for mode ("topk" | "pick") and any
relationships (e.g., feature_ids vs feature_id) so readers and linters can
validate the field meanings. Ensure the docstring follows pydocstyle/Google
conventions and mentions defaults where relevant.

Source: Coding guidelines

---

99-107: ⚡ Quick win

Add return type hint to features endpoint.

     `@app.get`("/features")
-    def features():
+    def features() -> list[dict]:

As per coding guidelines, use Pyright for type checking in Python files following pyproject.toml configuration.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`
around lines 99 - 107, The endpoint function features lacks a return type hint;
update its signature to include a typed return such as def features() ->
List[Dict[str, Any]]: and add the necessary imports (from typing import List,
Dict, Any) at the top of the module, or alternatively define and use a pydantic
model and set response_model on `@app.get`; modify the function signature and
imports so Pyright type checking passes while keeping the existing logic in
features().

Source: Coding guidelines

---

109-154: ⚡ Quick win

Add return type hint to annotate endpoint.

     `@app.post`("/annotate")
-    def annotate(req: AnnotateRequest):
+    def annotate(req: AnnotateRequest) -> dict:

As per coding guidelines, use Pyright for type checking in Python files following pyproject.toml configuration.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`
around lines 109 - 154, The annotate endpoint lacks a return type hint which
fails Pyright checks; update the annotate function signature (def annotate(req:
AnnotateRequest)) to include an explicit return type like -> Dict[str, Any] (or
a proper TypedDict/AnnotateResponse if available), and add the corresponding
typing import (e.g., from typing import Dict, Any) at the top of the module so
Pyright accepts the annotated return for the function annotate and its returned
JSON structure.

Source: Coding guidelines

---

86-97: ⚡ Quick win

Add return type hint to health endpoint.

     `@app.get`("/health")
-    def health():
+    def health() -> dict:

As per coding guidelines, use Pyright for type checking in Python files following pyproject.toml configuration.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`
around lines 86 - 97, The health endpoint lacks a return type hint; update the
health function signature (def health) to declare a typed return such as ->
Dict[str, Any] or -> dict[str, Any] and add the corresponding import (from
typing import Any, Dict) so Pyright can validate the returned mapping built from
engine (engine.ready, engine.layer, engine.n_features, engine.labels,
engine.sae_ckpt_path, engine.organism_tags, engine.device); keep the returned
structure unchanged and ensure the type hint covers the mixed value types.

Source: Coding guidelines

---

156-172: ⚡ Quick win

Add return type hint to generate endpoint.

     `@app.post`("/generate")
-    def generate(req: GenerateRequest):
+    def generate(req: GenerateRequest) -> dict:

As per coding guidelines, use Pyright for type checking in Python files following pyproject.toml configuration.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`
around lines 156 - 172, Add an explicit return type annotation to the FastAPI
endpoint function generate (def generate(req: GenerateRequest) -> Any) and
import Any from typing; update the signature so Pyright knows the endpoint's
return type (e.g., def generate(req: GenerateRequest) -> Any:), leaving the body
and exception handling (engine.generate call and HTTPException raises)
unchanged.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/scripts/launch_inference.sh`:
- Around line 17-21: The script currently embeds development-only absolute
defaults for VENV, EVO2_CKPT_DIR, SAE_CKPT_PATH, and FEATURE_ANNOTATIONS which
will break elsewhere; remove those hardcoded paths and instead either (a) set
VENV to a relative default like RECIPE_DIR/.venv and leave
EVO2_CKPT_DIR/SAE_CKPT_PATH/FEATURE_ANNOTATIONS unset, or (b) require these env
vars be provided and add an explicit validation block that checks VENV,
EVO2_CKPT_DIR, SAE_CKPT_PATH, and FEATURE_ANNOTATIONS (while allowing
EMBEDDING_LAYER to keep a sane numeric default), and if any are missing print a
clear error naming the missing variable(s) and exit non‑zero; update the code
references to VENV, EVO2_CKPT_DIR, SAE_CKPT_PATH, FEATURE_ANNOTATIONS, and
EMBEDDING_LAYER accordingly.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py`:
- Around line 57-67: Add a Google-style docstring to the _engine function
describing its purpose, parameters, and return value: explain that _engine
constructs and returns an Evo2SAE instance, document each parameter passed to
Evo2SAE (evo2_ckpt_dir, sae_ckpt_path, layer, device, max_seq_len,
feature_annotations) with types and brief descriptions, and state the return
type (Evo2SAE). Place the docstring immediately below the def _engine(args):
line using the standard Google style (Args:, Returns:) so tools and linters can
pick it up.
- Around line 34-55: The function _add_common is missing a Google-style
docstring; add a concise Google-style docstring immediately below the def
_add_common(p: argparse.ArgumentParser) -> None: line describing the function’s
purpose (registers shared CLI arguments), the parameter p (an
argparse.ArgumentParser), and any side effects/returns (modifies the parser in
place, returns None). Use the Google docstring sections: Args and Returns, and
keep wording aligned with surrounding code style.
- Around line 70-87: Add a Google-style docstring to _read_fasta describing
parameters (path), return values (ids, seqs), behavior (supports gzipped files)
and exceptions; and fix the header-parsing edge case by replacing the brittle
line that does line[1:].split()[0] with logic that strips the leading ">" and
whitespace, uses .split() safely (e.g., parts = line[1:].strip().split(); name =
parts[0] if parts else f"seq_{len(ids)}") so headers like "> " don't raise
IndexError and still produce a generated id when no token is present.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/core.py`:
- Around line 168-189: The docstring and logic in the method that reads
self.feature_annotations (variables: labels, peaks, path, path.suffix) claim to
support parquet/tsv/csv/json but only handle parquet; update the code in the
function (the block starting with path = Path(self.feature_annotations)) to
detect .csv/.tsv (use csv or pandas.read_csv), .json (json.load or
pandas.read_json), and parse the same columns ("feature_id", "label" or
"annotation", "max_activation") into labels and peaks just like the parquet
branch, and for any other suffix emit an explicit logger.warning stating the
format is unsupported and return empty labels/peaks; ensure you reuse the same
keys/behavior (casting ids to int, labels to str, peaks to float) as done in the
pq branch so the rest of the code remains compatible.
- Around line 352-366: The code indexes SAE tensors using incoming feature IDs
(see fids, features and usages of self.sae.encoder.weight /
self.sae.decoder.weight) without validation; add explicit bounds and type checks
before any tensor indexing inside the block that builds specs (validate each fid
is an integer >=0 and < self.sae.encoder.weight.size(0) and similarly valid for
decoder indexing), and if invalid raise a ValueError with a clear message so the
/generate handler returns 400; perform these checks at the start of the with
self._lock block (before accessing self.sae.* tensors) or filter/convert
f["feature_id"] to int safely and validate before using it in specs construction
(references: fids, features, self.sae.encoder.weight, self.sae.decoder.weight,
self.layer).

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`:
- Around line 124-131: The code currently treats any non-"pick" mode as "topk";
update the conditional around req.mode in server.py to explicitly handle "pick"
and "topk" only and raise an HTTPException(400, "invalid mode; allowed values:
'pick', 'topk'") for any other value. Concretely, change the if/else to if
req.mode == "pick": ... elif req.mode == "topk": compute k and call
engine.top_features(...); else: raise the 400 error so typos or unsupported
modes are rejected (refer to req.mode, engine.top_features, chosen).
- Line 84: The CORS middleware is currently set to allow all origins via
app.add_middleware(CORSMiddleware, allow_origins=["*"]) which is too permissive
for production; update the server startup to read an environment variable (e.g.,
CORS_ALLOWED_ORIGINS or CORS_ALLOWED_ORIGIN) and use that to populate
allow_origins (parse a comma-separated list into a list), defaulting to a safe
value like an empty list or localhost for dev, and ensure
allow_methods/allow_headers remain appropriate; locate the use of
app.add_middleware and replace the hardcoded ["*"] with the parsed config so
deployments can restrict origins without code changes.
- Around line 23-34: Add a second blank line after the import block (the line
ending with "from .core import Evo2SAE, clean_dna") so there are two blank lines
before the next top-level statement (e.g., the logger = logging.getLogger(...)
or any subsequent definitions); this aligns with the isort rule and ensures the
import section (including Evo2SAE and clean_dna) is separated from module-level
code.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/tests/test_server.py`:
- Around line 99-105: Update the test_endpoints_503_until_ready to also assert
that the /generate endpoint returns 503 when the engine is not ready: in the
existing test that creates FakeEngine (eng.ready = False), using
TestClient(build_app(eng)) add a POST request to "/generate" with a
representative JSON payload (similar shape to other tests, e.g. prompt/sequence
fields) and assert c.post("/generate", json=...).status_code == 503 so /generate
is covered like /features and /annotate.

---

Nitpick comments:
In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py`:
- Around line 37-50: Default file paths for CLI args (--sae-ckpt-path,
--feature-annotations and EVO2_CKPT_DIR env fallback) are hardcoded to
/data/interp/evo2/...; remove or replace these with portable defaults by making
the argparse defaults None (or point to a user/home-relative path) and rely on
environment variables (SAE_CKPT_PATH, FEATURE_ANNOTATIONS, EVO2_CKPT_DIR) or
explicit CLI input, and update the code that consumes these values (where these
args are referenced) to validate and raise a clear error if no path is provided;
target the add_argument calls for "--sae-ckpt-path", "--feature-annotations" and
the EVO2_CKPT_DIR default.

In
`@bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py`:
- Around line 51-55: Add Google-style (pydocstyle) docstrings describing each
field on the FeatureClamp Pydantic model: update the class docstring for
FeatureClamp (subclassing BaseModel) to include an Args section documenting
feature_id (int) and strength (float) with concise descriptions and
units/semantics (e.g., feature index and target steering strength, default 1.0).
Keep the top-line summary intact and ensure the Args block follows Google style
so linters accept it.
- Around line 58-68: Add Google-style (pydocstyle) docstrings for the
GenerateRequest datamodel: add a class docstring describing the purpose of
GenerateRequest and include an Args section that documents each attribute
(prompt, organism, tag, features: list[FeatureClamp], n_tokens, temperature,
top_k, compare_baseline) with types and brief descriptions (e.g., prompt: input
sequence string; organism: organism context or "None (raw DNA)"; tag: optional
user tag; features: SAE FeatureClamp list used for clamping; n_tokens: number of
tokens to generate; temperature: sampling temperature; top_k: top-k sampling
value; compare_baseline: whether to compare to baseline). Ensure the formatting
follows Google-style pydocstyle conventions and place the docstring immediately
under the class GenerateRequest declaration.
- Around line 39-48: The AnnotateRequest Pydantic model lacks Google-style field
docstrings; update the class docstring for AnnotateRequest to include a
Google-style "Attributes:" section that documents each field (sequence,
organism, tag, mode, k, feature_ids, feature_id), describing purpose,
types/constraints and allowed values for mode ("topk" | "pick") and any
relationships (e.g., feature_ids vs feature_id) so readers and linters can
validate the field meanings. Ensure the docstring follows pydocstyle/Google
conventions and mentions defaults where relevant.
- Around line 99-107: The endpoint function features lacks a return type hint;
update its signature to include a typed return such as def features() ->
List[Dict[str, Any]]: and add the necessary imports (from typing import List,
Dict, Any) at the top of the module, or alternatively define and use a pydantic
model and set response_model on `@app.get`; modify the function signature and
imports so Pyright type checking passes while keeping the existing logic in
features().
- Around line 109-154: The annotate endpoint lacks a return type hint which
fails Pyright checks; update the annotate function signature (def annotate(req:
AnnotateRequest)) to include an explicit return type like -> Dict[str, Any] (or
a proper TypedDict/AnnotateResponse if available), and add the corresponding
typing import (e.g., from typing import Dict, Any) at the top of the module so
Pyright accepts the annotated return for the function annotate and its returned
JSON structure.
- Around line 86-97: The health endpoint lacks a return type hint; update the
health function signature (def health) to declare a typed return such as ->
Dict[str, Any] or -> dict[str, Any] and add the corresponding import (from
typing import Any, Dict) so Pyright can validate the returned mapping built from
engine (engine.ready, engine.layer, engine.n_features, engine.labels,
engine.sae_ckpt_path, engine.organism_tags, engine.device); keep the returned
structure unchanged and ensure the type hint covers the mixed value types.
- Around line 156-172: Add an explicit return type annotation to the FastAPI
endpoint function generate (def generate(req: GenerateRequest) -> Any) and
import Any from typing; update the signature so Pyright knows the endpoint's
return type (e.g., def generate(req: GenerateRequest) -> Any:), leaving the body
and exception handling (engine.generate call and HTTPException raises)
unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3e499152-13aa-449c-b7ad-6b67a8279836

📥 Commits

Reviewing files that changed from the base of the PR and between e407165 and de81106.

📒 Files selected for processing (8)

• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/pyproject.toml
• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/scripts/launch_inference.sh
• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/__init__.py
• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/cli.py
• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/core.py
• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/src/evo2_sae/server.py
• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/tests/test_server.py
• bionemo-recipes/interpretability/sparse_autoencoders/recipes/evo2/tests/test_steering.py

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.