adapt load/saving, preprocessing, app, readme, modelcard

MODELCARD.md → MODEL_CARD.md

@@ -1,7 +1,7 @@
 # Model card - tox21_rf_classifier
 ### Model details
 - Model name: Random Forest Tox21 Baseline
-- Developer: ML-JKU
+- Developer: JKU (Linz)
 - Paper URL: https://link.springer.com/article/10.1023/A:1010933404324
 - Model type / architecture:
     - Random Forest implemented using sklearn.RandomForestClassifier.
@@ -29,7 +29,3 @@ Tox21 training and validation sets.
 
 ### Evaluation data
 Tox21 test set.
-
-### Hugging Face Space details
-- Space: MASKED-FOR-REVIEW
-- Git commit hash: MASKED-FOR-REVIEW

README.md

@@ -6,31 +6,39 @@ colorTo: purple
 sdk: docker
 pinned: false
 license: cc-by-nc-4.0
-short_description: This is a RF classifier for the Tox21 test dataset
+short_description: Random Forest Baseline for Tox21
 ---
 
 # Tox21 Random Forest Classifier
 
-This repository hosts a Hugging Face Space that provides an examplary API for submitting models to the [Tox21 Leaderboard](https://huggingface.co/spaces/ml-jku/tox21_leaderboard).
+This repository hosts a Hugging Face Space that provides an API for submitting models to the [Tox21 Leaderboard](https://huggingface.co/spaces/tschouis/tox21_leaderboard).
 
+Here **Random Forest (RF)** models are trained on the Tox21 dataset, and the trained models are provided for inference. For each of the twelve toxic effects, a separate RF model is trained. The input to the model is a **SMILES** string of the small molecule, and the output are 12 numeric values for each of the toxic effects of the Tox21 dataset. 
 
-Here **Random Forest (RF)** models are trained on the Tox21 dataset, and the trained models are provided for 
-inference. For each of the twelve toxic effects, a separate RF model is trained. The input to the model
-is a **SMILES** string of the small molecule, and the output are 12 numeric values for 
-each of the toxic effects of the Tox21 dataset. 
-
+**Important:** For leaderboard submission, your Space needs to include training code. The file `train.py` should train the model using the config specified inside the `config/` folder and save the final model parameters into a file inside the `checkpoints/` folder. The model should be trained using the [Tox21_dataset](https://huggingface.co/datasets/tschouis/tox21) provided on Hugging Face. The datasets can be loaded like this:
+```python
+from datasets import load_dataset
+ds = load_dataset("ml-jku/tox21", token=token)
+train_df = ds["train"].to_pandas()
+val_df = ds["validation"].to_pandas()
+```
 
-**Important:** For leaderboard submission, your Space does not need to include training code. It only needs to implement inference in the `predict()` function inside `predict.py`. The `predict()` function must keep the provided skeleton: it should take a list of SMILES strings as input and return a prediction dictionary as output, with SMILES and targets as keys. Therefore, any preprocessing of SMILES strings must be executed on-the-fly during inference.
+Additionally, the Space needs to implement inference in the `predict()` function inside `predict.py`. The `predict()` function must keep the provided skeleton: it should take a list of SMILES strings as input and return a nested prediction dictionary as output, with SMILES as keys and dictionaries containing targetname-prediction pairs as values. Therefore, any preprocessing of SMILES strings must be executed on-the-fly during inference.
 
 # Repository Structure
 - `predict.py` - Defines the `predict()` function required by the leaderboard (entry point for inference).
 - `app.py` - FastAPI application wrapper (can be used as-is).
+- `preprocess.py` - preprocesses SMILES strings to generate feature descriptors and saves results as NPZ files in `data/`.
+- `train.py` - trains and saves a model using the config in the `config/` folder.
+- `config/` - the config file used by `train.py`. 
+- `logs/` - all the logs of `train.py`, the saved model, and predictions on the validation set.
+- `data/` - RF uses numerical data. During preprocessing in `preprocess.py` two NPZ files containing molecule features are created and saved here.
+- `checkpoints/` - the saved model that is used in `predict.py` is here.
 
 - `src/` - Core model & preprocessing logic:
-    - `data.py` - SMILES preprocessing pipeline
-    - `model.py` - Random Forest classifier wrapper
-    - `train.py` - Script to train the classifier
-    - `utils.py` – Constants and Helper functions
+    - `preprocess.py` - SMILES preprocessing logic
+    - `model.py` - RF model class with processing, saving and loading logic
+    - `utils.py` - utility functions
 
 # Quickstart with Spaces
 
@@ -44,13 +52,17 @@ You can easily adapt this project in your own Hugging Face account:
 
 - Modify `predict()` inside `predict.py` to perform model inference while keeping the function skeleton unchanged to remain compatible with the leaderboard.
 
+- Modify `train.py` and/or `preprocess.py` according to your model and preprocessing pipeline.
+
+- Modify the file inside `config/` to contain all hyperparameters that are set in `train.py`.
+
 That’s it, your model will be available as an API endpoint for the Tox21 Leaderboard.
 
 # Installation
 To run (and train) the random forest, clone the repository and install dependencies:
 
 ```bash
-git clone https://huggingface.co/spaces/ml-jku/tox21_rf_classifier
+git clone https://huggingface.co/spaces/tschouis/tox21_rf_classifier
 cd tox_21_rf_classifier
 
 conda create -n tox21_rf_cls python=3.11
@@ -58,22 +70,6 @@ conda activate tox21_rf_cls
 pip install -r requirements.txt
 ```
 
-# Training
-
-To train the Random Forest model from scratch:
-
-```bash
-python -m src/train.py
-```
-
-This will:
-
-1. Load and preprocess the Tox21 training dataset.
-2. Train a Random Forest classifier.
-3. Save the trained model to the assets/ folder.
-4. Evaluate the trained Random Forest classifier on the validation split. 
-
-
 # Inference
 
 For inference, you only need `predict.py`.
@@ -101,8 +97,6 @@ The output will be a nested dictionary in the format:
 
 # Notes
 
-- Only adapting `predict.py` for your model inference is required for leaderboard submission.
-
-- Training (`src/train.py`) is provided for reproducibility.
+- Adapting `predict.py`, `train.py`, `config/`, and `checkpoints/` is required for leaderboard submission.
 
-- Preprocessing (here inside `src/data.py`) must be applied at inference time, not just training.
+- Preprocessing must be done inside `predict.py` not just `train.py`.

app.py

@@ -44,7 +44,7 @@ def root():
 @app.get("/metadata")
 def metadata():
     return {
-        "name": "Tox21RandomForest",
+        "name": "RF",
         "version": "1.0.0",
         "max_batch_size": 256,
         "tox_endpoints": [
@@ -74,5 +74,5 @@ def predict(request: Request):
     predictions = predict_func(request.smiles)
     return {
         "predictions": predictions,
-        "model_info": {"name": "Tox21RandomForest", "version": "1.0.0"},
+        "model_info": {"name": "RF", "version": "1.0.0"},
     }

config/config.json

@@ -17,9 +17,11 @@
         "use": "true",
         "min_var": 0.01,
         "max_corr": 0.95,
-        "feature_keys": ["ecfps", "tox", "maccs", "rdkit_descrs"],
-        "independent_keys": "false",
-        "max_features": -1
+        "max_features": -1,
+        "min_var__feature_keys": ["ecfps", "tox", "maccs", "rdkit_descrs"],
+        "max_corr__feature_keys": ["ecfps", "tox", "maccs", "rdkit_descrs"],
+        "min_var__independent_keys": "false",
+        "max_corr__independent_keys": "false"
     },
     "feature_quantilization": {
         "use": "true",

predict.py

@@ -57,12 +57,11 @@ def predict(
         scaler=config["scaler"],
     )
 
-    state = joblib.load(config["ckpt_path"])
-    model.set_state(state)
+    model.load(config["ckpt_path"])
     print(f"Loaded model from {config['ckpt_path']}")
 
     state = joblib.load(config["preprocessor_path"])
-    preprocessor.__setstate__(state)
+    preprocessor.set_state(state)
     print(f"Loaded preprocessor from {config['preprocessor_path']}")
 
     # make predicitons

src/model.py

@@ -6,6 +6,7 @@ SMILES and target names as keys.
 
 # ---------------------------------------------------------------------------------------
 # Dependencies
+import joblib
 import numpy as np
 from sklearn.ensemble import RandomForestClassifier
 
@@ -33,17 +34,21 @@ class Tox21RFClassifier:
             for task in self.tasks
         }
 
-    def set_state(self, state: dict) -> None:
-        """Sets the state of the model
+    def load(self, path: str) -> None:
+        """Load model from filepath
 
         Args:
-            state (dict): models state dict
+            path (str): filepath to model checkpoint
         """
-        self.models = state
+        self.models = joblib.load(path)
 
-    def get_state(self) -> None:
-        """Return model state dict"""
-        return self.models
+    def save(self, path: str) -> None:
+        """Save model to filepath
+
+        Args:
+            path (str): filepath to model checkpoint
+        """
+        joblib.dump(self.models, path)
 
     def fit(self, task: str, X: np.ndarray, y: np.ndarray) -> None:
         """Train the random forest for a given task

src/preprocess.py

@@ -78,7 +78,7 @@ class SubSampler(TransformerMixin, BaseEstimator):
         _X = X.copy()
         _y = y.copy() if y is not None else None
 
-        if self.max_samples > 0:
+        if self.max_samples > 0 and _X.shape[0] > self.max_samples:
             resample_idxs = np.random.choice(
                 np.arange(_X.shape[0]), size=(self.max_samples,), replace=True
             )
@@ -127,21 +127,32 @@ class FeatureSelector(FeatureDictMixin, TransformerMixin, BaseEstimator):
         max_corr=1.0,
         max_features=-1,
         feature_keys=None,
-        independent_keys=False,
+        min_var__feature_keys=None,
+        max_corr__feature_keys=None,
+        max_features__feature_keys=None,
+        min_var__independent_keys=False,
+        max_corr__independent_keys=False,
+        max_features__independent_keys=False,
     ):
         self.min_var = min_var
         self.max_corr = max_corr
         self.max_features = max_features
-        self.independent_keys = independent_keys
-        self._feature_mask = None
+
+        self.min_var__feature_keys = min_var__feature_keys
+        self.max_corr__feature_keys = max_corr__feature_keys
+        self.max_features__feature_keys = max_features__feature_keys
+
+        self.min_var__independent_keys = min_var__independent_keys
+        self.max_corr__independent_keys = max_corr__independent_keys
+        self.max_features__independent_keys = max_features__independent_keys
 
         super().__init__(feature_keys=feature_keys)
 
-    def _get_min_var_feature_mask(self, X: np.ndarray) -> np.ndarray:
+    def _get_min_var_mask(self, X: np.ndarray, *args) -> np.ndarray:
         var_thresh = VarianceThreshold(threshold=self.min_var)
         return var_thresh.fit(X).get_support()  # mask
 
-    def _get_max_corr_feature_mask(
+    def _get_max_corr_mask(
         self, X: np.ndarray, prev_feature_mask: np.ndarray
     ) -> np.ndarray:
         _prev_feature_mask = prev_feature_mask.copy()
@@ -156,45 +167,86 @@ class FeatureSelector(FeatureDictMixin, TransformerMixin, BaseEstimator):
         _prev_feature_mask[_prev_feature_mask] = to_keep
         return _prev_feature_mask
 
+    def _get_max_features_mask(
+        self, X: np.ndarray, prev_feature_mask: np.ndarray
+    ) -> np.ndarray:
+        _prev_feature_mask = prev_feature_mask.copy()
+        # select features with at least max_var variation
+        feature_vars = np.nanvar(X[:, _prev_feature_mask], axis=0)
+        order = np.argsort(feature_vars)[: -(self.max_features + 1) : -1]
+        keep_feat_idx = np.arange(len(_prev_feature_mask))[order]
+        _prev_feature_mask = np.isin(
+            np.arange(len(_prev_feature_mask)), keep_feat_idx, assume_unique=True
+        )
+        return _prev_feature_mask
+
+    def apply_filter(self, filter, X, prev_feature_mask):
+        mask = prev_feature_mask.copy()
+        func = self.__getattribute__(f"_get_{filter}_mask")
+        feature_keys = self.__getattribute__(f"{filter}__feature_keys")
+
+        if self.__getattribute__(f"{filter}__independent_keys"):
+            for key in feature_keys:
+                key_mask = self._curr_keys == key
+                mask[key_mask] = func(X[:, key_mask], mask[key_mask])
+
+        else:
+            feature_key_mask = np.isin(self._curr_keys, feature_keys)
+            mask[feature_key_mask] = func(
+                X[:, feature_key_mask], mask[feature_key_mask]
+            )
+        return mask
+
     def fit(self, X: dict[str, np.ndarray]):
         _X = self.dict_to_array(X)
         feature_mask = np.ones((_X.shape[1]), dtype=bool)
 
         # select features with at least min_var variation
         if self.min_var > 0.0:
-            if self.independent_keys:
-                for key in self.feature_keys:
+            if self.min_var__independent_keys:
+                for key in self.min_var__feature_keys:
                     key_mask = self._curr_keys == key
-                    subset = _X[:, key_mask]
-                    feature_mask[key_mask] = self._get_min_var_feature_mask(subset)
+                    feature_mask[key_mask] = self._get_min_var_mask(_X[:, key_mask])
 
             else:
-                feature_mask = self._get_min_var_feature_mask(_X)
+                feature_key_mask = np.isin(self._curr_keys, self.min_var__feature_keys)
+                feature_mask[feature_key_mask] = self._get_min_var_mask(
+                    _X[:, feature_key_mask]
+                )
 
         # select features with at least max_var variation
         if self.max_corr < 1.0:
-            if self.independent_keys:
-                for key in self.feature_keys:
+            if self.max_corr__independent_keys:
+                for key in self.max_corr__feature_keys:
                     key_mask = self._curr_keys == key
                     subset = _X[:, key_mask]
-                    feature_mask[key_mask] = self._get_max_corr_feature_mask(
+                    feature_mask[key_mask] = self._get_max_corr_mask(
                         subset, feature_mask[key_mask]
                     )
             else:
-                feature_mask = self._get_max_corr_feature_mask(_X, feature_mask)
+                feature_key_mask = np.isin(self._curr_keys, self.max_corr__feature_keys)
+                feature_mask[feature_key_mask] = self._get_max_corr_mask(
+                    _X[:, feature_key_mask], feature_mask[feature_key_mask]
+                )
 
         if self.max_features == 0:
             raise ValueError(
                 f"max_features (={self.max_features}) must be -1 or larger 0."
             )
         elif self.max_features > 0:
-            # select features with at least max_var variation
-            feature_vars = np.nanvar(_X[:, feature_mask], axis=0)
-            order = np.argsort(feature_vars)[: -(self.max_features + 1) : -1]
-            keep_feat_idx = np.arange(feature_mask)[order]
-            feature_mask = np.isin(
-                np.arange(feature_mask), keep_feat_idx, assume_unique=True
-            )
+            if self.max_features__independent_keys:
+                for key in self.max_features__feature_keys:
+                    key_mask = self._curr_keys == key
+                    feature_mask[key_mask] = self._get_max_features_mask(
+                        _X[:, key_mask], feature_mask[key_mask]
+                    )
+            else:
+                feature_key_mask = np.isin(
+                    self._curr_keys, self.max_features__feature_keys
+                )
+                feature_mask[feature_key_mask] = self._get_max_features_mask(
+                    _X[:, feature_key_mask], feature_mask[feature_key_mask]
+                )
 
         self._feature_mask = feature_mask
         self.is_fitted_ = True
@@ -278,6 +330,7 @@ class FeaturePreprocessor(TransformerMixin, BaseEstimator):
 
         self.feature_selection_config = copy.deepcopy(feature_selection_config)
         self.use_feat_selec = self.feature_selection_config.pop("use")
+        self.feature_selection_config["feature_keys"] = descriptors
         self.feature_selector = FeatureSelector(**self.feature_selection_config)
 
         self.max_samples = max_samples
@@ -330,10 +383,8 @@ class FeaturePreprocessor(TransformerMixin, BaseEstimator):
 
         if self.use_feat_quant:
             _X = self.quantile_creator.transform(_X)
-
         if self.use_feat_selec:
             _X = self.feature_selector.transform(_X)
-
         _X = np.concatenate([_X[descr] for descr in self.descriptors], axis=1)
         _X = self.scaler.transform(_X)
 

src/utils.py

@@ -32,8 +32,6 @@ TASKS = [
     "SR-p53",
 ]
 
-KNOWN_DESCR = ["ecfps", "tox", "maccs", "rdkit_descrs"]
-
 USED_200_DESCR = [
     0,
     1,

train.py

@@ -46,10 +46,10 @@ def main(config):
     )
 
     logger.info(f"Config: {config}")
-    model_config_repr = "Model configs: \n" + "\n".join(
+    model_config_repr = "Model config: \n" + "\n".join(
         [str(val) for val in config["model_config"].values()]
     )
-    logger.info(f"Model configs: \n{model_config_repr}")
+    logger.info(f"Model config: \n{model_config_repr}")
 
     # seeding
     random.seed(config["seed"])
@@ -111,8 +111,7 @@ def main(config):
     logger.info(log_text)
 
     if config["ckpt_path"]:
-        state = model.get_state()
-        joblib.dump(state, config["ckpt_path"])
+        model.save(config["ckpt_path"])
         logger.info(f"Save model as: {config['ckpt_path']}")
 
     if config["preprocessor_path"]: