Skip to content

ESM Encoder Class Creation#37

Merged
vratins merged 25 commits into
mainfrom
dev_esm_encoder
Mar 18, 2026
Merged

ESM Encoder Class Creation#37
vratins merged 25 commits into
mainfrom
dev_esm_encoder

Conversation

@vratins

@vratins vratins commented Feb 26, 2026

Copy link
Copy Markdown
Contributor
  • 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

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 ESMEncoder that returns precomputed ESM embeddings as scalar node features (no vectors, no edges).
  • Updates BaseProteinEncoder.forward to return (s, V, pp_edge_attr) and updates encoder/unit tests accordingly.
  • Adds use_edge_update option 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.

Comment thread src/slae_encoder.py Outdated
Comment thread src/esm_encoder.py Outdated
Comment thread src/gvp_encoder.py
Comment thread src/__init__.py Outdated
Comment thread src/gvp_encoder.py
Comment thread src/encoder_base.py
Comment thread src/encoder_base.py Outdated
Comment thread src/encoder_base.py
Comment thread src/gvp_encoder.py Outdated
Comment thread src/gvp_encoder.py
Comment thread src/__init__.py Outdated
Comment thread src/encoder_base.py
Comment thread src/encoder_base.py
Comment thread src/encoder_base.py
Comment thread src/encoder_base.py Outdated
Comment thread src/gvp_encoder.py Outdated
Comment thread src/gvp_encoder.py Outdated
Comment thread src/gvp_encoder.py Outdated
Comment thread src/slae_encoder.py Outdated
Comment thread src/slae_encoder.py Outdated
@vratins

vratins commented Mar 16, 2026

Copy link
Copy Markdown
Contributor Author

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:

  • encoder_base.py
  • gvp_encoder.py

@vratins vratins mentioned this pull request Mar 16, 2026

@marcuscollins marcuscollins left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lots of comments, but nothing blocking. Please go through and either address them or make issues to track any requests.

Comment thread .github/workflows/build.yml
Comment thread .github/workflows/build.yml
Comment thread src/__init__.py
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,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

None of these as statements are necessary, please remove them.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would lead to unused import issues when linting -- do you recommend adding noqa comments instead?

Comment thread src/encoder_base.py Outdated
Comment thread src/encoder_base.py Outdated
Comment thread tests/test_encoder.py Outdated
Comment thread tests/test_encoder.py Outdated
Comment thread tests/test_encoder.py
Comment thread src/gvp_encoder.py Outdated
Comment thread src/gvp_encoder.py Outdated
@vratins
vratins merged commit ea8ba8b into main Mar 18, 2026
3 checks passed
@vratins vratins linked an issue Mar 18, 2026 that may be closed by this pull request
@vratins
vratins deleted the dev_esm_encoder branch March 19, 2026 15:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add a toggle for edge updates in the GVP fix relative imports use loguru for logging instead of print statements

5 participants