diff --git a/README.md b/README.md index c442872..5aeaeb5 100644 --- a/README.md +++ b/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) diff --git a/chebifier/cli.py b/chebifier/cli.py index 608d150..57cb377 100644 --- a/chebifier/cli.py +++ b/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", diff --git a/chebifier/prediction_models/c3p_predictor.py b/chebifier/prediction_models/c3p_predictor.py index e8b3f6b..2053072 100644 --- a/chebifier/prediction_models/c3p_predictor.py +++ b/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) diff --git a/chebifier/utils.py b/chebifier/utils.py index 0734a61..8513707 100644 --- a/chebifier/utils.py +++ b/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( diff --git a/configs/example_config.yml b/configs/example_config.yml index 3a67698..44d3409 100644 --- a/configs/example_config.yml +++ b/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