Skip to content

feat: Ease MPNN command line usage #367

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
roccomoretti wants to merge 2 commits into
base: production
Choose a base branch
from
Open

feat: Ease MPNN command line usage #367

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2f54bd4
feat: Ease MPNN command line usage
roccomoretti Jul 29, 2026
f55267d
Ruff format
roccomoretti Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 8 additions & 7 deletions models/mpnn/src/mpnn/inference_engines/mpnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,23 +50,24 @@ def __init__(
device: str | torch.device | None = None,
):
# Store raw configuration
self.model_type = model_type
self.is_legacy_weights = is_legacy_weights
self.out_directory = out_directory
self.write_fasta = write_fasta
self.write_structures = write_structures

# allow null for checkpoint path when foundry-installed
# TODO: Currently this assumes the model type is the key in the registered path. Rework needed
self.checkpoint_path = (
str(
REGISTERED_CHECKPOINTS[
self.model_type.replace("_", "")
].get_default_path()
)
str(REGISTERED_CHECKPOINTS[model_type.replace("_", "")].get_default_path())
if not checkpoint_path
else checkpoint_path
)
# The default weights sets are all legacy types
self.is_legacy_weights = True if not checkpoint_path else is_legacy_weights
# the soluble_mpnn type just changes the default weights -- otherwise it behaves like regular protein_mpnn
if model_type == "soluble_mpnn":
self.model_type = "protein_mpnn"
else:
self.model_type = model_type

# Determine the device (supports XPU, CUDA, and CPU).
if device is not None:
Expand Down
11 changes: 2 additions & 9 deletions models/mpnn/src/mpnn/utils/inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
parser.add_argument(
"--model_type",
type=str,
choices=["protein_mpnn", "ligand_mpnn"],
choices=["protein_mpnn", "ligand_mpnn", "soluble_mpnn"],
help="Model type to use.",
default=MPNN_GLOBAL_INFERENCE_DEFAULTS["model_type"],
)
Expand Down Expand Up @@ -591,17 +591,10 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
return json.load(f)

# Build a single-input JSON object from CLI
if (
args.model_type is None
or args.checkpoint_path is None
or args.is_legacy_weights is None
or args.structure_path is None
):
if args.model_type is None or args.structure_path is None:
raise ValueError(
"When --config_json is not provided, "
"--model_type, "
"--checkpoint_path, "
"--is_legacy_weights, "
"--structure_path "
"must all be specified."
)
Expand Down
Loading