diff --git a/README.md b/README.md index c442872..af5e4c9 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) @@ -151,12 +151,13 @@ Currently, the following models are supported: | Model | Description | #Classes | Publication | Repository | |-------|-------------|----------|-----------------------------------------------------------------------|----------------------------------------------------------------------------------------| -| `electra` | A transformer-based deep learning model trained on ChEBI SMILES strings. | 1531* | [Glauer, Martin, et al., 2024: Chebifier: Automating semantic classification in ChEBI to accelerate data-driven discovery, Digital Discovery 3 (2024) 896-907](https://pubs.rsc.org/en/content/articlehtml/2024/dd/d3dd00238a) | [python-chebai](https://github.com/ChEB-AI/python-chebai) | -| `resgated` | A Residual Gated Graph Convolutional Network trained on ChEBI molecules. | 1531* | | [python-chebai-graph](https://github.com/ChEB-AI/python-chebai-graph) | -| `gat` | A Graph Attention Network trained on ChEBI molecules. | 1531* | | [python-chebai-graph](https://github.com/ChEB-AI/python-chebai-graph) | -| `chemlog_peptides` | A rule-based model specialised on peptide classes. | 18 | [Flügel, Simon, et al., 2025: ChemLog: Making MSOL Viable for Ontological Classification and Learning, arXiv](https://arxiv.org/abs/2507.13987) | [chemlog-peptides](https://github.com/sfluegel05/chemlog-peptides) | -| `chemlog_element`, `chemlog_organox` | Extensions of ChemLog for classes that are defined either by the presence of a specific element or by the presence of an organic bond. | 118 + 37 | | [chemlog-extra](https://github.com/ChEB-AI/chemlog-extra) | -| `c3p` | A collection _Chemical Classifier Programs_, generated by LLMs based on the natural language definitions of ChEBI classes. | 338 | [Mungall, Christopher J., et al., 2025: Chemical classification program synthesis using generative artificial intelligence, Journal of Cheminsformatics](https://link.springer.com/article/10.1186/s13321-025-01092-3) | [c3p](https://github.com/chemkg/c3p) | +| `electra` | A transformer-based deep learning model trained on ChEBI SMILES strings. | 1,766/2,117* | [Glauer, Martin, et al., 2024: Chebifier: Automating semantic classification in ChEBI to accelerate data-driven discovery, Digital Discovery 3 (2024) 896-907](https://pubs.rsc.org/en/content/articlehtml/2024/dd/d3dd00238a) | [python-chebai](https://github.com/ChEB-AI/python-chebai) | +| `resgated` | A Residual Gated Graph Convolutional Network trained on ChEBI molecules. | 1,766/2,117* | [Khedekar, Aditya Ganesh, 2026: Integrating Chemical Knowledge into Graph Neural Networks, Master Thesis](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf) | [python-chebai-graph](https://github.com/ChEB-AI/python-chebai-graph) | +| `gat` | A Graph Attention Network trained on ChEBI molecules. | 1,766/2,117* | [Khedekar, Aditya Ganesh, 2026: Integrating Chemical Knowledge into Graph Neural Networks, Master Thesis](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf) | [python-chebai-graph](https://github.com/ChEB-AI/python-chebai-graph) | +| `chemlog_peptides` | A rule-based model specialised on peptide classes. | 18 | [Flügel, Simon, et al., 2026: Defining Peptides in ChEBI, Jorunal of Cheminformatics](https://link.springer.com/article/10.1186/s13321-026-01196-4) | [chemlog-peptides](https://github.com/sfluegel05/chemlog-peptides) | +| `chemlog_element`, `chemlog_organox` | Extensions of ChemLog for classes that are defined either by the presence of a specific element or by the presence of an organic bond. | 118 + 37 | [Flügel, Simon, et al., 2025: ChemLog: Making MSOL Viable for Ontological Classification and Learning, arXiv](https://arxiv.org/abs/2507.13987) | [chemlog-extra](https://github.com/ChEB-AI/chemlog-extra) | +| `c3p` | A collection _Chemical Classifier Programs_, generated by LLMs based on the natural language definitions of ChEBI classes. | 338 | [Mungall, Christopher J., et al., 2025: Chemical classification program synthesis using generative artificial intelligence, Journal of Cheminformatics](https://link.springer.com/article/10.1186/s13321-025-01092-3) | [c3p](https://github.com/chemkg/c3p) | +| `lopster` | Rules for 36 ChEBI classes, focusing on classes that cannot be expressed in OWL | 36 | [Magka, Despoina, et al., 2014: A rule-based ontological framework for the classification of molecules, Journal of Biomedical Semantics](https://link.springer.com/article/10.1186/2041-1480-5-17) | [original implementation](https://github.com/magkades/lopster) - Chebifier uses an updated version integrated into [chemlog](https://github.com/sfluegel05/chemlog-peptides) | In addition, Chebifier also includes a ChEBI lookup that automatically retrieves the ChEBI superclasses for a class matched by a SMILES string. This is not activated by default, but can be included by adding @@ -168,196 +169,36 @@ chebi_lookup: to your configuration file. ### The ensemble -For an extended description of the ensemble, see [Flügel, Simon, et al., 2025: Chebifier 2: An Ensemble for Chemistry](https://ceur-ws.org/Vol-4064/SymGenAI4Sci-paper4.pdf). - -ensemble_architecture - -Given a sample (i.e., a SMILES string) and models $m_1, m_2, \ldots, m_n$, the ensemble works as follows: -1. Get predictions from each model $m_i$ for the sample. -2. For each class $c$, aggregate predictions $p_c^{m_i}$ from all models that made a prediction for that class. -The aggregation happens separately for all positive predictions (i.e., $p_c^{m_i} \geq 0.5$) and all negative predictions -($p_c^{m_i} < 0.5$). If the aggregated value is larger for the positive predictions than for the negative predictions, -the ensemble makes a positive prediction for class $c$: - -image - - -Here, confidence is the model's (self-reported) confidence in its prediction. Each model has its own -decision threshold $t_{m_i}$, calibrated on the validation set (see below), and confidence measures -how far the prediction sits from that threshold — scaled separately on each side, so that a -maximally confident negative ($p = 0$) and a maximally confident positive ($p = 1$) both count 1: -$ -\text{confidence}_c^{m_i} = \begin{cases} -(t_{m_i} - p_c^{m_i}) / t_{m_i} & \text{if } p_c^{m_i} < t_{m_i} \\ -(p_c^{m_i} - t_{m_i}) / (1 - t_{m_i}) & \text{otherwise} -\end{cases} -$ -For example, for a model with $t_{m_i} = 0.5$ and a positive prediction of $p_c^{m_i} = 0.55$, the -confidence is $(0.55 - 0.5)/0.5 = 0.1$. One could say that the model is not very confident in its -prediction and very close to switching to a negative prediction. If another model is very sure about -its negative prediction with $p_c^{m_j} = 0.1$ (and $t_{m_j} = 0.5$), the confidence is -$(0.5 - 0.1)/0.5 = 0.8$. Therefore, if in doubt, we are more confident in the negative prediction. - -The two-sided scaling matters whenever a model's threshold is not 0.5: with $t_{m_i} = 0.2$, a -negative prediction only has a range of $0.2$ to move in and a positive one a range of $0.8$, so -without rescaling the positive side would systematically outweigh the negative side. - -Confidence is used by the weighted voting ensembles (`wmv-conf` and `wmv-f1`). If the `ensemble_type` -is set to `mv`, all votes count the same (confidence is fixed to 1), which gives an unweighted -majority-voting baseline. - -The`model_weight` can be set for each model in the configuration file (default: 1). This is used to favor a certain -model independently of a given class. -`Trust` is based on the model's performance on a validation set. After training, we evaluate the Machine Learning models -on a validation set for each class. If the `ensemble_type` is set to `wmv-f1`, the trust is calculated as F1-score $^{6.25}$. -For `mv` and `wmv-conf`, the trust is set to 1 for all models. - -#### Learned aggregation (`ltr` and `des`) - -Two further `ensemble_type`s replace the fixed voting rule by a model that is fitted on the -validation split. Both restrict themselves to a candidate set (per molecule, the union of each -base learner's top-`candidate_k` classes) and both emit the same net score as the voting -ensembles, so inconsistency resolution and the decision threshold apply unchanged. - -- `ltr` — **learning to rank**, an adaptation of - [GOLabeler](https://doi.org/10.1093/bioinformatics/bty130): the base learner scores for a - (molecule, class) pair become the feature vector of a LambdaMART ranker (LightGBM) that ranks - ChEBI classes per molecule. Features are the raw base learner scores plus the number of covering - models and the max/mean/std over them; a global cutoff on the ranker score is calibrated on a - held-out 20% of the validation split. Feature column *j* is always base learner *j*, so the - ranker can learn which model to trust — but the raw scores say nothing about the class being - scored. `class_stats` (on by default) adds that: one column per base learner holding its - validation F1 *for this class* (the same quantity `wmv-f1` weights by), plus the class prevalence - and its number of positives. To keep the labels of the scored molecules out of the features, the - statistics used during training are estimated on the training molecules only, while prediction - uses the statistics of the whole validation split. Set `class_stats=False` for the plain - GOLabeler feature set; that also skips the per-model threshold calibration the F1 scores need. -- `des` — **dynamic ensemble selection**, an adaptation of - [META-DES.H](https://arxiv.org/pdf/1811.01742): a `GaussianNB` meta-classifier estimates, per - (molecule, class, base learner), how competent that base learner is *for this molecule*, and only - the competent ones vote, weighted by that competence. Competence is described by the paper's five - meta-feature sets over two neighbourhoods — the `region_size` nearest molecules by Tanimoto - similarity on ECFP4, and the `profile_size` nearest output profiles. Because the neighbourhoods - are looked up at prediction time, calibration stores the reference predictions, labels and - fingerprints in the ensemble directory (~1 GB for a 20-model ensemble on ChEBI50). - The meta-features are otherwise purely behavioural — one meta-classifier is fitted over all - (molecule, class, base learner) rows pooled, and the paper's input identifies neither the base - learner nor the class, so competence is a function of local track record alone. `use_model_id` - (on by default) appends a one-hot encoding of the base learner, which lets the meta-classifier - express "model A is the stronger one here" instead of only "whichever model this is, it behaves - like *this*"; `use_model_id=False` restores the published feature set. - `meta_classifier="mlp"` replaces `GaussianNB` with a standardised two-layer `MLPClassifier`, - which drops the feature-independence assumption — a poor fit for these meta-features, since the - `region_size` correctness flags are strongly correlated with each other and with their own mean. - The MLP is fitted in one pass over the meta-training set rather than chunk-wise, which the - consensus filter keeps small (~130k rows for 8 base learners on ChEBI25 3-STAR); - `max_meta_samples` caps it if a larger ensemble overflows memory. - Two further options control the reference set rather than the meta-classifier. - `morgan_radius` / `morgan_bits` / `morgan_chirality` set the fingerprint the region of competence - is measured on. Plain ECFP4 cannot separate stereoisomers, which are distinct ChEBI classes, so - 6.6% of ChEBI25 3-STAR validation molecules share a fingerprint with one carrying different - labels; `morgan_chirality` is therefore on by default, which halves that to 3.9%. Widening - `morgan_bits` changes nothing — the degeneracy is structural, not hash collisions. - `full_dsel=True` stores the whole validation split as the reference set instead of only the 80% - that the meta-classifier is fitted on, for denser neighbourhoods at prediction time. - - Note that the region of competence excludes the query molecule itself during calibration but - not during prediction, where the query is genuinely unseen. Predicting for the validation split - therefore lets ~80% of molecules retrieve themselves as their own nearest neighbour, which makes - any validation-split metric for `des` optimistic. Use the test split. - -Both calibrate their hyperparameters by 5-fold cross-validation on the validation split, scoring -macro-F1 on each held-out fold (the cutoff is tuned on a fold-internal dev set, so the reported -score is not tuned on the fold it is measured on). Only the parameters that moved the result in -previous experiments are searched: `candidate_k` for `ltr`, and `region_size` / `profile_size` / -`vote` for `des`. The ranker's own tree hyperparameters, and `des`'s consensus and competence -thresholds, sit on a plateau and are left at their published values. Passing any searched parameter -to the constructor skips the search for it — `chebifier build` takes constructor arguments as -`-ep key=value`, e.g. -`-ep candidate_k=70 -ep class_stats=1` or `-ep region_size=7 -ep meta_classifier=mlp`. Arguments -that change the stored model are recorded in the ensemble's metadata, so `chebifier evaluate` picks -them up on its own. `scripts/reproduce_ablation_3star.ps1` compares the optional features above -against their baselines this way. Results are written to `hyperparameter_search.csv` and -`best_hyperparameters.csv` in the ensemble directory, as for `wmv-f1`. +The ensemble collects per-class scores from every base learner and turns them into one decision per +class, selected with `-t`/`--ensemble-type`. For an extended description, see +[Flügel, Simon, et al., 2025: Chebifier 2: An Ensemble for Chemistry](https://ceur-ws.org/Vol-4064/SymGenAI4Sci-paper4.pdf). + +ensemble architecture overview + + +| Strategy | How it works | +|----------|--------------| +| `mv` | Plain majority vote; every model that predicted a class counts equally. | +| `wmv-conf` | Majority vote weighted by each model's confidence, i.e. how far its score sits from its calibrated decision threshold (scaled per side so a maximally confident positive and negative both count 1). | +| `wmv-f1` (default) | Confidence weighting plus a per-class trust term, the model's validation F1 raised to the power 6.25. | +| `ltr` | A LambdaMART ranker (adapting [GOLabeler](https://doi.org/10.1093/bioinformatics/bty130)) fitted on the validation split ranks classes per molecule from the base learner scores. Optionally adds per-class validation statistics as features (`class_stats`). | +| `des` | Dynamic ensemble selection (adapting [META-DES.H](https://arxiv.org/pdf/1811.01742)): a meta-classifier estimates each base learner's local competence per molecule, and only the competent ones vote. | + +Each model also carries a `model_weight` (configurable, default 1) that scales its vote independently +of the class. `ltr` and `des` calibrate their hyperparameters by 5-fold cross-validation on the +validation split; `chebifier build` takes their constructor arguments as `-ep key=value`. All +strategies emit the same net score, so inconsistency resolution and the decision threshold apply +unchanged. ### Inconsistency resolution -After a decision has been made for each class independently, the consistency of the predictions with regard to the ChEBI hierarchy -and disjointness axioms is checked. This is -done in 3 steps: -- (1) First, the hierarchy is corrected. For each pair of classes $A$ and $B$ where $A$ is a subclass of $B$ (following -the is-a relation in ChEBI), we set the ensemble prediction of $A$ to that of $B$ if $B$ is the more -_confident_ of the two, and $B$ to that of $A$ otherwise. Confidence is the distance from the decision -threshold, scaled separately on each side of it so that a maximally confident negative and a maximally -confident positive both count $1$ — the same measure `wmv-conf` weights its votes by. For example, if -$A$ scores $0.6$ and $B$ scores $0.1$ at a threshold of $0.5$, $B$ is the more confident one -($0.8$ against $0.2$), so $A$ is lowered to $0.1$ and neither class is predicted. -- (2) Next, we check for disjointness. This is not specified directly in ChEBI, but in an additional ChEBI module ([chebi-disjoints.owl](https://ftp.ebi.ac.uk/pub/databases/chebi/ontology/)). -We have extracted these disjointness axioms into a CSV file and added some more disjointness axioms ourselves (see -`data>disjoint_chebi.csv` and `data>disjoint_additional.csv`). If two classes $A$ and $B$ are disjoint and we predict -both, we select one with the higher class score and set the other to 0. -- (3) Since the second step might have introduced new inconsistencies into the hierarchy, we repeat the first step, but -with a small change. For a pair of classes $A \subseteq B$ with predictions $1$ and $0$, instead of setting $B$ to $1$, -we now set $A$ to $0$. This has the advantage that we cannot introduce new disjointness-inconsistencies and don't have -to repeat step 2. - -#### Alternative methods - -The method above is `--inconsistency-resolution score-based` (`-ir`, the default). Two alternative -families from the literature are available at the same point in the pipeline; all of them consume a -net score and return a net score, so the decision threshold applies unchanged. Scores are -probabilities in $[0, 1]$, with $0.5$ meaning "undecided"; the decision itself is made at the -operating point the ensemble reports as `decision_threshold`, which is not always $0.5$. - -- `ilr-godel`, `ilr-lukasiewicz` — **Iterative Local Refinement** - ([Daniele et al. 2023](https://doi.org/10.1007/s10994-023-06310-3)). Subsumption becomes the - implication $A \rightarrow B$ and disjointness the formula $\neg (A \wedge B)$, both as hard - constraints ($\hat t = 1$). Each constraint is repaired by its *minimal refinement function* — - the closest truth vector satisfying it — and the repairs are iterated to a fixpoint instead of - running the fixed 3-step schedule above. The two variants differ in how they split a violation: - Gödel is winner-take-all (it raises the parent to the child, and zeroes the weaker side of a - disjoint pair), whereas Łukasiewicz shares the correction — a disjointness violation with scores - $0.8$ and $0.7$ becomes $0.55$ and $0.45$ rather than $0.8$ and $0$. -- `hex` — **HEX graphs** - ([Deng et al. 2014](https://doi.org/10.1007/978-3-319-10590-1_4)). A CRF over binary label - vectors in which hierarchy edges forbid $(B, A) = (0, 1)$ and exclusion edges forbid - $(1, 1)$. Illegal states have probability zero, so the marginals satisfy - $P(A) \le P(B)$ for $A \subseteq B$ and $P(A) + P(B) \le 1$ for disjoint $A, B$ by construction. - -`ilr-godel` and `ilr-lukasiewicz` are tuned with `alpha`, `max_iter` and `tol`, passed as -`-irp alpha=0.5`. `scripts/calibrate_resolution.py` grid-searches resolution parameters against a -validation split; the grid per method is defined in its `GRIDS` dict. Note that a monotone -reparametrisation of the scores cannot change `ilr-godel`'s decisions: every Gödel operation is -order-preserving, so it cannot move a score across the boundary. - -#### Why HEX needs an approximation - -Applied as published, HEX inference is intractable here. Its cost is bounded by -$O(\min(|V|2^w, |V|2^{\Omega}))$, and on a 2117-class ChEBI label set the maximum overlap is -$\Omega = 2115$ and the junction tree width is $\le 62$, with over 5 million legal states in the -largest cliques — the paper's efficiency argument assumes labels are mostly mutually exclusive, -whereas ChEBI labels overwhelmingly overlap (~25 classes hold per molecule). Exact junction-tree -inference is therefore not an option at this scale, so `hex` deviates from the published method; -this should be reported as such. - -`hex` (`chebifier/hex_bounded.py`) replaces exact inference with a **branch-and-bound over partial -assignments**. Each search node fixes some labels on and some off, leaving the rest free, and -yields an interval $[\mathrm{lb}, \mathrm{ub}]$ that provably brackets every label's true marginal. -Fixing a label propagates through hierarchy and exclusion edges to a fixpoint, so infeasible -branches are pruned immediately. The node with the largest slack is expanded first, for at most -`budget` expansions (default 2000). If the search exhausts the frontier within that budget the -intervals collapse and the result is exact; otherwise they stay open and the bounds remain valid -but loose. Search also stops early once every label's interval lies entirely on one side of the -decision threshold, since further refinement cannot change any decision. - -The smoother returns the **lower** bound. A label whose interval still straddles the threshold is -therefore decided negative — ties go against predicting the class — and the number of such -labels is accumulated in `n_uncertified`. Pass `budget` and `processes` (molecules are bounded in -parallel across a worker pool) with `-irp budget=4000`. `threshold` defaults to the ensemble's -operating point and only needs to be set explicitly to override it. +After each class has been decided independently, the predictions are reconciled with the ChEBI +hierarchy (is-a) and disjointness axioms (`data/disjoint_chebi.csv` and `data/disjoint_additional.csv`). The method is chosen with +`-ir`/`--inconsistency-resolution`, or disabled with `--no-resolve-inconsistencies`; each consumes a +net score and returns one, so the decision threshold applies unchanged. + +| Method | How it works | +|--------|--------------| +| `score-based` (default) | Repairs hierarchy then disjointness violations by keeping the more confident class of each conflicting pair (confidence = distance from the decision threshold). A final hierarchy pass lowers children rather than raising parents, so no new disjointness conflicts appear. | +| `ilr-godel` | Iterative Local Refinement ([Daniele et al. 2023](https://doi.org/10.1007/s10994-023-06310-3)) with Gödel logic: each subsumption/disjointness constraint is repaired winner-take-all and iterated to a fixpoint. | +| `ilr-lukasiewicz` | The same ILR framework with Łukasiewicz logic, which shares the correction between the two conflicting classes instead of winner-take-all (e.g. scores 0.8/0.7 become 0.55/0.45). | +| `hex` | HEX-graph constrained inference ([Deng et al. 2014](https://doi.org/10.1007/978-3-319-10590-1_4)). Exact inference is intractable on ChEBI's heavily overlapping labels, so this is a bounded branch-and-bound approximation whose intervals decide ties negatively. | diff --git a/chebifier/cli.py b/chebifier/cli.py index 608d150..969058d 100644 --- a/chebifier/cli.py +++ b/chebifier/cli.py @@ -1,4 +1,6 @@ import importlib.resources +import json +import math import os from typing import Literal @@ -6,6 +8,7 @@ import numpy as np import pandas as pd import yaml +from chebi_utils.obo_extractor import get_hierarchy_subgraph from chebi_utils.read_molecule import smiles_or_inchi_to_mol from chebifier.build_ensemble import EnsembleBuilder @@ -32,16 +35,19 @@ def read_molecules(molecules, molecule_file): - """Collect molecules from CLI arguments and/or a file (one molecule per line) and convert them to RDKit mol objects.""" + """Collect molecules from CLI arguments and/or a file (one molecule per line), returning the raw + inputs alongside the RDKit mol objects they were parsed into.""" raw_inputs = list(molecules) if molecule_file: with open(molecule_file, "r", encoding="utf-8") as f: raw_inputs.extend([line.strip() for line in f if line.strip()]) - return [smiles_or_inchi_to_mol(raw_input) for raw_input in raw_inputs] + return raw_inputs, [smiles_or_inchi_to_mol(raw_input) for raw_input in raw_inputs] -def build_base_learners(ensemble_config, prediction_cache_dir=None, split=None): +def build_base_learners( + ensemble_config, prediction_cache_dir=None, split=None, chebi_graph_file=None +): """Instantiate the base learners described by an ensemble configuration file. If prediction_cache_dir and split are given, models whose predictions for that split are @@ -71,7 +77,7 @@ def build_base_learners(ensemble_config, prediction_cache_dir=None, split=None): base_learners[model_name] = None continue if chebi_graph is None: - chebi_graph = load_chebi_graph() + chebi_graph = load_chebi_graph(chebi_graph_file) if "hugging_face" in model_config: hugging_face_kwargs = download_model_files(model_config["hugging_face"]) else: @@ -205,7 +211,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", @@ -323,6 +329,11 @@ def build( ) builder.build_ensemble() + with open( + os.path.join(ensemble_dir, "ensemble_classes.txt"), "w", encoding="utf-8" + ) as f: + f.writelines(f"{cls}\n" for cls in validation_labels.columns) + @cli.command() @base_learner_options @@ -441,7 +452,6 @@ def output_path(dir_, variant): chebi_graph = load_chebi_graph() disjoint_files = get_disjoint_files() - # TODO: Hugging Face support eval_data, eval_labels = load_dataset( data_path, split=split, @@ -542,6 +552,30 @@ def collect_classes(ensemble_config, prediction_cache_dir, data_path, split, out print(f"Saved {len(classes)} classes to {output}.") +def class_name(chebi_graph, chebi_id): + node = chebi_graph.nodes.get(str(chebi_id)) + if node is None or not node.get("name"): + return f"CHEBI:{chebi_id}" + return node["name"] + + +def jsonable(value): + """JSON has no NaN, and a base learner that did not cover a class reports exactly that.""" + value = float(value) + return None if math.isnan(value) else value + + +def most_specific(predicted_classes, hierarchy): + """The predicted classes that have no predicted subclass, i.e. the lowest classes the + prediction reaches in the hierarchy. is-a edges point from child to parent, so the predecessors + of a class are its subclasses.""" + predicted = [cls for cls in predicted_classes if cls in hierarchy] + subgraph = hierarchy.subgraph(predicted) + return [ + cls for cls in predicted if not any(True for _ in subgraph.predecessors(cls)) + ] + + @cli.command() @ensemble_options @click.option( @@ -587,6 +621,21 @@ def collect_classes(ensemble_config, prediction_cache_dir, data_path, split, out help="Probability a class has to exceed to be predicted (default: the operating point the " "ensemble reports)", ) +@click.option( + "--chebi-graph", + "chebi_graph_file", + type=click.Path(exists=True), + default=None, + help="Local ChEBI graph (pickled networkx graph) to run on (default: downloaded from " + "Hugging Face)", +) +@click.option( + "--attribution/--no-attribution", + "explain", + default=False, + help="Include per-class explanations, i.e. the share of the decision each base learner " + "holds (default: False)", +) def predict( ensemble_config, ensemble_type, @@ -598,19 +647,31 @@ def predict( inconsistency_resolution, ir_param, decision_threshold, + chebi_graph_file, + explain, output, ): """Predict ChEBI classes for a list of SMILES / InChI strings.""" - molecules_list = read_molecules(molecules, molecule_file) + raw_inputs, molecules_list = read_molecules(molecules, molecule_file) if not molecules_list: click.echo("No molecules provided. Use --molecules or --molecule-file.") return if ensemble_dir is None: ensemble_dir = download_ensemble_calibration() - base_learners = build_base_learners(ensemble_config) + base_learners = build_base_learners( + ensemble_config, chebi_graph_file=chebi_graph_file + ) ensemble_model = build_ensemble_model(ensemble_type, ensemble_dir, ensemble_config) + classes_file = os.path.join(ensemble_dir, "ensemble_classes.txt") + if not os.path.exists(classes_file): + print( + f"Warning: no ensemble_classes.txt in {ensemble_dir}, deriving classes from " + "the base learner predictions instead." + ) + classes_file = None + predictions = predict_molecules( base_learners, ensemble_model, @@ -620,10 +681,83 @@ def predict( inconsistency_resolution=inconsistency_resolution, inconsistency_resolution_params=parse_ir_params(ir_param), decision_threshold=decision_threshold, + classes=read_classes(classes_file), + attribution=explain, + chebi_graph_file=chebi_graph_file, ) - print(f"Predictions: {predictions}") - # TODO: turn the aggregated predictions into ChEBI classes per molecule, print them / save to output + chebi_graph = load_chebi_graph(chebi_graph_file) + hierarchy = get_hierarchy_subgraph(chebi_graph) + predicted_classes = predictions["predicted_classes"] + class_decisions = predictions["class_decisions"] + complete_failure = predictions["complete_failure"] + net_score = predictions["net_score"] + attribution = predictions.get("attribution") + attribution_models = predictions.get("attribution_models") + base_learner_predictions = predictions.get("base_learner_predictions", {}) + positive = predictions.get("positive_mask") + negative = predictions.get("negative_mask") + + results = [] + for i, raw_input in enumerate(raw_inputs): + if molecules_list[i] is None or complete_failure[i]: + print(f"[{i + 1}] {raw_input}: no prediction") + result = { + "input": raw_input, + "predicted_parents": None, + "direct_parents": None, + } + if explain: + result["explanations"] = None + results.append(result) + continue + class_indices = class_decisions[i].nonzero().flatten().tolist() + predicted = [predicted_classes[j] for j in class_indices] + direct = most_specific(predicted, hierarchy) + direct_set = set(direct) + print( + f"[{i + 1}] {raw_input}: {len(predicted)} predicted class(es), " + f"{len(direct)} most specific (*)" + ) + for cls in predicted: + marker = "*" if cls in direct_set else " " + print(f" {marker} CHEBI:{cls} {class_name(chebi_graph, cls)}") + + result = { + "input": raw_input, + "predicted_parents": predicted, + "direct_parents": [[cls, class_name(chebi_graph, cls)] for cls in direct], + } + if explain: + explanations = {} + for j in class_indices: + cls = predicted_classes[j] + models = {} + if attribution is not None: + for m, model_name in enumerate(attribution_models): + # which way the model voted, against its own threshold; models that did + # not cover the class cast no vote and hold no share of the decision + vote = int(positive[i, j, m]) - int(negative[i, j, m]) + if vote: + models[model_name] = { + "attribution": jsonable(attribution[i, j, m]), + "vote": vote, + "prediction": jsonable( + base_learner_predictions[model_name][i, j] + ), + } + explanations[cls] = { + "name": class_name(chebi_graph, cls), + "score": jsonable(net_score[i, j]), + "models": models, + } + result["explanations"] = explanations + results.append(result) + + if output is not None: + with open(output, "w", encoding="utf-8") as f: + json.dump(results, f, indent=2) + print(f"Saved predictions for {len(results)} molecules to {output}.") if __name__ == "__main__": diff --git a/chebifier/ensemble/weighted_majority_ensemble.py b/chebifier/ensemble/weighted_majority_ensemble.py index 2e48268..a886cf6 100644 --- a/chebifier/ensemble/weighted_majority_ensemble.py +++ b/chebifier/ensemble/weighted_majority_ensemble.py @@ -67,16 +67,22 @@ def _save_classwise_f1(self, validation_predictions, validation_labels): ) def _load_classwise_f1(self, model_name: str, num_classes: int) -> torch.Tensor: - if self.model_f1_scores is not None: + if self.model_f1_scores is not None and model_name in self.model_f1_scores: return self.model_f1_scores[model_name] classwise_f1_path = Path(self.ensemble_dir) / f"{model_name}_classwise_f1.txt" if classwise_f1_path.exists(): with open(classwise_f1_path, "r", encoding="utf-8") as f: - return torch.tensor([float(x) for x in f.read().splitlines()]) + f1 = torch.tensor([float(x) for x in f.read().splitlines()]) + self.model_f1_scores = self.model_f1_scores or {} + self.model_f1_scores[model_name] = f1 + return f1 print( f"No class-wise F1 scores for {model_name} in {self.ensemble_dir}, voting with full trust." ) - return torch.ones(num_classes) + f1 = torch.ones(num_classes) + self.model_f1_scores = self.model_f1_scores or {} + self.model_f1_scores[model_name] = f1 + return f1 def _load_hyperparameters(self) -> tuple[float, int]: """Hyperparameters set explicitly take precedence, otherwise the optimal values found during @@ -160,6 +166,8 @@ def _score_hyperparameters( self.weighting_strength = weighting_strength self.weighting_exponent = weighting_exponent scores = [] + old_model_f1_scores = self.model_f1_scores + old_thresholds = self.prediction_thresholds for thresholds, classwise_f1, test_idx in folds: self.prediction_thresholds = thresholds self.model_f1_scores = classwise_f1 @@ -178,8 +186,8 @@ def _score_hyperparameters( fold_labels = validation_labels[test_idx] fold_f1 = self.classwise_f1(decisions, fold_labels) scores.append(fold_f1[fold_labels.sum(dim=0) > 0].mean().item()) - self.prediction_thresholds = None - self.model_f1_scores = None + self.prediction_thresholds = old_thresholds + self.model_f1_scores = old_model_f1_scores return scores def _optimize_hyperparameters(self, validation_predictions, validation_labels): diff --git a/chebifier/predict.py b/chebifier/predict.py index af4ef9c..409ff8c 100644 --- a/chebifier/predict.py +++ b/chebifier/predict.py @@ -24,7 +24,8 @@ def apply_inconsistency_resolution( """Resolve inconsistencies in batches - the smoother materialises a (batch_size, n_classes, n_classes) tensor, which does not fit into memory for a whole dataset split. """ - smoother.set_label_names(class_names) + if getattr(smoother, "label_names", None) != class_names: + smoother.set_label_names(class_names) net_score = aggregated_predictions["net_score"] valid = aggregated_predictions.get("has_valid_predictions") attribution = aggregated_predictions.get("attribution") @@ -51,6 +52,36 @@ def apply_inconsistency_resolution( return aggregated_predictions +_SMOOTHER_CACHE = {} + + +def get_smoother(inconsistency_resolution, chebi_graph, disjoint_files, params): + """Build a smoother, reusing a previously built one where possible. + + Building a smoother parses the disjointness files and (once the label names are known) + computes the transitive closure of the hierarchy, which is far more expensive than the + resolution itself - and none of it depends on the molecules being predicted. + """ + key = ( + inconsistency_resolution, + id(chebi_graph), + tuple(str(file) for file in disjoint_files), + repr(sorted(params.items())), + ) + if key not in _SMOOTHER_CACHE: + # the graph is kept alive alongside the smoother so that its id stays unique + _SMOOTHER_CACHE[key] = ( + chebi_graph, + get_smoother_class(inconsistency_resolution)( + chebi_graph=chebi_graph, + label_names=None, + disjoint_files=disjoint_files, + **params, + ), + ) + return _SMOOTHER_CACHE[key][1] + + def base_learner_cache_path( prediction_cache_dir: str, model_name: str, split: str ) -> str: @@ -188,6 +219,8 @@ def aggregate_predictions( aggregated_predictions = ensemble_model.predict( test_predictions, molecules, **({"attribution": True} if attribution else {}) ) + if attribution: + aggregated_predictions["base_learner_predictions"] = test_predictions # net_score, has_valid_predictions, intermediate_results_dict return aggregated_predictions, predicted_classes @@ -199,6 +232,7 @@ def resolve_and_decide( inconsistency_resolution_params: Optional[dict] = None, decision_threshold: float = 0.5, chebi_graph=None, + chebi_graph_file: Optional[str] = None, disjoint_files=None, ) -> dict: """Resolve inconsistencies in aggregated predictions and turn them into class decisions. @@ -208,20 +242,18 @@ def resolve_and_decide( `aggregated_predictions` is not modified, so the same aggregation can be passed to several resolution variants. Pass `inconsistency_resolution=None` or "none" to skip the resolution, - and chebi_graph / disjoint_files to avoid reloading them for every variant. + and chebi_graph / disjoint_files to avoid reloading them for every variant. `chebi_graph_file` + loads the hierarchy from a local file instead of Hugging Face. """ aggregated_predictions = dict(aggregated_predictions) if inconsistency_resolution not in (None, "none"): if chebi_graph is None: - chebi_graph = load_chebi_graph() + chebi_graph = load_chebi_graph(chebi_graph_file) if disjoint_files is None: disjoint_files = get_disjoint_files() params = inconsistency_resolution_params or {} - smoother = get_smoother_class(inconsistency_resolution)( - chebi_graph=chebi_graph, - label_names=None, - disjoint_files=disjoint_files, - **params, + smoother = get_smoother( + inconsistency_resolution, chebi_graph, disjoint_files, params ) if "threshold" not in params and hasattr(smoother, "threshold"): smoother.threshold = decision_threshold @@ -261,6 +293,7 @@ def predict( classes: Optional[list[str]] = None, split: str = "test", attribution: bool = False, + chebi_graph_file: Optional[str] = None, ) -> dict: """ Get end-to-end predictions from base learners and an ensemble model. @@ -281,8 +314,13 @@ def predict( split (str): Name of the dataset split, used to separate cached base learner predictions of different splits within the same cache directory. attribution (bool): Also report, per (molecule, class), the share of the decision each base - learner is responsible for (summing to 1 over the base learners). Only supported by the - voting ensembles and the score-based inconsistency resolution. + learner is responsible for (summing to 1 over the base learners), together with the raw + base learner predictions it was derived from (`base_learner_predictions`, one + (num_molecules, num_classes) tensor per model). Only supported by the voting ensembles + and the score-based inconsistency resolution. + chebi_graph_file (Optional[str]): Local ChEBI graph (pickled networkx graph) the + inconsistency resolution runs on. If None (the default), it is downloaded from + Hugging Face. Returns: dict: A dictionary containing the final predictions and optionally the smoothed predictions. @@ -303,6 +341,7 @@ def predict( inconsistency_resolution if resolve_inconsistencies else "none" ), inconsistency_resolution_params=inconsistency_resolution_params, + chebi_graph_file=chebi_graph_file, decision_threshold=( ensemble_model.decision_threshold if decision_threshold is None 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/prediction_models/chebi_lookup.py b/chebifier/prediction_models/chebi_lookup.py index 4f6dcd6..58a8db6 100644 --- a/chebifier/prediction_models/chebi_lookup.py +++ b/chebifier/prediction_models/chebi_lookup.py @@ -25,7 +25,7 @@ def __init__( or "ChEBI Lookup: If the SMILES is equivalent to a ChEBI entry, retrieve the classification of that entry." ) self.chebi_version = chebi_version - self.chebi_graph = kwargs.get("chebi_graph", load_chebi_graph()) + self.chebi_graph = kwargs.get("chebi_graph") or load_chebi_graph() self.lookup_table = self.get_inchikey_lookup() def get_inchikey_lookup(self): diff --git a/chebifier/utils.py b/chebifier/utils.py index 0734a61..4af7443 100644 --- a/chebifier/utils.py +++ b/chebifier/utils.py @@ -14,8 +14,13 @@ CHEBI_VERSION = 252 +@functools.lru_cache(maxsize=None) def load_chebi_graph(filename=None): - """Load ChEBI graph from Hugging Face (if filename is None) or local file""" + """Load ChEBI graph from Hugging Face (if filename is None) or local file. + + Cached: unpickling the graph is slow, and callers that key their own caches on the graph + (get_superclasses, predict.get_smoother) only hit them if they get the same object back. + """ if filename is None: print("Loading ChEBI graph from Hugging Face...") file = download_model_files( @@ -68,10 +73,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 diff --git a/pyproject.toml b/pyproject.toml index 67a1900..3b21543 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,7 @@ chebifier = ["*.yml"] dev = ["black", "isort", "pre-commit"] models = [ - "chebai_graph==1.2.0", + "chebai_graph==1.2.3", "chebai==1.4.0", "chemlog-extra==1.1", # includes ChemLog / Lopster "c3p"