Update README.md

README.md

@@ -18,7 +18,154 @@ language:
 - en
 ---
 
-# MolE - Antimicrobial Prediction
+# MolE -- Antimicrobial Prediction
+
+## Short description
+
+MolE learns task-independent molecular representations of chemicals vis Graph Isomorphism Networks (GINs). Combined with an XGBoost classifier it estimates the probability of a compound
+inhibiting bacterial growth. The model was developed by Roberto Olayo Alarcon et al. and more information can be found in the GitHub repository and the accompanying paper.
+
+
+## Model versions
+
+-   **MolE Antimicrobial Prediction** (March 2024)
+    Pretrained representation model + XGBoost classifier trained on
+    antimicrobial screening data (Maier et al., 2018).
+    Repository: <https://github.com/rolayoalarcon/mole_antimicrobial_potential>\
+    Hugging Face Hub: `virtual-human-chc/MolE`
+
+
+## Long description
+
+MolE integrates molecular graph-based representation learning with
+gradient-boosted decision trees for predicting antimicrobial potential.
+The approach involves: 1. **Representation learning:** A graph neural
+network (GINet) trained on 100,000 randomly sampled compounds to derive
+molecular embeddings from SMILES strings. 2. **Prediction:** These
+embeddings are used as input to an **XGBoost** model that predicts
+antimicrobial activity scores across 40 bacterial strains, based on data from *Maier et al., 2018*.
+
+The model was developed by **Roberto Olayo Alarcon et al.**.
+Further information is available in the [paper](https://www.nature.com/articles/s41467-025-58804-4).
+
+
+## Metadata
+
+### Input
+
+-   **Description:** Table of chemicals with their SMILES
+    representations.
+-   **Shape:** `[n, 2]`
+-   **Data format:**
+    -   `chem_name` *(str)*: Name of the molecule (e.g., *Halicin*).
+    -   `smiles` *(str)*: SMILES representation of the molecule (e.g.,
+        `C1=C(SC(=N1)SC2=NN=C(S2)N)[N+](=O)[O-]`).
+-   **Example input file:** `examples/input/examples_molecules.tsv`
+
+
+### Model
+
+-   **Modality:** Probability of microbial growth inhibition.
+-   **Scale:** Combination of bacterial strain and compound.
+-   **Description:**
+    -   The model computes antimicrobial predictive probabilities for **40 bacterial strains** contained in [Maier et al., 2018](https://www.nature.com/articles/nature25979), using a two-step process:
+        1.  Generate molecular embeddings with a pre-trained **GINet**
+            representation model (`model.pth`, `config.yaml`).
+        2.  Predict antimicrobial properties with a **trained XGBoost**
+            classifier (`MolE-XGBoost-08.03.2024_14.20.pkl`).
+    -   The scores reflect the likelihood that a compound inhibits the
+        growth of a given bacterial strain.
+-   **Training data:** 100000 randomly sampled compounds from ChemBERTa for the pretraining of MolE and data from *Maier et al., 2018* containing the influence of 1197 marketed drugs on the growth of 40 bacterial strains for the XGBoost classifier
+-   **Publication:** [Nature Communications (2025)](https://www.nature.com/articles/s41467-025-58804-4)
+
+
+### Output
+
+-   **Description:** For each compound, the model predicts growth inhibition scores for 40 different bacterial strains.
+-   **Output format:** table
+    -   **Shape:** `[n × 40, 2]`
+    -   **Columns**
+        -   `pred_id` *(str)*: Combination of the molecule name and
+            bacterial strain (e.g.,
+            `Halicin:Akkermansia muciniphila (NT5021)`).
+        -   `antimicrobial_predictive_probability` *(float)*: Predicted
+            probability that the compound inhibits microbial growth.
+-   **Example output file:**
+    `examples/output/example_molecules_prediction.tsv`
+
+
+## Installation
+
+Install the conda environment with all dependencies:
+
+``` bash
+# Create the conda environment called virtual-human-chc-mole
+conda env create -f environment.yaml
+
+# Activate the environment
+conda activate virtual-human-chc-mole
+```
+
+## Example
+
+### Minimal inference example
+
+``` python
+import torch
+import yaml
+import pickle
+import pandas as pd
+from huggingface_hub import hf_hub_download
+from mole_package import ginet_concat, mole_antimicrobial_prediction, mole_representation, dataset_representation
+
+class MolE:
+    def __init__(self, device='auto'):
+        repo = "virtual-human-chc/MolE"
+        self.device = "cuda:0" if device == "auto" and torch.cuda.is_available() else "cpu"
+
+        # Download + load
+        cfg = yaml.safe_load(open(hf_hub_download(repo, "config.yaml")))
+        self.model = ginet_concat.GINet(**cfg["model"]).to(self.device)
+        self.model.load_state_dict(torch.load(hf_hub_download(repo, "model.pth"), map_location=self.device))
+        self.xgb = pickle.load(open(hf_hub_download(repo, "MolE-XGBoost-08.03.2024_14.20.pkl"), "rb"))
+
+    def predict_from_smiles(self, smiles_tsv):
+        smiles_df = mole_representation.read_smiles(smiles_tsv, "smiles", "chem_name")
+        emb = dataset_representation.batch_representation(smiles_df, self.model, "smiles", "chem_name", device=self.device)
+        X_input = mole_antimicrobial_prediction.add_strains(
+            emb, "data/01.prepare_training_data/maier_screening_results.tsv.gz"
+        )
+        probs = self.xgb.predict_proba(X_input)[:, 1]
+        return pd.DataFrame(
+            {"antimicrobial_predictive_probability": probs},
+            index=X_input.index
+        )
+        
+# Run inference
+mole = MolE()
+pred = mole.predict_from_smiles("examples/input/examples_molecules.tsv")
+print(pred)
+```
+
+
+## References
+
+1.  Roberto Olayo Alarcon et al., *MolE: Graph-based molecular representation learning for antimicrobial discovery*, [Nature Communications (2025)](https://www.nature.com/articles/s41467-025-58804-4).
+2.  Maier et al., *Extensive antimicrobial screening of small molecules*, *Nature* (2018), <https://www.nature.com/articles/nature25979>.
+3.  GitHub repository: <https://github.com/rolayoalarcon/MolE>.
+4.  GitHub repository: <https://github.com/rolayoalarcon/mole_antimicrobial_potential>.
+5.  Model weights (Zenodo DOI): <https://doi.org/10.5281/zenodo.10803099>.
+
+
+## Copyright
+
+Code derived from <https://github.com/rolayoalarcon/MolE> is licensed
+under the **MIT License**, © 2024 Roberto Olayo Alarcon.
+Model weights are licensed under **Creative Commons Attribution 4.0
+International (CC BY 4.0)**, © 2024 Roberto Olayo Alarcon.
+Additional code © 2025 Maksim Pavlov, licensed under MIT.
+
+<!-- # MolE - Antimicrobial Prediction
 
 This model uses MolE's pre-trained representation to train XGBoost models to predict the antimicrobial activity of compounds based on their molecular structure. The model was developed by Roberto Olayo Alarcon et al. and more information can be found in the [GitHub repository](https://github.com/rolayoalarcon/MolE) and the [accompanying paper](https://www.nature.com/articles/s41467-025-58804-4).
 
@@ -92,4 +239,4 @@ molecule prioritization, reflecting the chance of the given molecule having grow
 
 ## Copyright
  
-Code derived from https://github.com/rolayoalarcon/MolE is licensed under the MIT license, Copyright (c) 2024 Roberto Olayo Alarcon. The [model weights](https://doi.org/10.5281/zenodo.10803099) are licensed under [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0/legalcode), Copyright (c) 2024 Roberto Olayo Alarcon. The other code is licensed under the MIT license, Copyright (c) 2025 Maksim Pavlov.
+Code derived from https://github.com/rolayoalarcon/MolE is licensed under the MIT license, Copyright (c) 2024 Roberto Olayo Alarcon. The [model weights](https://doi.org/10.5281/zenodo.10803099) are licensed under [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0/legalcode), Copyright (c) 2024 Roberto Olayo Alarcon. The other code is licensed under the MIT license, Copyright (c) 2025 Maksim Pavlov. -->