cleanup; remove eval from train.py

config/config.json

@@ -5,7 +5,7 @@
     "model_path": "checkpoints/rf_alltasks.joblib",
     "data_folder": "data_tox21/",
     "log_folder": "logs/",
-    "debug": 1, 
+    "debug": "false", 
     "task_configs": {
         "NR-AR": {
             "max_depth": "none",

predict.py

@@ -10,8 +10,8 @@ from collections import defaultdict
 
 import json
 import numpy as np
-
 from tqdm import tqdm
+
 from src.preprocess import create_descriptors
 from src.utils import TASKS, normalize_config
 from src.model import Tox21RFClassifier
@@ -73,7 +73,7 @@ def predict(
 from testing import test_eval
 
 with open(CONFIG_FILE, "r") as f:
-    cfg = json.load(f)
-cfg = normalize_config(cfg)
+    config = json.load(f)
+config = normalize_config(config)
 
-test_eval(predict, debug=cfg["debug"], use_only_clean=False, use_only_first=False)
+test_eval(predict, debug=config["debug"], use_only_clean=False, use_only_first=False)

preprocess.py

@@ -15,7 +15,6 @@ from src.preprocess import create_descriptors, get_tox21_split
 from src.utils import (
     TASKS,
     HF_TOKEN,
-    write_pickle,
     create_dir,
 )
 

src/model.py

@@ -11,10 +11,10 @@ import joblib
 
 import numpy as np
 
+from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.preprocessing import StandardScaler
 from sklearn.feature_selection import VarianceThreshold
-from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.preprocessing import StandardScaler
 from statsmodels.distributions.empirical_distribution import ECDF
 
 from .utils import TASKS

src/preprocess.py

@@ -18,11 +18,7 @@ from rdkit import Chem, DataStructs
 from rdkit.Chem import Descriptors, rdFingerprintGenerator, MACCSkeys
 from rdkit.Chem.rdchem import Mol
 
-from .utils import (
-    USED_200_DESCR,
-    TOX_SMARTS_PATH,
-    Standardizer,
-)
+from .utils import USED_200_DESCR, TOX_SMARTS_PATH, Standardizer
 
 
 def create_cleaned_mol_objects(smiles: list[str]) -> tuple[list[Mol], np.ndarray]:

train.py

@@ -9,8 +9,6 @@ import logging
 import argparse
 
 import numpy as np
-
-from sklearn.metrics import roc_auc_score
 from datetime import datetime
 
 from src.model import Tox21RFClassifier
@@ -77,18 +75,6 @@ def main(cfg):
     labels = np.concatenate([train_y, val_y], axis=0)
     logger.info(f"Train data shape: {data.shape}")
 
-    full_data = np.load(
-        "data/tox21_descriptors.npz",
-        allow_pickle=True,
-    )
-
-    test_mask = full_data["sets"] == "test"
-    test_data = full_data["features"][test_mask]
-    test_labels = full_data["labels"][test_mask]
-    assert (
-        data.shape[1] == test_data.shape[1]
-    ), "different number of features found in train and test set!"
-
     if cfg["model_path"]:
         logger.info(
             f"Fitted RandomForestClassifier will be saved as: {cfg['model_path']}"
@@ -113,6 +99,8 @@ def main(cfg):
         task_labels = task_labels[label_mask].astype(int)
 
         model.fit(task, task_data, task_labels)
+        if cfg["debug"]:
+            break
 
     log_text = f"Finished training."
     logger.info(log_text)
@@ -121,29 +109,6 @@ def main(cfg):
         model.save_model(cfg["model_path"])
         logger.info(f"Save model as: {cfg['model_path']}")
 
-    del model
-    model = Tox21RFClassifier()
-    model.load_model(cfg["model_path"])
-
-    logger.info("Evaluate model")
-    results = {}
-    preds = np.empty_like(test_labels, dtype=np.float32)
-    for i, task in enumerate(model.tasks):
-        task_labels = test_labels[:, i]
-        label_mask = ~np.isnan(task_labels)
-
-        task_labels = task_labels[label_mask].astype(int)
-
-        pred = model.predict(task, test_data)
-        results[task] = [roc_auc_score(y_true=task_labels, y_score=pred[label_mask])]
-        preds[:, i] = pred.copy()
-
-    logger.info("Results:")
-    logger.info(results)
-    logger.info(
-        f"Average: {sum([score[0] for score in results.values()]) / len(results)}"
-    )
-
 
 if __name__ == "__main__":
     args = parser.parse_args()