pavm595/MolE-antimicrobial

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 and the accompanying paper.

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.

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:

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 are licensed under Creative Commons Attribution 4.0 International, Copyright (c) 2024 Roberto Olayo Alarcon. The other code is licensed under the MIT license, Copyright (c) 2025 Maksim Pavlov.