From 82b4ba4d9e0b4378e2ded02258853e2fe02849d5 Mon Sep 17 00:00:00 2001
From: Rocco Moretti <rmorettiase@gmail.com>
Date: Wed, 29 Jul 2026 18:14:27 -0500
Subject: [PATCH 1/4] feat: MPNN options: multistruct and json/option combo

For MPNN, add the ability to specify multiple input structures
on the command line. --structure_path can now take multiple
inputs (e.g. `--structure_path *.cif`)

Additionally, add the ability to combine the --config_json setting
with command line parameters.

The philosophy is that --config_json can be set up with "standard"
settings, and then command line parameters
(most notably --structure_path) can then be used to launch multiple
different runs from the same JSON file settings.

Right now the settings in config.json take presidence, and the settings
on the command line won't overwrite them. Command line parameters are
only used for filling in values missing from config.json.
(Not sure if this is the best approach, but it does mirror
the current approch of ignoring the command line completely if
--config_json is set.)

For --structure_path, if multiple structures are provided, any block in
the JSON `inputs` setting which does not set `structure_path` will be
duplicated for each structure. (If no block is provided, a default block
will be generated.)

Additionally, to better harmonize JSON/command line usage, JSON-like
settings of "true"/"false"/"null" with command line parameters are now
treated like their Python "True"/"False"/"None" equivalents.

Note that technically this could be considered a breaking change,
as command lines with both --config_json and other parameters will
not behave the same. (As would command line with the JSON-like values.)
I don't think this is a big issue, as those command lines weren't really
well-formed to start with.
---
 models/mpnn/src/mpnn/utils/inference.py     | 260 +++++++++++++-------
 models/mpnn/tests/test_inference_helpers.py |   5 +-
 models/mpnn/tests/test_inference_utils.py   |  53 +++-
 3 files changed, 227 insertions(+), 91 deletions(-)

diff --git a/models/mpnn/src/mpnn/utils/inference.py b/models/mpnn/src/mpnn/utils/inference.py
index 19092aa5..225f0bc7 100644
--- a/models/mpnn/src/mpnn/utils/inference.py
+++ b/models/mpnn/src/mpnn/utils/inference.py
@@ -97,10 +97,13 @@
 
 
 def str2bool(v: str) -> bool:
-    """Helper function to parse boolean CLI args."""
-    if v in ("True", "1"):
+    """Helper function to parse boolean CLI args.
+
+    Accept Python, JSON or C++ style values
+    """
+    if v in ("True", "true", "1"):
         return True
-    elif v in ("False", "0"):
+    elif v in ("False", "false", "0"):
         return False
     else:
         raise argparse.ArgumentTypeError(f"Boolean value expected, got {v!r}")
@@ -111,11 +114,20 @@ def none_or_type(v: Any, specified_type: Callable[[Any], Any]) -> Any | None:
     CLI type parser that turns 'None' into None. Otherwise, returns the value
     cast to the given type. This function is useful for the parser/pipeline
     override arguments where None has a special meaning (use default behavior).
+
+    Accept both Python like 'None' and JSON-like 'null'.
     """
-    if v == "None":
+    if v == "None" or v == "null": # Python-like or JSON like None/null value.
         return None
     return specified_type(v)
 
+class TrackUserSetting(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        setattr(namespace, self.dest, values)
+
+        if not hasattr(namespace, 'user_set'):
+            setattr(namespace, "user_set", {})
+        namespace.user_set[self.dest] = True
 
 def build_arg_parser() -> argparse.ArgumentParser:
     """Build the MPNN inference arg parser."""
@@ -133,6 +145,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "flags are parsed but ignored."
         ),
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["config_json"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Model Type and Weights ---------------- #
@@ -142,12 +155,14 @@ def build_arg_parser() -> argparse.ArgumentParser:
         choices=["protein_mpnn", "ligand_mpnn"],
         help="Model type to use.",
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["model_type"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--checkpoint_path",
         type=str,
         help="Path to model checkpoint.",
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["checkpoint_path"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--is_legacy_weights",
@@ -155,6 +170,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         choices=[True, False],
         help="Whether to interpret checkpoint as legacy-weight ordering.",
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["is_legacy_weights"],
+        action=TrackUserSetting,
     )
 
     # --------------- Output controls ---------------- #
@@ -163,6 +179,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         type=str,
         help="Output directory for CIF/FASTA.",
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["out_directory"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--write_fasta",
@@ -170,6 +187,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         choices=[True, False],
         help="Whether to write FASTA outputs.",
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["write_fasta"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--write_structures",
@@ -177,6 +195,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         choices=[True, False],
         help="Whether to write designed structures (CIF).",
         default=MPNN_GLOBAL_INFERENCE_DEFAULTS["write_structures"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Structure Path and Name ---------------- #
@@ -184,13 +203,16 @@ def build_arg_parser() -> argparse.ArgumentParser:
         "--structure_path",
         type=str,
         help="Path to structure file (CIF or PDB).",
+        nargs="+",
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["structure_path"],
+        action="extend", # Don't need to track usage
     )
     parser.add_argument(
         "--name",
         type=str,
         help="Optional name / label for the input.",
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["name"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Sampling Parameters ---------------- #
@@ -199,6 +221,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         type=int,
         help="Random seed for sampling.",
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["seed"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--batch_size",
@@ -208,12 +231,14 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "the effective repeat_sample_num passed to the pipeline."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["batch_size"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--number_of_batches",
         type=int,
         help="Number of batches of size batch_size to draw.",
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["number_of_batches"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Parser overrides ---------------- #
@@ -227,6 +252,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "use the parser default behavior."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["remove_ccds"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--remove_waters",
@@ -238,6 +264,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "has special behavior: use the parser default behavior."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["remove_waters"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Pipeline Setup Overrides ---------------- #
@@ -249,6 +276,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "has special behavior: use the pipeline default behavior."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["occupancy_threshold_sidechain"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--occupancy_threshold_backbone",
@@ -258,6 +286,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "has special behavior: use the pipeline default behavior."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["occupancy_threshold_backbone"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--undesired_res_names",
@@ -268,6 +297,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "pipeline default behavior."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["undesired_res_names"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Scalar User Settings ---------------- #
@@ -276,6 +306,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         type=float,
         help=("Structure noise (Angstroms) used in user settings."),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["structure_noise"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--decode_type",
@@ -290,6 +321,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "for all previous positions when predicting each residue."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["decode_type"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--causality_pattern",
@@ -315,6 +347,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "itself (as a destination node)."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["causality_pattern"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--initialize_sequence_embedding_with_ground_truth",
@@ -334,6 +367,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS[
             "initialize_sequence_embedding_with_ground_truth"
         ],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--features_to_return",
@@ -344,6 +378,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             '["mask_for_loss"], "decoder_features": ["log_probs"]}\''
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["features_to_return"],
+        action=TrackUserSetting,
     )
     # Only applicable for LigandMPNN.
     parser.add_argument(
@@ -355,6 +390,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "for LigandMPNN."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["atomize_side_chains"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Design scope (mutually exclusive) ---------------- #
@@ -366,6 +402,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             'List of residue IDs to fix: e.g. \'["A35","B40","C52"]\' or "A35,B40,C52"'
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["fixed_residues"],
+        action=TrackUserSetting,
     )
     design_group.add_argument(
         "--designed_residues",
@@ -375,18 +412,21 @@ def build_arg_parser() -> argparse.ArgumentParser:
             'e.g. \'["A35","B40","C52"]\' or "A35,B40,C52"'
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["designed_residues"],
+        action=TrackUserSetting,
     )
     design_group.add_argument(
         "--fixed_chains",
         type=str,
         help=('List of chain IDs to fix: e.g. \'["A","B"]\' or "A,B"'),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["fixed_chains"],
+        action=TrackUserSetting,
     )
     design_group.add_argument(
         "--designed_chains",
         type=str,
         help=('List of chain IDs to design: e.g. \'["A","B"]\' or "A,B"'),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["designed_chains"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Bias, Omission, and Pair Bias ---------------- #
@@ -395,6 +435,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         type=str,
         help='Bias dict: e.g. \'{"ALA": -1.0, "GLY": 0.5}\'',
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["bias"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--bias_per_residue",
@@ -403,12 +444,14 @@ def build_arg_parser() -> argparse.ArgumentParser:
             'Per-residue bias dict: e.g. \'{"A35": {"ALA": -2.0}}\'. Overwrites --bias.'
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["bias_per_residue"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--omit",
         type=str,
         help=('List of residue types to omit: e.g. \'["ALA","GLY","UNK"]\'.'),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["omit"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--omit_per_residue",
@@ -418,6 +461,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             'e.g. \'{"A35": ["ALA","GLY","UNK"]}\'. Overwrites --omit.'
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["omit_per_residue"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--pair_bias",
@@ -428,6 +472,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             '\'{"ALA": {"GLY": -0.5}, "GLY": {"ALA": -0.5}}\''
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["pair_bias"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--pair_bias_per_residue_pair",
@@ -445,6 +490,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             'is the innermost dict (e.g. {"GLY": -1.0} ).'
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["pair_bias_per_residue_pair"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Temperature ---------------- #
@@ -453,6 +499,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         type=float,
         help=("Temperature for sampling."),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["temperature"],
+        action=TrackUserSetting,
     )
     parser.add_argument(
         "--temperature_per_residue",
@@ -462,6 +509,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "--temperature."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["temperature_per_residue"],
+        action=TrackUserSetting,
     )
 
     # ---------------- Symmetry ---------------- #
@@ -475,6 +523,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             '\'[["A35","B35"],["A40","B40","C40"]]\''
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["symmetry_residues"],
+        action=TrackUserSetting,
     )
     sym_group.add_argument(
         "--homo_oligomer_chains",
@@ -487,6 +536,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             '\'[["A","B","C"]]\''
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["homo_oligomer_chains"],
+        action=TrackUserSetting,
     )
 
     # Symmetry weights
@@ -500,6 +550,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
             "Ignored if homo_oligomer_chains is used."
         ),
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["symmetry_residues_weights"],
+        action=TrackUserSetting,
     )
 
     return parser
@@ -578,96 +629,135 @@ def _absolute_path_or_none(path_str: str | None) -> str | None:
         return None
     return str(Path(path_str).expanduser().resolve())
 
+def safe_set_in_config(
+    config: dict[str, Any],
+    args: argparse.Namespace,
+    name: str,
+    param: str | None = None,
+    format: str | None = None,
+) -> None:
+    """
+    If the given entry of name isn't already set in the top-level of the config dictionary,
+    set it from the parameter in args, but only if it's been annotated as user-set.
+
+    Note that every option used with this function must be annotated with `action=TrackUserSetting`
+    """
+    if name in config:
+        return
+    if param is None:
+        param = name
+    if param in args.user_set and args.user_set[param]:
+        if format is None:
+            config[name] = getattr(args, param)
+        elif format == "json":
+            config[name] = parse_json_like(getattr(args, param))
+        elif format == "list":
+            config[name] = parse_list_like(getattr(args, param))
 
 def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
     """Convert CLI args into the top-level JSON config dict."""
-    # If a config JSON is provided, load and return it directly. Ignore the
-    # other CLI args.
+
+    # In case we haven't provided any of the user-set options, make sure we have a dictionary
+    if not hasattr(args, "user_set"):
+       setattr( args, "user_set", {} )
+
+    # If set, use the config_json as the template
     if args.config_json:
         config_path = _absolute_path_or_none(args.config_json)
         # args.config_json is truthy here, so the absolute path is never None.
         assert config_path is not None
         with open(config_path, "r") as f:
-            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
-    ):
-        raise ValueError(
-            "When --config_json is not provided, "
-            "--model_type, "
-            "--checkpoint_path, "
-            "--is_legacy_weights, "
-            "--structure_path "
-            "must all be specified."
-        )
+            config = json.load(f)
+        if not isinstance(config, dict):
+            raise TypeError("The top level of config_json must be a JSON object (dictionary).")
+    else:
+        # Bare-bones config -- we'll populate it from options/defaults
+        config = {}
 
-    config: dict[str, Any] = {
-        # Model Type and Weights
-        "model_type": args.model_type,
-        "checkpoint_path": args.checkpoint_path,
-        "is_legacy_weights": args.is_legacy_weights,
-        # Output controls
-        "out_directory": args.out_directory,
-        "write_fasta": args.write_fasta,
-        "write_structures": args.write_structures,
-        # Singleton inputs list (CLI only supports single input at a time).
-        "inputs": [
-            {
-                # Structure Path and Name
-                "structure_path": args.structure_path,
-                "name": args.name,
-                # Sampling Parameters
-                "seed": args.seed,
-                "batch_size": args.batch_size,
-                "number_of_batches": args.number_of_batches,
-                # Parser Overrides
-                "remove_ccds": parse_list_like(args.remove_ccds),
-                "remove_waters": args.remove_waters,
-                # Pipeline Setup Overrides
-                "occupancy_threshold_sidechain": args.occupancy_threshold_sidechain,
-                "occupancy_threshold_backbone": args.occupancy_threshold_backbone,
-                "undesired_res_names": parse_list_like(args.undesired_res_names),
-                # Scalar User Settings
-                "structure_noise": args.structure_noise,
-                "decode_type": args.decode_type,
-                "causality_pattern": args.causality_pattern,
-                "initialize_sequence_embedding_with_ground_truth": args.initialize_sequence_embedding_with_ground_truth,
-                "features_to_return": parse_json_like(args.features_to_return),
-                # Only applicable for LigandMPNN
-                "atomize_side_chains": args.atomize_side_chains,
-                # Design scope - if all None, design all residues
-                "fixed_residues": parse_list_like(args.fixed_residues),
-                "designed_residues": parse_list_like(args.designed_residues),
-                "fixed_chains": parse_list_like(args.fixed_chains),
-                "designed_chains": parse_list_like(args.designed_chains),
-                # Bias, Omission, and Pair Bias
-                "bias": parse_json_like(args.bias),
-                "bias_per_residue": parse_json_like(args.bias_per_residue),
-                "omit": parse_json_like(args.omit),
-                "omit_per_residue": parse_json_like(args.omit_per_residue),
-                "pair_bias": parse_json_like(args.pair_bias),
-                "pair_bias_per_residue_pair": parse_json_like(
-                    args.pair_bias_per_residue_pair
-                ),
-                # Temperature
-                "temperature": args.temperature,
-                "temperature_per_residue": parse_json_like(
-                    args.temperature_per_residue
-                ),
-                # Symmetry
-                "symmetry_residues": parse_json_like(args.symmetry_residues),
-                "symmetry_residues_weights": parse_json_like(
-                    args.symmetry_residues_weights
-                ),
-                "homo_oligomer_chains": parse_json_like(args.homo_oligomer_chains),
-            }
-        ],
-    }
+    # Global settings
+    # Model Type and Weights
+    safe_set_in_config(config, args, "model_type")
+    safe_set_in_config(config, args, "checkpoint_path")
+    safe_set_in_config(config, args, "is_legacy_weights")
+    # Output controls
+    safe_set_in_config(config, args, "out_directory")
+    safe_set_in_config(config, args, "write_fasta")
+    safe_set_in_config(config, args, "write_structures")
+
+    if "model_type" not in config or config["model_type"] is None:
+        raise ValueError("model_type must be specified.")
+
+    # Provide a skeleton input list, to (hopefully) be populated from --structure_path
+    if "inputs" in config and not isinstance(config["inputs"], list):
+        raise TypeError("The inputs entry of config_json must be a list.")
+    if "inputs" not in config or len(config["inputs"]) == 0:
+        config["inputs"] = [{}]
+
+    for entry in config["inputs"]:
+        if not isinstance(entry, dict):
+            raise TypeError("Each entry in the inputs list of config_json must be a JSON object (dictionary)")
+
+    # Find the template input dicts in the config (either from config_json, or the stub config) and populate them with --structure_path values
+    structure_free = [ d for d in config["inputs"] if "structure_path" not in d ]
+    if len(structure_free) != 0:
+        if len(args.structure_path) == 0:
+            raise ValueError("structure_path must be specified, either in the config_json or the command line")
+
+        # combinitorial assortment
+        new_inputs = []
+        for entry in structure_free:
+            for structure in args.structure_path:
+                new_input = copy.deepcopy( entry )
+                new_input["structure_path"] = structure
+                new_inputs.append( new_input )
+
+        # TODO: Add some sort of check/validation/adjustment for output naming in the combinitorial case
+        already_has_structure = [ d for d in config["inputs"] if "structure_path" in d ]
+        config["inputs"] = already_has_structure + new_inputs
+    elif args.structure_path and len(args.structure_path) != 0:
+        raise ValueError("Command line structure specified with --structure_path, but all provided input entries already have structures specified.")
+
+    for entry in config["inputs"]:
+        # Name
+        safe_set_in_config(entry, args, "name")
+        # Sampling Parameters
+        safe_set_in_config(entry, args, "seed")
+        safe_set_in_config(entry, args, "batch_size")
+        safe_set_in_config(entry, args, "number_of_batches")
+        # Parser Overrides
+        safe_set_in_config(entry, args, "remove_ccds", format="list")
+        safe_set_in_config(entry, args, "remove_waters")
+        # Pipeline Setup Overrides
+        safe_set_in_config(entry, args, "occupancy_threshold_sidechain")
+        safe_set_in_config(entry, args, "occupancy_threshold_backbone")
+        safe_set_in_config(entry, args, "undesired_res_names", format="list")
+        # Scalar User Settings
+        safe_set_in_config(entry, args, "structure_noise")
+        safe_set_in_config(entry, args, "decode_type")
+        safe_set_in_config(entry, args, "causality_pattern")
+        safe_set_in_config(entry, args, "initialize_sequence_embedding_with_ground_truth")
+        safe_set_in_config(entry, args, "features_to_return", format="json")
+        # Only applicable for LigandMPNN
+        safe_set_in_config(entry, args, "atomize_side_chains")
+        # Design scope - if all None, design all residues
+        safe_set_in_config(entry, args, "fixed_residues", format="list")
+        safe_set_in_config(entry, args, "designed_residues", format="list")
+        safe_set_in_config(entry, args, "fixed_chains", format="list")
+        safe_set_in_config(entry, args, "designed_chains", format="list")
+        # Bias, Omission, and Pair Bias
+        safe_set_in_config(entry, args, "bias", format="json")
+        safe_set_in_config(entry, args, "bias_per_residue", format="json")
+        safe_set_in_config(entry, args, "omit", format="json")
+        safe_set_in_config(entry, args, "omit_per_residue", format="json")
+        safe_set_in_config(entry, args, "pair_bias", format="json")
+        safe_set_in_config(entry, args, "pair_bias_per_residue_pair", format="json")
+        # Temperature
+        safe_set_in_config(entry, args, "temperature")
+        safe_set_in_config(entry, args, "temperature_per_residue", format="json")
+        # Symmetry
+        safe_set_in_config(entry, args, "symmetry_residues", format="json")
+        safe_set_in_config(entry, args, "symmetry_residues_weights", format="json")
+        safe_set_in_config(entry, args, "homo_oligomer_chains", format="json")
 
     return config
 
diff --git a/models/mpnn/tests/test_inference_helpers.py b/models/mpnn/tests/test_inference_helpers.py
index 48d13daf..7832094d 100644
--- a/models/mpnn/tests/test_inference_helpers.py
+++ b/models/mpnn/tests/test_inference_helpers.py
@@ -18,13 +18,13 @@
 
 
 @pytest.mark.parametrize(
-    ("value", "expected"), [("True", True), ("1", True), ("False", False), ("0", False)]
+    ("value", "expected"), [("True", True), ("true",True), ("1", True), ("False", False), ("false", False), ("0", False)]
 )
 def test_str2bool_valid(value: str, expected: bool):
     assert str2bool(value) is expected
 
 
-@pytest.mark.parametrize("value", ["true", "false", "yes", "", "2"])
+@pytest.mark.parametrize("value", ["yes", "", "2"])
 def test_str2bool_invalid_raises(value: str):
     with pytest.raises(argparse.ArgumentTypeError):
         str2bool(value)
@@ -32,6 +32,7 @@ def test_str2bool_invalid_raises(value: str):
 
 def test_none_or_type_sentinel_and_cast():
     assert none_or_type("None", int) is None
+    assert none_or_type("null", int) is None
     assert none_or_type("5", int) == 5
     assert none_or_type("1.5", float) == 1.5
 
diff --git a/models/mpnn/tests/test_inference_utils.py b/models/mpnn/tests/test_inference_utils.py
index 0c08ef4e..61310d3b 100644
--- a/models/mpnn/tests/test_inference_utils.py
+++ b/models/mpnn/tests/test_inference_utils.py
@@ -100,8 +100,10 @@ def _make_simple_inference_output() -> MPNNInferenceOutput:
     "value, expected",
     [
         ("True", True),
+        ("true", True),
         ("1", True),
         ("False", False),
+        ("false", False),
         ("0", False),
     ],
 )
@@ -109,7 +111,7 @@ def test_str2bool_valid(value: str, expected: bool) -> None:
     assert str2bool(value) is expected
 
 
-@pytest.mark.parametrize("value", ["yes", "no", "true ", ""])
+@pytest.mark.parametrize("value", ["yes", "no", ""])
 def test_str2bool_invalid_raises(value: str) -> None:
     with pytest.raises(argparse.ArgumentTypeError):
         _ = str2bool(value)
@@ -121,6 +123,9 @@ def test_str2bool_invalid_raises(value: str) -> None:
         ("None", int, None),
         ("None", float, None),
         ("None", str, None),
+        ("null", int, None),
+        ("null", float, None),
+        ("null", str, None),
         ("1", int, 1),
         ("1.5", float, 1.5),
         ("foo", str, "foo"),
@@ -203,7 +208,7 @@ def test_build_arg_parser_smoke() -> None:
     assert args.model_type == "protein_mpnn"
     assert args.checkpoint_path == "/tmp/ckpt.pt"
     assert args.is_legacy_weights is True
-    assert args.structure_path == "/tmp/structure.cif"
+    assert args.structure_path == ["/tmp/structure.cif",]
 
 
 @pytest.mark.parametrize(
@@ -338,7 +343,7 @@ def test_cli_to_json_builds_single_input_and_parses_fields() -> None:
     assert inp["features_to_return"] == {"input_features": ["mask_for_loss"]}
     assert inp["atomize_side_chains"] is False
     assert inp["fixed_residues"] == ["A1", "A2"]
-    assert inp["designed_residues"] is None
+    #assert inp["designed_residues"] is None # Not set
     assert inp["bias"] == {"ALA": -1.0}
     assert inp["omit"] == ["UNK"]
     assert inp["temperature"] == 0.1
@@ -349,7 +354,7 @@ def test_cli_to_json_builds_single_input_and_parses_fields() -> None:
     # Structure path and name are carried through; full absolutization happens
     # later in MPNNInferenceInput.post_process_inputs.
     assert inp["structure_path"] == "/tmp/structure.cif"
-    assert "name" in inp  # may be None; defaults later
+    #assert "name" in inp  # may be None; defaults later
 
 
 def test_cli_to_json_with_config_json_file(tmp_path: Path) -> None:
@@ -371,6 +376,46 @@ def test_cli_to_json_with_config_json_file(tmp_path: Path) -> None:
     config = cli_to_json(args)
     assert config == original
 
+def test_cli_to_json_with_both_config_json_and_cmdline(tmp_path: Path) -> None:
+    config_path = tmp_path / "config.json"
+    original = {
+        "model_type": "ligand_mpnn",
+        "checkpoint_path": "/ckpt.pt",
+        "is_legacy_weights": False,
+        "write_fasta": False,
+        "inputs": [{"temperature":0.4, "seed":43}],
+    }
+    config_path.write_text(json.dumps(original))
+
+    parser = build_arg_parser()
+    args = parser.parse_args([
+        "--config_json", str(config_path),
+        "--temperature", "0.3",
+        "--structure_path", "input.cif", "ligand.pdb",
+        "--is_legacy_weights", "true",
+        "--out_directory", "out/",
+        "--structure_noise", "0.5",
+        "--structure_path", "one_more.cif",
+    ])
+
+    config = cli_to_json(args)
+
+    assert config["model_type"] == "ligand_mpnn"
+    assert config["write_fasta"] == False
+    assert config["is_legacy_weights"] == False # From config.json, not overridden
+    assert config["out_directory"] == "out/" # Picked up from options
+
+    inputs = config["inputs"]
+    assert len(inputs) == 3
+    structures = []
+    for inp in inputs:
+        structures.append( inp["structure_path"] )
+        assert inp["seed"] == 43
+        assert inp["temperature"] == 0.4 # From template
+        assert inp["structure_noise"] == 0.5 # From options
+
+    assert sorted(structures) == sorted( ["input.cif", "ligand.pdb", "one_more.cif"] )
+
 
 ###############################################################################
 # MPNNInferenceInput: ID parsing and masking

From c74d682b480c6d4d223973e37c1c6a5bb66358be Mon Sep 17 00:00:00 2001
From: Rocco Moretti <rmorettiase@gmail.com>
Date: Wed, 29 Jul 2026 18:33:49 -0500
Subject: [PATCH 2/4] Format

---
 models/mpnn/src/mpnn/utils/inference.py     | 40 ++++++++++------
 models/mpnn/tests/test_inference_helpers.py | 10 +++-
 models/mpnn/tests/test_inference_utils.py   | 51 +++++++++++++--------
 3 files changed, 68 insertions(+), 33 deletions(-)

diff --git a/models/mpnn/src/mpnn/utils/inference.py b/models/mpnn/src/mpnn/utils/inference.py
index 225f0bc7..56348aba 100644
--- a/models/mpnn/src/mpnn/utils/inference.py
+++ b/models/mpnn/src/mpnn/utils/inference.py
@@ -117,18 +117,20 @@ def none_or_type(v: Any, specified_type: Callable[[Any], Any]) -> Any | None:
 
     Accept both Python like 'None' and JSON-like 'null'.
     """
-    if v == "None" or v == "null": # Python-like or JSON like None/null value.
+    if v == "None" or v == "null":  # Python-like or JSON like None/null value.
         return None
     return specified_type(v)
 
+
 class TrackUserSetting(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):
         setattr(namespace, self.dest, values)
 
-        if not hasattr(namespace, 'user_set'):
+        if not hasattr(namespace, "user_set"):
             setattr(namespace, "user_set", {})
         namespace.user_set[self.dest] = True
 
+
 def build_arg_parser() -> argparse.ArgumentParser:
     """Build the MPNN inference arg parser."""
     parser = argparse.ArgumentParser(
@@ -205,7 +207,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
         help="Path to structure file (CIF or PDB).",
         nargs="+",
         default=MPNN_PER_INPUT_INFERENCE_DEFAULTS["structure_path"],
-        action="extend", # Don't need to track usage
+        action="extend",  # Don't need to track usage
     )
     parser.add_argument(
         "--name",
@@ -629,6 +631,7 @@ def _absolute_path_or_none(path_str: str | None) -> str | None:
         return None
     return str(Path(path_str).expanduser().resolve())
 
+
 def safe_set_in_config(
     config: dict[str, Any],
     args: argparse.Namespace,
@@ -654,12 +657,13 @@ def safe_set_in_config(
         elif format == "list":
             config[name] = parse_list_like(getattr(args, param))
 
+
 def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
     """Convert CLI args into the top-level JSON config dict."""
 
     # In case we haven't provided any of the user-set options, make sure we have a dictionary
     if not hasattr(args, "user_set"):
-       setattr( args, "user_set", {} )
+        setattr(args, "user_set", {})
 
     # If set, use the config_json as the template
     if args.config_json:
@@ -669,7 +673,9 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
         with open(config_path, "r") as f:
             config = json.load(f)
         if not isinstance(config, dict):
-            raise TypeError("The top level of config_json must be a JSON object (dictionary).")
+            raise TypeError(
+                "The top level of config_json must be a JSON object (dictionary)."
+            )
     else:
         # Bare-bones config -- we'll populate it from options/defaults
         config = {}
@@ -695,27 +701,33 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
 
     for entry in config["inputs"]:
         if not isinstance(entry, dict):
-            raise TypeError("Each entry in the inputs list of config_json must be a JSON object (dictionary)")
+            raise TypeError(
+                "Each entry in the inputs list of config_json must be a JSON object (dictionary)"
+            )
 
     # Find the template input dicts in the config (either from config_json, or the stub config) and populate them with --structure_path values
-    structure_free = [ d for d in config["inputs"] if "structure_path" not in d ]
+    structure_free = [d for d in config["inputs"] if "structure_path" not in d]
     if len(structure_free) != 0:
         if len(args.structure_path) == 0:
-            raise ValueError("structure_path must be specified, either in the config_json or the command line")
+            raise ValueError(
+                "structure_path must be specified, either in the config_json or the command line"
+            )
 
         # combinitorial assortment
         new_inputs = []
         for entry in structure_free:
             for structure in args.structure_path:
-                new_input = copy.deepcopy( entry )
+                new_input = copy.deepcopy(entry)
                 new_input["structure_path"] = structure
-                new_inputs.append( new_input )
+                new_inputs.append(new_input)
 
         # TODO: Add some sort of check/validation/adjustment for output naming in the combinitorial case
-        already_has_structure = [ d for d in config["inputs"] if "structure_path" in d ]
+        already_has_structure = [d for d in config["inputs"] if "structure_path" in d]
         config["inputs"] = already_has_structure + new_inputs
     elif args.structure_path and len(args.structure_path) != 0:
-        raise ValueError("Command line structure specified with --structure_path, but all provided input entries already have structures specified.")
+        raise ValueError(
+            "Command line structure specified with --structure_path, but all provided input entries already have structures specified."
+        )
 
     for entry in config["inputs"]:
         # Name
@@ -735,7 +747,9 @@ def cli_to_json(args: argparse.Namespace) -> dict[str, Any]:
         safe_set_in_config(entry, args, "structure_noise")
         safe_set_in_config(entry, args, "decode_type")
         safe_set_in_config(entry, args, "causality_pattern")
-        safe_set_in_config(entry, args, "initialize_sequence_embedding_with_ground_truth")
+        safe_set_in_config(
+            entry, args, "initialize_sequence_embedding_with_ground_truth"
+        )
         safe_set_in_config(entry, args, "features_to_return", format="json")
         # Only applicable for LigandMPNN
         safe_set_in_config(entry, args, "atomize_side_chains")
diff --git a/models/mpnn/tests/test_inference_helpers.py b/models/mpnn/tests/test_inference_helpers.py
index 7832094d..19769e97 100644
--- a/models/mpnn/tests/test_inference_helpers.py
+++ b/models/mpnn/tests/test_inference_helpers.py
@@ -18,7 +18,15 @@
 
 
 @pytest.mark.parametrize(
-    ("value", "expected"), [("True", True), ("true",True), ("1", True), ("False", False), ("false", False), ("0", False)]
+    ("value", "expected"),
+    [
+        ("True", True),
+        ("true", True),
+        ("1", True),
+        ("False", False),
+        ("false", False),
+        ("0", False),
+    ],
 )
 def test_str2bool_valid(value: str, expected: bool):
     assert str2bool(value) is expected
diff --git a/models/mpnn/tests/test_inference_utils.py b/models/mpnn/tests/test_inference_utils.py
index 61310d3b..4aad8ae7 100644
--- a/models/mpnn/tests/test_inference_utils.py
+++ b/models/mpnn/tests/test_inference_utils.py
@@ -208,7 +208,9 @@ def test_build_arg_parser_smoke() -> None:
     assert args.model_type == "protein_mpnn"
     assert args.checkpoint_path == "/tmp/ckpt.pt"
     assert args.is_legacy_weights is True
-    assert args.structure_path == ["/tmp/structure.cif",]
+    assert args.structure_path == [
+        "/tmp/structure.cif",
+    ]
 
 
 @pytest.mark.parametrize(
@@ -343,7 +345,7 @@ def test_cli_to_json_builds_single_input_and_parses_fields() -> None:
     assert inp["features_to_return"] == {"input_features": ["mask_for_loss"]}
     assert inp["atomize_side_chains"] is False
     assert inp["fixed_residues"] == ["A1", "A2"]
-    #assert inp["designed_residues"] is None # Not set
+    # assert inp["designed_residues"] is None # Not set
     assert inp["bias"] == {"ALA": -1.0}
     assert inp["omit"] == ["UNK"]
     assert inp["temperature"] == 0.1
@@ -354,7 +356,7 @@ def test_cli_to_json_builds_single_input_and_parses_fields() -> None:
     # Structure path and name are carried through; full absolutization happens
     # later in MPNNInferenceInput.post_process_inputs.
     assert inp["structure_path"] == "/tmp/structure.cif"
-    #assert "name" in inp  # may be None; defaults later
+    # assert "name" in inp  # may be None; defaults later
 
 
 def test_cli_to_json_with_config_json_file(tmp_path: Path) -> None:
@@ -376,6 +378,7 @@ def test_cli_to_json_with_config_json_file(tmp_path: Path) -> None:
     config = cli_to_json(args)
     assert config == original
 
+
 def test_cli_to_json_with_both_config_json_and_cmdline(tmp_path: Path) -> None:
     config_path = tmp_path / "config.json"
     original = {
@@ -383,38 +386,48 @@ def test_cli_to_json_with_both_config_json_and_cmdline(tmp_path: Path) -> None:
         "checkpoint_path": "/ckpt.pt",
         "is_legacy_weights": False,
         "write_fasta": False,
-        "inputs": [{"temperature":0.4, "seed":43}],
+        "inputs": [{"temperature": 0.4, "seed": 43}],
     }
     config_path.write_text(json.dumps(original))
 
     parser = build_arg_parser()
-    args = parser.parse_args([
-        "--config_json", str(config_path),
-        "--temperature", "0.3",
-        "--structure_path", "input.cif", "ligand.pdb",
-        "--is_legacy_weights", "true",
-        "--out_directory", "out/",
-        "--structure_noise", "0.5",
-        "--structure_path", "one_more.cif",
-    ])
+    args = parser.parse_args(
+        [
+            "--config_json",
+            str(config_path),
+            "--temperature",
+            "0.3",
+            "--structure_path",
+            "input.cif",
+            "ligand.pdb",
+            "--is_legacy_weights",
+            "true",
+            "--out_directory",
+            "out/",
+            "--structure_noise",
+            "0.5",
+            "--structure_path",
+            "one_more.cif",
+        ]
+    )
 
     config = cli_to_json(args)
 
     assert config["model_type"] == "ligand_mpnn"
     assert config["write_fasta"] == False
-    assert config["is_legacy_weights"] == False # From config.json, not overridden
-    assert config["out_directory"] == "out/" # Picked up from options
+    assert config["is_legacy_weights"] == False  # From config.json, not overridden
+    assert config["out_directory"] == "out/"  # Picked up from options
 
     inputs = config["inputs"]
     assert len(inputs) == 3
     structures = []
     for inp in inputs:
-        structures.append( inp["structure_path"] )
+        structures.append(inp["structure_path"])
         assert inp["seed"] == 43
-        assert inp["temperature"] == 0.4 # From template
-        assert inp["structure_noise"] == 0.5 # From options
+        assert inp["temperature"] == 0.4  # From template
+        assert inp["structure_noise"] == 0.5  # From options
 
-    assert sorted(structures) == sorted( ["input.cif", "ligand.pdb", "one_more.cif"] )
+    assert sorted(structures) == sorted(["input.cif", "ligand.pdb", "one_more.cif"])
 
 
 ###############################################################################

From d0a373382a0fa44ec2e8c43c1572d8b5a2b0e436 Mon Sep 17 00:00:00 2001
From: Rocco Moretti <rmorettiase@gmail.com>
Date: Wed, 29 Jul 2026 18:38:38 -0500
Subject: [PATCH 3/4] Ruff fixes

---
 models/mpnn/tests/test_inference_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/models/mpnn/tests/test_inference_utils.py b/models/mpnn/tests/test_inference_utils.py
index 4aad8ae7..217f1e8e 100644
--- a/models/mpnn/tests/test_inference_utils.py
+++ b/models/mpnn/tests/test_inference_utils.py
@@ -414,8 +414,8 @@ def test_cli_to_json_with_both_config_json_and_cmdline(tmp_path: Path) -> None:
     config = cli_to_json(args)
 
     assert config["model_type"] == "ligand_mpnn"
-    assert config["write_fasta"] == False
-    assert config["is_legacy_weights"] == False  # From config.json, not overridden
+    assert not config["write_fasta"]
+    assert not config["is_legacy_weights"]  # From config.json, not overridden
     assert config["out_directory"] == "out/"  # Picked up from options
 
     inputs = config["inputs"]

From 57e0458bb245f3b6b4c07192d2947233a7880729 Mon Sep 17 00:00:00 2001
From: Rocco Moretti <rmorettiase@gmail.com>
Date: Wed, 29 Jul 2026 18:44:01 -0500
Subject: [PATCH 4/4] Fix mypy test.

---
 models/mpnn/src/mpnn/utils/inference.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/models/mpnn/src/mpnn/utils/inference.py b/models/mpnn/src/mpnn/utils/inference.py
index 56348aba..5bf2db24 100644
--- a/models/mpnn/src/mpnn/utils/inference.py
+++ b/models/mpnn/src/mpnn/utils/inference.py
@@ -123,7 +123,13 @@ def none_or_type(v: Any, specified_type: Callable[[Any], Any]) -> Any | None:
 
 
 class TrackUserSetting(argparse.Action):
-    def __call__(self, parser, namespace, values, option_string=None):
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: Any,
+        option_string: str | None = None,
+    ) -> None:
         setattr(namespace, self.dest, values)
 
         if not hasattr(namespace, "user_set"):