V2: improve CLI for multitask

chemprop/cli/predict.py

@@ -61,6 +61,11 @@ def add_predict_args(parser: ArgumentParser) -> ArgumentParser:
         required=True,
         help="Path to a pretrained model checkpoint (.ckpt) or a pretrained model file (.pt).",
     )
+    parser.add_argument(
+        "--target-columns",
+        nargs="+",
+        help="Column names to save the predictions to. If not provided, the predictions will be saved to columns named 'pred_0', 'pred_1', etc.",
+    )
 
     # TODO: add uncertainty and calibration
     # unc_args = parser.add_argument_group("uncertainty and calibration args")
@@ -264,7 +269,17 @@ def main(args):
     preds = torch.concat(predss, 0)
     if isinstance(model.predictor, MulticlassClassificationFFN):
         preds = torch.argmax(preds, dim=-1)
-    target_columns = [f"pred_{i}" for i in range(preds.shape[1])] # TODO: need to improve this
+
+    if args.target_columns is not None:
+        assert (
+            len(args.target_columns) == model.n_tasks
+        ), "Number of target columns must match the number of tasks."
+        target_columns = args.target_columns
+    else:
+        target_columns = [
+            f"pred_{i}" for i in range(preds.shape[1])
+        ]  # TODO: need to improve this for cases like multi-task MVE and multi-task multiclass
+
     df_test[target_columns] = preds
     if args.output.suffix == ".pkl":
         df_test = df_test.reset_index(drop=True)

chemprop/cli/train.py

@@ -788,7 +788,7 @@ def main(args):
         if args.save_smiles_splits:
             save_smiles_splits(args, output_dir, train_dset, val_dset, test_dset)
 
-        if args.task_type == "regression":
+        if "regression" in args.task_type:
             scaler = train_dset.normalize_targets()
             val_dset.normalize_targets(scaler)
             logger.info(f"Train data: mean = {scaler.mean_} | std = {scaler.scale_}")

chemprop/models/model.py

@@ -10,7 +10,7 @@
 
 from chemprop.data import TrainingBatch, BatchMolGraph
 from chemprop.nn.metrics import Metric
-from chemprop.nn import MessagePassing, Aggregation, Predictor, LossFunction
+from chemprop.nn import MessagePassing, Aggregation, Predictor, LossFunction, MveFFN
 from chemprop.schedulers import NoamLR
 
 

chemprop/nn/predictors.py

@@ -115,13 +115,22 @@ def __init__(
         dropout: float = 0,
         activation: str = "relu",
         criterion: LossFunction | None = None,
-        loc: float | Tensor = 0,
-        scale: float | Tensor = 1,
+        loc: float | Tensor = 0.,
+        scale: float | Tensor = 1.,
     ):
         super().__init__(n_tasks, input_dim, hidden_dim, n_layers, dropout, activation, criterion)
 
-        self.register_buffer("loc", torch.tensor(loc).view(1, -1))
-        self.register_buffer("scale", torch.tensor(scale).view(1, -1))
+        if isinstance(loc, float):
+            loc = torch.ones(1, self.n_tasks) * loc
+        else:
+            loc = torch.tensor(loc).view(1, -1)
+        self.register_buffer("loc", loc)
+
+        if isinstance(scale, float):
+            scale = torch.ones(1, self.n_tasks) * scale
+        else:
+            scale = torch.tensor(scale).view(1, -1)
+        self.register_buffer("scale", scale)
 
     def forward(self, Z: Tensor) -> Tensor:
         Y = super().forward(Z)

tests/cli/test_cli_regression_mol_multitask.py

@@ -0,0 +1,60 @@
+"""This tests the CLI functionality of training and predicting a regression model on a single molecule.
+"""
+
+import pytest
+
+from chemprop.cli.main import main
+
+pytestmark = pytest.mark.CLI
+
+
+@pytest.fixture
+def data_path(data_dir):
+    return str(data_dir / "regression" / "mol_multitask.csv")
+
+
+@pytest.fixture
+def model_path(data_dir):
+    return str(data_dir / "example_model_v2_regression_mol_multitask.pt")
+
+
+def test_train_quick(monkeypatch, data_path):
+    args = ["chemprop", "train", "-i", data_path, "--epochs", "1", "--num-workers", "0"]
+
+    with monkeypatch.context() as m:
+        m.setattr("sys.argv", args)
+        main()
+
+
+def test_predict_quick(monkeypatch, data_path, model_path):
+    args = [
+        "chemprop",
+        "predict",
+        "-i",
+        data_path,
+        "--model-path",
+        model_path,
+        "--target-columns",
+        "mu",
+        "alpha",
+        "homo",
+        "lumo",
+        "gap",
+        "r2",
+        "zpve",
+        "cv",
+        "u0",
+        "u298",
+        "h298",
+        "g298",
+    ]
+
+    with monkeypatch.context() as m:
+        m.setattr("sys.argv", args)
+        main()
+
+    args = ["chemprop", "predict", "-i", data_path, "--model-path", model_path]
+
+    with monkeypatch.context() as m:
+        m.setattr("sys.argv", args)
+        main()