BioXMol — molecular encoder

A graph neural network that turns a SMILES string into a fixed-length embedding. It was trained with soft contrastive learning to align molecular structure with biological readouts (Cell Painting morphology from JUMP-CP and L1000 transcriptomics), so the embeddings carry biological signal useful for tasks like drug-induced liver injury (DILI) prediction. This repo ships the molecule encoder only — you provide SMILES.

This is a custom GNN, not a transformers model, so AutoModel.from_pretrained does not work. Use the block below; it depends only on torch and rdkit.

How to use

Install:

pip install torch rdkit safetensors huggingface_hub numpy

Copy this whole block and run it — it downloads everything it needs from this repo:

import os, sys, torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

REPO = "arslanmasood/BioXMol"

# 1. download the code + weights from this repo
for f in ["modeling_bioxmol.py", "featurizer.py"]:
    code_dir = os.path.dirname(hf_hub_download(REPO, f))
sys.path.insert(0, code_dir)
from modeling_bioxmol import GatedGraphNeuralNetwork
from featurizer import smiles2graph

# 2. load the encoder
model = GatedGraphNeuralNetwork(n_edge=1, in_dim=75, n_conv=6, fc_dims=[1024, 128])
model.load_state_dict(load_file(hf_hub_download(REPO, "bioxmol_soft_seed0.safetensors")))
model.eval()

# 3. encode a list of SMILES -> embeddings
@torch.no_grad()
def encode(smiles_list, layer="first_fc"):
    graphs = [smiles2graph(s) for s in smiles_list]
    n = max(nf.shape[0] for _, nf in graphs)            # pad to the largest molecule
    A, F, M = [], [], []
    for adj, nf in graphs:
        adj = torch.as_tensor(adj, dtype=torch.float)
        nf = torch.as_tensor(nf, dtype=torch.float)
        a = torch.zeros(n, n);           a[:adj.shape[0], :adj.shape[1]] = adj
        f = torch.zeros(n, nf.shape[1]); f[:nf.shape[0]] = nf
        m = torch.zeros(n, 1);           m[:nf.shape[0]] = 1.0   # 1 = real atom
        A.append(a); F.append(f); M.append(m)
    return model.embed(torch.stack(A), torch.stack(F), torch.stack(M), layer=layer)

# example
emb = encode(["CCO", "CC(=O)Oc1ccccc1C(=O)O"])
print(emb.shape)   # torch.Size([2, 1024])

layer can be

• "GNN" (75-d),
• "first_fc" (1024-d, the default and recommended),
• "second_fc" (128-d).

Training data

Pretrained on JUMP-CP (Cell Painting) and LINCS L1000 (transcriptomics); evaluated on DILIRank 2.0. Each dataset has its own terms — check them before redistributing derived artifacts.

License

CC-BY-NC-4.0 (attribution, non-commercial). The training code inherits GPLv3 from the upstream framework and is distributed separately.

Citation

@misc{masood2026bioxmol,
  title  = {Unifying Disjoint Phenotypic Contexts: A Multimodal Soft Contrastive Approach to Identify DILI Activity Cliffs},
  author = {Masood, Muhammad Arslan and Cui, Tianyu and Heinonen, Markus and Kaski, Samuel},
  year   = {2026},
  doi    = {10.21203/rs.3.rs-9875956/v1}
}