Update README.md

README.md

@@ -1,22 +1,86 @@
-# 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.
-
-## Files:
-
-- `model.pth` - the pre-trained representation model's weights
-- `config.yaml` - model configuration
-- `MolE-XGBoost-08.03.2024_14.20.pkl` - pretrained XGBoost model
-
-## Usage
-
-Not ready yet.
-
-## Publication
-For more information about MolE, and how we use it to predict antimicrobial activity, you can check out the paper in Nature Communications:  
-[**Pre-trained molecular representations enable antimicrobial discovery**](https://www.nature.com/articles/s41467-025-58804-4) 
-
-## GitHub
-
-The code is available here:
-[**Link to GitHub repo**](https://github.com/rolayoalarcon/mole_antimicrobial_potential)
+---
+tags:
+- pytorch
+- pyg
+- graph-neural-networks
+- machine-learning
+- barlow-twins
+- graph-isomorphism-network
+- molecular-biology
+- computational-biology
+- antibiotics
+- antimicrobial-discovery
+- high-throughput-screening
+- virtual-drug-screening
+- haicu
+library_name: pytorch
+language:
+- en
+---
+
+# 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).
+
+## Files:
+
+- `model.pth` - the pre-trained representation model's weights
+- `config.yaml` - model configuration
+- `MolE-XGBoost-08.03.2024_14.20.pkl` - pretrained XGBoost model
+
+## Usage
+
+### Inference Example
+
+Below is a minimal example showing how to load and run inference with **MolE** directly from the Hugging Face Hub.
+
+```python
+import torch, yaml, pickle, pandas as pd
+from huggingface_hub import hf_hub_download
+import mole_representation, mole_antimicrobial_prediction
+
+class MolE:
+    def __init__(self, device='auto'):
+        repo = "pavm595/MolE-antimicrobial"
+        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 = mole_representation.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 = mole_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:
+
+```python
+mole = MolE()
+pred = mole.predict_from_smiles("examples/input/examples_molecules.tsv")
+print(pred)
+```
+
+## Metadata
+
+### Input
+
+The input is a TSV file with two columns: `chem_name` and `smiles`. The column 'chem_name' contains the name of the molecule from PubChem, e.g. Halicin, and the column 'smiles' contains the chemical formula in SMILES format, e.g. `C1=C(SC(=N1)SC2=NN=C(S2)N)[N+](=O)[O-]`. An example input is the file `examples\input\example_molecules.tsv`.
+
+### Output
+
+The output is a TSV file with two columns: `pred_id` and `antimicrobial_predictive_probability`. The column `pred_id` contains a given molecule and a bacteria, e.g. Halicin:Akkermansia muciniphila (NT5021), and the column `antimicrobial_predictive_probability` contains antimicrobial potential (AP) scores for
+molecule prioritization, reflecting the chance of the given molecule having growth inhibition effect on the corresponding bacteria, e.g. 0.021192694. An example output is `examples/output/example_molecules_prediction.tsv`.
+
+## 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.