ESM Encoder Class Creation#37
Conversation
vratins
commented
Feb 26, 2026
- Add ESMEncoder for pre-computed ESM3 embeddings
- Update encoder interface to return edge features as third output
- Add use_edge_update parameter to GVPEncoder for optional edge updates
- Use constants for defaults
- Update tests
…butes, using constants, fixing tests
There was a problem hiding this comment.
Pull request overview
Adds a new cached-embedding encoder (ESM) and extends the encoder interface so encoders can optionally return protein–protein edge features, with GVP edge-updates made configurable.
Changes:
- Introduces
ESMEncoderthat returns precomputed ESM embeddings as scalar node features (no vectors, no edges). - Updates
BaseProteinEncoder.forwardto return(s, V, pp_edge_attr)and updates encoder/unit tests accordingly. - Adds
use_edge_updateoption and cached-edge-feature passthrough/computation changes in the GVP encoder stack.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
tests/test_encoder.py |
Updates tests for the new 3-output encoder interface and adds ESM encoder coverage. |
src/encoder_base.py |
Changes the abstract encoder interface to return optional PP edge features as a third output. |
src/gvp_encoder.py |
Adds use_edge_update, supports cached edge features, and returns edge features optionally. |
src/slae_encoder.py |
Updates SLAE encoder to comply with the new interface (returns None for edge features). |
src/esm_encoder.py |
Adds the new ESM cached-embedding encoder implementation and registration. |
src/__init__.py |
Updates package-level imports to trigger registration of the new/renamed encoder modules. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
In the last commit I made slight formatting changes to files to ensure that ruff/ty/build checks pass. The only files where there are comments that are addressed and functional changes to review are:
|
marcuscollins
left a comment
There was a problem hiding this comment.
Lots of comments, but nothing blocking. Please go through and either address them or make issues to track any requests.
| from src.encoder_base import BaseProteinEncoder, build_encoder, register_encoder | ||
| from src import gvp_encoder as gvp_encoder | ||
| from src.encoder_base import ( | ||
| BaseProteinEncoder as BaseProteinEncoder, |
There was a problem hiding this comment.
None of these as statements are necessary, please remove them.
There was a problem hiding this comment.
This would lead to unused import issues when linting -- do you recommend adding noqa comments instead?
… is created for ty check
…into dev_esm_encoder
…into dev_esm_encoder