documentation updates

README.md

@@ -34,14 +34,14 @@ pip install torch==2.12.0 torch_scatter torch_geometric -f https://data.pyg.org/
 ## Usage
 
 ```bash
-# Predict for one or more SMILES / InChI strings (default config: eval)
+# Predict for one or more SMILES / InChI strings (default config: web)
 python -m chebifier predict -m "CC(=O)OC1=CC=CC=C1C(=O)O" -m "C1=CC=C(C=C1)C(=O)O"
 
 # Predict for molecules listed in a file (one SMILES / InChI per line)
 python -m chebifier predict -f smiles.txt
 
-# Use the web ensemble, or your own configuration file
-python -m chebifier predict -e web -m "CC(=O)O"
+# Use the eval ensemble, or your own configuration file
+python -m chebifier predict -e eval -m "CC(=O)O"
 python -m chebifier predict -e configs/my_config.yml -f smiles.txt
 
 # Get all available options
@@ -50,8 +50,8 @@ python -m chebifier predict --help
 
 ### Advanced CLI
 
-The ensemble configuration is selected with `--ensemble-config`: `eval` or `web` (both downloaded from
-[Hugging Face](https://huggingface.co/datasets/chebai/chebifier), `eval` is the default) or a path to your own
+The ensemble configuration is selected with `--ensemble-config`: `web` or `eval` (both downloaded from
+[Hugging Face](https://huggingface.co/datasets/chebai/chebifier), `web` is the default) or a path to your own
 configuration file. Create your own file to change which models are included in the ensemble or how they are weighted.
 
 Trained deep learning models are automatically downloaded from [Hugging Face](https://huggingface.co/chebai).
@@ -78,7 +78,7 @@ my_gat:
 
 You can also supply your own model checkpoints (see `configs/example_config.yml` for an example).
 
-The base learners are selected with `-e`/`--ensemble-config` (default `eval`). The deep learning
+The base learners are selected with `-e`/`--ensemble-config` (default `web`). The deep learning
 base learners and the ensemble's calibration for the standard `eval`/`web` configs are downloaded
 from Hugging Face automatically on first use. To use a calibration of your own (e.g. one you built
 yourself, see below), pass its directory with `-d`/`--ensemble-dir`.
@@ -92,11 +92,11 @@ from chebifier.cli import build_base_learners, build_ensemble_model
 from chebifier.predict import predict
 from chebifier.utils import download_ensemble_calibration
 
-# Base learners from the "eval" config ("web" or a path to your own config also work).
-base_learners = build_base_learners("eval")
+# Base learners from the "web" config ("eval" or a path to your own config also work).
+base_learners = build_base_learners("web")
 # download_ensemble_calibration() fetches the standard calibration from Hugging Face; pass your own
 # directory instead to use a calibration you built yourself.
-ensemble = build_ensemble_model("wmv-f1", download_ensemble_calibration(), "eval")
+ensemble = build_ensemble_model("wmv-f1", download_ensemble_calibration(), "web")
 
 smiles_list = ["CC(=O)OC1=CC=CC=C1C(=O)O", "C1=CC=C(C=C1)C(=O)O"]
 result = predict(base_learners, ensemble, smiles_list)

chebifier/cli.py

@@ -205,7 +205,7 @@ def base_learner_options(command):
                 type=str,
                 default=None,
                 help="Ensemble configuration: 'web' or 'eval' (downloaded from Hugging Face) or a "
-                "path to a custom config file listing the base learners (default: eval)",
+                "path to a custom config file listing the base learners (default: web)",
             ),
             click.option(
                 "--prediction-cache-dir",

chebifier/prediction_models/c3p_predictor.py

@@ -4,6 +4,7 @@
 
 import tqdm
 from chebi_utils.read_molecule import smiles_or_inchi_to_mol
+from rdkit import Chem
 
 from chebifier import modelwise_smiles_lru_cache
 from chebifier.prediction_models import BasePredictor
@@ -59,12 +60,16 @@ def __init__(
         self.keep_classes_outside_graph = keep_classes_outside_graph
 
     @modelwise_smiles_lru_cache.batch_decorator
-    def predict_list(self, smiles_list: list[str]) -> list:
+    def predict_list(self, smiles_list: list[str | Chem.Mol]) -> list:
         from c3p import classifier as c3p_classifier
 
         _patch_c3p(c3p_classifier)
-        # C3P only takes SMILES, while the evaluation datasets hand out RDKit molecules
-        mol_list = [smiles_or_inchi_to_mol(smiles) for smiles in smiles_list]
+        # smiles_list is not named correctly - it can by an InChI, SMILES or mol object
+        # -> convert all to Mol, then to SMILES
+        mol_list = [
+            smiles_or_inchi_to_mol(smiles) if isinstance(smiles, str) else smiles
+            for smiles in smiles_list
+        ]
         smiles_list = [to_smiles(molecule) for molecule in mol_list]
         result_list = []
         for batch_start in tqdm.tqdm(
@@ -81,10 +86,7 @@ def predict_list(self, smiles_list: list[str]) -> list:
             )
 
         # Look up the position of each SMILES via a dict instead of scanning smiles_list
-        # for every result (C3P returns one result per class and molecule, so the scan
-        # made reformatting quadratic in the number of molecules). Repeated SMILES map to
-        # all of their positions, which list.index could not do (it always returned the
-        # first one, leaving the later rows without any predictions).
+        # for every result
         indices_by_smiles: dict[str, list[int]] = {}
         for idx, smiles in enumerate(smiles_list):
             indices_by_smiles.setdefault(smiles, []).append(idx)

chebifier/utils.py

@@ -68,10 +68,10 @@ def load_ensemble_config(ensemble_config=None):
     """Resolve an ensemble configuration to a config dict.
 
     'web' and 'eval' are downloaded from the chebifier Hugging Face dataset, anything else is
-    treated as a path to a config file. None defaults to 'eval'.
+    treated as a path to a config file. None defaults to 'web'.
     """
     if ensemble_config is None:
-        ensemble_config = "eval"
+        ensemble_config = "web"
     if ensemble_config in DEFAULT_CONFIGS:
         filename = DEFAULT_CONFIGS[ensemble_config]
         print(

configs/example_config.yml

@@ -1,24 +1,41 @@
+chemlog:
+  type: chemlog
+  model_weight: 100
 
-chemlog_peptides:
-    type: chemlog_peptides
-    model_weight: 100 # if chemlog is available, it always gets chosen
-my_resgated:
-  type: resgated
-  ckpt_path: my_resgated.ckpt # checkpoint trained with chebai
-  molecular_properties: # list of properties used during training
-      - chebai_graph.preprocessing.properties.AtomType
-      - chebai_graph.preprocessing.properties.NumAtomBonds
-      - chebai_graph.preprocessing.properties.AtomCharge
-      - chebai_graph.preprocessing.properties.AtomAromaticity
-      - chebai_graph.preprocessing.properties.AtomHybridization
-      - chebai_graph.preprocessing.properties.AtomNumHs
-      - chebai_graph.preprocessing.properties.BondType
-      - chebai_graph.preprocessing.properties.BondInRing
-      - chebai_graph.preprocessing.properties.BondAromaticity
-      - chebai_graph.preprocessing.properties.RDKit2DNormalized
-  #classwise_weights_path: my_resgated_metrics.json # can be calculated with chebai.results.generate_class_properties
+my_gat_aug:
+  type: gat
+  ckpt_path: my_gat_aug.ckpt
+  dataset_cls: chebai_graph.preprocessing.datasets.ChEBI25_WFGE_WGN_AsPerNodeType
+  molecular_properties:
+    - chebai_graph.preprocessing.properties.AtomNodeLevel
+    - chebai_graph.preprocessing.properties.AugAtomAromaticity
+    - chebai_graph.preprocessing.properties.AugAtomCharge
+    - chebai_graph.preprocessing.properties.AugAtomHybridization
+    - chebai_graph.preprocessing.properties.AugAtomNumHs
+    - chebai_graph.preprocessing.properties.AugAtomType
+    - chebai_graph.preprocessing.properties.AugNumAtomBonds
+    - chebai_graph.preprocessing.properties.AtomFunctionalGroup
+    - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG
+    - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG
+    - chebai_graph.preprocessing.properties.IsFGAlkyl
+    - chebai_graph.preprocessing.properties.AugRDKit2DNormalized
+    - chebai_graph.preprocessing.properties.BondLevel
+    - chebai_graph.preprocessing.properties.AugBondAromaticity
+    - chebai_graph.preprocessing.properties.AugBondInRing
+    - chebai_graph.preprocessing.properties.AugBondType
 
-my_electra:
-  type: electra
-  ckpt_path: my_electra.ckpt
-  #classwise_weights_path: my_electra_metrics.json # can be calculated with chebai.results.generate_class_properties
+my_gat:
+  type: gat
+  ckpt_path: my_gat.ckpt
+  dataset_cls: chebai_graph.preprocessing.datasets.ChEBI25GraphProperties
+  molecular_properties:
+    - chebai_graph.preprocessing.properties.AtomType
+    - chebai_graph.preprocessing.properties.NumAtomBonds
+    - chebai_graph.preprocessing.properties.AtomCharge
+    - chebai_graph.preprocessing.properties.AtomAromaticity
+    - chebai_graph.preprocessing.properties.AtomHybridization
+    - chebai_graph.preprocessing.properties.AtomNumHs
+    - chebai_graph.preprocessing.properties.BondType
+    - chebai_graph.preprocessing.properties.BondInRing
+    - chebai_graph.preprocessing.properties.BondAromaticity
+    - chebai_graph.preprocessing.properties.RDKit2DNormalized