feat: Ease MPNN command line usage

There were checks which prevented people from omitting --checkpoint_path when not using --config_json.
Additionally, simplify use of SolubleMPNN with the default parameter set.

Author: roccomoretti

models/mpnn/src/mpnn/inference_engines/mpnn.py

@@ -50,8 +50,6 @@ 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
@@ -61,12 +59,23 @@ def __init__(
         self.checkpoint_path = (
             str(
                 REGISTERED_CHECKPOINTS[
-                    self.model_type.replace("_", "")
+                    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:

models/mpnn/src/mpnn/utils/inference.py

@@ -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"],
     )
@@ -593,15 +593,11 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
     # 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
     ):
         raise ValueError(
             "When --config_json is not provided, "
             "--model_type, "
-            "--checkpoint_path, "
-            "--is_legacy_weights, "
             "--structure_path "
             "must all be specified."
         )