bayesianflow-for-chem 1.2.0__py3-none-any.whl

bayesianflow_for_chem/train.py

@@ -0,0 +1,243 @@
+# -*- coding: utf-8 -*-
+# Author: Nianze A. Tao (Omozawa Sueno)
+"""
+Define ChemBFN and regressor models for training.
+"""
+from pathlib import Path
+from typing import Dict, Tuple, Union, Optional
+import torch
+import torch.optim as op
+import torch.nn.functional as F
+from loralib import lora_state_dict, mark_only_lora_as_trainable
+from torch import Tensor
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from lightning import LightningModule
+from .model import ChemBFN, MLP
+from .scorer import Scorer
+
+DEFAULT_MODEL_HPARAM = {"lr": 5e-5, "lr_warmup_step": 1000, "uncond_prob": 0.2}
+DEFAULT_REGRESSOR_HPARAM = {
+    "mode": "regression",
+    "lr_scheduler_factor": 0.8,
+    "lr_scheduler_patience": 20,
+    "lr_warmup_step": 1000,
+    "max_lr": 1e-4,
+    "freeze": False,
+}
+
+
+class Model(LightningModule):
+    def __init__(
+        self,
+        model: ChemBFN,
+        mlp: Optional[MLP] = None,
+        scorer: Optional[Scorer] = None,
+        hparam: Dict[str, Union[int, float]] = DEFAULT_MODEL_HPARAM,
+    ) -> None:
+        """
+        A `~lightning.LightningModule` wrapper of bayesian flow network for chemistry model.\n
+        This module is used in training stage only. By calling `Model(...).export_model(YOUR_WORK_DIR)` after training,
+        the model(s) will be saved to `YOUR_WORK_DIR/model.pt` and (if exists) `YOUR_WORK_DIR/mlp.pt`.
+
+        :param model: `~bayesianflow_for_chem.model.ChemBFN` instance.
+        :param mlp: `~bayesianflow_for_chem.model.MLP` instance or `None`.
+        :param scorer: `~bayesianflow_for_chem.scorer.Scorer` instance or `None`.
+        :param hparam: a `dict` instance of hyperparameters. See `bayesianflow_for_chem.train.DEFAULT_MODEL_HPARAM`.
+        :type model: bayesianflow_for_chem.model.ChemBFN
+        :type mlp: bayesianflow_for_chem.model.MLP | None
+        :type scorer: bayesianflow_for_chem.scorer.Scorer | None
+        :type hparam: dict
+        """
+        super().__init__()
+        self.model = model
+        self.mlp = mlp
+        self.scorer = scorer
+        self.save_hyperparameters(hparam, ignore=["model", "mlp", "scorer"])
+        if model.lora_enabled:
+            mark_only_lora_as_trainable(self.model)
+        self.use_scorer = self.scorer is not None
+
+    def training_step(self, batch: Dict[str, Tensor]) -> Tensor:
+        x = batch["token"]
+        t = torch.rand((x.shape[0], 1, 1), device=x.device)
+        if "mask" in batch:
+            mask = batch["mask"]
+        else:
+            mask = None
+        if self.mlp is not None:
+            y = batch["value"]
+            y = self.mlp.forward(y)
+            if y.dim() == 2:
+                y = y[:, None, :]
+            y_mask = F.dropout(torch.ones_like(t), self.hparams.uncond_prob, True, True)
+            y_mask = (y_mask != 0).float()
+            loss, p = self.model.cts_loss(x, t, y * y_mask, mask, self.use_scorer)
+        else:
+            loss, p = self.model.cts_loss(x, t, None, mask, self.use_scorer)
+        self.log("continuous_time_loss", loss.item())
+        if self.use_scorer:
+            scorer_loss = self.scorer.calc_score_loss(p)
+            self.log(f"{self.scorer.name}_loss", scorer_loss.item())
+            loss += scorer_loss * self.scorer.eta
+        return loss
+
+    def configure_optimizers(self) -> Dict[str, op.AdamW]:
+        optimizer = op.AdamW(self.parameters(), lr=1e-8, weight_decay=0.01)
+        return {"optimizer": optimizer}
+
+    def optimizer_step(self, *args, **kwargs) -> None:
+        optimizer: op.AdamW = kwargs["optimizer"] if "optimizer" in kwargs else args[2]
+        # warm-up step
+        if self.trainer.global_step < self.hparams.lr_warmup_step:
+            lr_scale = int(self.trainer.global_step + 1) / self.hparams.lr_warmup_step
+            lr_scale = min(1.0, lr_scale)
+            for pg in optimizer.param_groups:
+                pg["lr"] = lr_scale * self.hparams.lr
+        super().optimizer_step(*args, **kwargs)
+        optimizer.zero_grad(set_to_none=True)
+
+    def export_model(self, workdir: Path) -> None:
+        """
+        Save the trained model.
+
+        :param workdir: the directory to save the model(s)
+        :type workdir: pathlib.Path
+        :return:
+        :rtype: None
+        """
+        if self.model.lora_enabled:
+            torch.save(
+                {
+                    "lora_nn": lora_state_dict(self.model),
+                    "lora_param": self.model.lora_param,
+                },
+                workdir / "lora.pt",
+            )
+        else:
+            torch.save(
+                {"nn": self.model.state_dict(), "hparam": self.model.hparam},
+                workdir / "model.pt",
+            )
+        if self.mlp is not None:
+            torch.save(
+                {"nn": self.mlp.state_dict(), "hparam": self.mlp.hparam},
+                workdir / "mlp.pt",
+            )
+
+
+class Regressor(LightningModule):
+    def __init__(
+        self,
+        model: ChemBFN,
+        mlp: MLP,
+        hparam: Dict[str, Union[str, int, float, bool]] = DEFAULT_REGRESSOR_HPARAM,
+    ) -> None:
+        """
+        A `~lightning.LightningModule` wrapper of bayesian flow network for chemistry regression model.\n
+        This module is used in training stage only. By calling `Regressor(...).export_model(YOUR_WORK_DIR)` after training,
+        the models will be saved to `YOUR_WORK_DIR/model.pt` and `YOUR_WORK_DIR/readout.pt`.
+
+        :param model: `~bayesianflow_for_chem.model.ChemBFN` instance.
+        :param mlp: `~bayesianflow_for_chem.model.MLP` instance.
+        :param hparam: a `dict` instance of hyperparameters. See `bayesianflow_for_chem.train.DEFAULT_REGRESSOR_HPARAM`.
+        :type model: bayesianflow_for_chem.model.ChemBFN
+        :type mlp: bayesianflow_for_chem.model.MLP
+        :type hparam: dict
+        """
+        super().__init__()
+        self.model = model
+        self.mlp = mlp
+        self.model.requires_grad_(not hparam["freeze"])
+        self.save_hyperparameters(hparam, ignore=["model", "mlp"])
+        if model.lora_enabled:
+            mark_only_lora_as_trainable(self.model)
+        assert hparam["mode"] in ("regression", "classification")
+
+    @staticmethod
+    def _mask_label(label: Tensor) -> Tuple[Tensor, Tensor]:
+        # find the unlabelled position(s)
+        label_mask = (label != torch.inf).float()
+        # masked the unlabelled position(s)
+        masked_label = label.masked_fill(label == torch.inf, 0)
+        return label_mask, masked_label
+
+    def training_step(self, batch: Dict[str, Tensor]) -> Tensor:
+        x, y = batch["token"], batch["value"]
+        z = self.model.inference(x, self.mlp)
+        if self.hparams.mode == "classification":
+            n_b, n_y = y.shape
+            z = z.reshape(n_b * n_y, -1)
+            loss = F.cross_entropy(z, y.reshape(-1).to(torch.long))
+        else:
+            y_mask, y = self._mask_label(y)
+            loss = F.mse_loss(z * y_mask, y, reduction="mean")
+        self.log("train_loss", loss.item())
+        return loss
+
+    def validation_step(self, batch: Dict[str, Tensor]) -> None:
+        x, y = batch["token"], batch["value"]
+        z = self.model.inference(x, self.mlp)
+        if self.hparams.mode == "classification":
+            n_b, n_y = y.shape
+            z = z.reshape(n_b * n_y, -1)
+            val_loss = 1 - (torch.argmax(z, -1) == y.reshape(-1)).float().mean()
+        else:
+            y_mask, y = self._mask_label(y)
+            val_loss = (z * y_mask - y).abs().sum() / y_mask.sum()
+        self.log("val_loss", val_loss.item())
+
+    def configure_optimizers(self) -> Dict:
+        optimizer = op.AdamW(self.parameters(), lr=1e-7, weight_decay=0.01)
+        lr_scheduler_config = {
+            "scheduler": ReduceLROnPlateau(
+                optimizer,
+                "min",
+                factor=self.hparams.lr_scheduler_factor,
+                patience=self.hparams.lr_scheduler_patience,
+                min_lr=1e-6,
+            ),
+            "interval": "epoch",
+            "monitor": "val_loss",
+            "frequency": 1,
+            "strict": True,
+        }
+        return {"optimizer": optimizer, "lr_scheduler": lr_scheduler_config}
+
+    def optimizer_step(self, *args, **kwargs) -> None:
+        optimizer: op.AdamW = kwargs["optimizer"] if "optimizer" in kwargs else args[2]
+        # warm-up step
+        if self.trainer.global_step < self.hparams.lr_warmup_step:
+            lr_scale = int(self.trainer.global_step + 1) / self.hparams.lr_warmup_step
+            lr_scale = min(1.0, lr_scale)
+            for pg in optimizer.param_groups:
+                pg["lr"] = lr_scale * self.hparams.max_lr
+        super().optimizer_step(*args, **kwargs)
+        optimizer.zero_grad(set_to_none=True)
+
+    def export_model(self, workdir: Path) -> None:
+        """
+        Save the trained model.
+
+        :param workdir: the directory to save the model
+        :type workdir: pathlib.Path
+        :return:
+        :rtype: None
+        """
+        torch.save(
+            {"nn": self.mlp.state_dict(), "hparam": self.mlp.hparam},
+            workdir / "readout.pt",
+        )
+        if not self.hparams.freeze:
+            if self.model.lora_enabled:
+                torch.save(
+                    {
+                        "lora_nn": lora_state_dict(self.model),
+                        "lora_param": self.model.lora_param,
+                    },
+                    workdir / "lora.pt",
+                )
+            else:
+                torch.save(
+                    {"nn": self.model.state_dict(), "hparam": self.model.hparam},
+                    workdir / "model_ft.pt",
+                )

bayesianflow_for_chem/vocab.txt

@@ -0,0 +1,246 @@
+<pad>
+<start>
+<end>
+>
+>>
+[
+]
+(
+)
+.
+=
+#
+-
++
+\
+/
+:
+~
+@
+?
+*
+$
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+H
+He
+Li
+Be
+B
+C
+N
+O
+F
+Ne
+Na
+Mg
+Al
+Si
+P
+S
+Cl
+Ar
+K
+Ca
+Sc
+Ti
+V
+Cr
+Mn
+Fe
+Co
+Ni
+Cu
+Zn
+Ga
+Ge
+As
+Se
+Br
+Kr
+Rb
+Sr
+Y
+Zr
+Nb
+Mo
+Tc
+Ru
+Rh
+Pd
+Ag
+Cd
+In
+Sn
+Sb
+Te
+I
+Xe
+Cs
+Ba
+Hf
+Ta
+W
+Re
+Os
+Ir
+Pt
+Au
+Hg
+Tl
+Pb
+Bi
+Po
+At
+Rn
+Fr
+Ra
+Rf
+Db
+Sg
+Bh
+Hs
+Mt
+Ds
+Rg
+Cn
+Nh
+Fl
+Mc
+Lv
+Ts
+Og
+La
+Ce
+Pr
+Nd
+Pm
+Sm
+Eu
+Gd
+Tb
+Dy
+Ho
+Er
+Tm
+Yb
+Lu
+Ac
+Th
+Pa
+U
+Np
+Pu
+Am
+Cm
+Bk
+Cf
+Es
+Fm
+Md
+No
+Lr
+b
+c
+n
+o
+s
+p
+%10
+%11
+%12
+%13
+%14
+%15
+%16
+%17
+%18
+%19
+%20
+%21
+%22
+%23
+%24
+%25
+%26
+%27
+%28
+%29
+%30
+%31
+%32
+%33
+%34
+%35
+%36
+%37
+%38
+%39
+%40
+%41
+%42
+%43
+%44
+%45
+%46
+%47
+%48
+%49
+%50
+%51
+%52
+%53
+%54
+%55
+%56
+%57
+%58
+%59
+%60
+%61
+%62
+%63
+%64
+%65
+%66
+%67
+%68
+%69
+%70
+%71
+%72
+%73
+%74
+%75
+%76
+%77
+%78
+%79
+%80
+%81
+%82
+%83
+%84
+%85
+%86
+%87
+%88
+%89
+%90
+%91
+%92
+%93
+%94
+%95
+%96
+%97
+%98
+%99

bayesianflow_for_chem-1.2.0.dist-info/METADATA

@@ -0,0 +1,162 @@
+Metadata-Version: 2.2
+Name: bayesianflow_for_chem
+Version: 1.2.0
+Summary: Bayesian flow network framework for Chemistry
+Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
+Author: Nianze A. Tao
+Author-email: tao-nianze@hiroshima-u.ac.jp
+License: AGPL-3.0 licence
+Project-URL: Source, https://github.com/Augus1999/bayesian-flow-network-for-chemistry
+Keywords: Chemistry,CLM,ChemBFN
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: GNU Affero General Public License v3
+Classifier: Natural Language :: English
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: rdkit>=2023.9.6
+Requires-Dist: torch>=2.3.1
+Requires-Dist: numpy>=1.26.4
+Requires-Dist: loralib>=0.1.2
+Requires-Dist: lightning>=2.2.0
+Requires-Dist: scikit-learn>=1.5.0
+Requires-Dist: typing_extensions>=4.8.0
+Provides-Extra: geo2seq
+Requires-Dist: pynauty>=2.8.8.1; extra == "geo2seq"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: project-url
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# ChemBFN: Bayesian Flow Network for Chemistry
+
+[![DOI](https://zenodo.org/badge/DOI/10.1021/acs.jcim.4c01792.svg)](https://doi.org/10.1021/acs.jcim.4c01792)
+[![arxiv](https://img.shields.io/badge/arXiv-2412.11439-red)](https://arxiv.org/abs/2412.11439)
+
+This is the repository of the PyTorch implementation of ChemBFN model.
+
+## Features
+
+ChemBFN provides the state-of-the-art functionalities of
+* SMILES or SELFIES-based *de novo* molecule generation
+* Protein sequence *de novo* generation
+* Classifier-free guidance conditional generation (single or multi-objective optimisation)
+* Context-guided conditional generation (inpaint)
+* Outstanding out-of-distribution chemical space sampling
+* Fast sampling via ODE solver
+* Molecular property and activity prediction finetuning
+* Reaction yield prediction finetuning
+
+in an all-in-one-model style.
+
+## News
+
+* [21/01/2025] Our first paper has been accepted by [JCIM](https://pubs.acs.org/doi/10.1021/acs.jcim.4c01792).
+* [17/12/2024] The second paper of out-of-distribution generation is available on [arxiv.org](https://arxiv.org/abs/2412.11439).
+* [31/07/2024] Paper is available on [arxiv.org](https://arxiv.org/abs/2407.20294).
+* [21/07/2024] Paper was submitted to arXiv.
+
+## Install
+
+```bash
+$ pip install -U bayesianflow_for_chemistry
+```
+
+## Usage
+
+You can find example scripts in [📁example](./example) folder.
+
+## Pre-trained Model
+
+You can find pretrained models in [release](https://github.com/Augus1999/bayesian-flow-network-for-chemistry/releases) or on our [🤗Hugging Face model page](https://huggingface.co/suenoomozawa/ChemBFN).
+
+## Dataset Handling
+
+We provide a Python class [`CSVData`](./bayesianflow_for_chem/data.py) to handle data stored in CSV or similar format containing headers to identify the entities. The following is a quickstart.
+
+1. Download your dataset file (e.g., ESOL form [MoleculeNet](https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/delaney-processed.csv)) and split the file:
+```python
+>>> from bayesianflow_for_chem.tool import split_data
+
+>>> split_data("delaney-processed.csv", method="scaffold")
+```
+
+2. Load the split data:
+```python
+>>> from bayesianflow_for_chem.data import smiles2token, collate, CSVData
+
+>>> dataset = CSVData("delaney-processed_train.csv")
+>>> dataset[0]
+{'Compound ID': ['Thiophene'], 
+'ESOL predicted log solubility in mols per litre': ['-2.2319999999999998'], 
+'Minimum Degree': ['2'], 
+'Molecular Weight': ['84.14299999999999'], 
+'Number of H-Bond Donors': ['0'], 
+'Number of Rings': ['1'], 
+'Number of Rotatable Bonds': ['0'], 
+'Polar Surface Area': ['0.0'], 
+'measured log solubility in mols per litre': ['-1.33'], 
+'smiles': ['c1ccsc1']}
+```
+
+3. Create a mapping function to tokenise the dataset and select values:
+```python
+>>> import torch
+
+>>> def encode(x):
+...   smiles = x["smiles"][0]
+...   value = [float(i) for i in x["measured log solubility in mols per litre"]]
+...   return {"token": smiles2token(smiles), "value": torch.tensor(value)}
+
+>>> dataset.map(encode)
+>>> dataset[0]
+{'token': tensor([  1, 151,  23, 151, 151, 154, 151,  23,   2]), 
+'value': tensor([-1.3300])}
+```
+
+4. Wrap the dataset in <u>torch.utils.data.DataLoader</u>:
+```python
+>>> dataloader = torch.utils.data.DataLoader(dataset, 32, collate_fn=collate)
+```
+
+## Cite This Work
+
+```bibtex
+@misc{2024chembfn,
+      title={A Bayesian Flow Network Framework for Chemistry Tasks}, 
+      author={Nianze Tao and Minori Abe},
+      year={2024},
+      eprint={2407.20294},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2407.20294}, 
+}
+```
+Out-of-distribution generation:
+```bibtex
+@misc{2024chembfn_ood,
+      title={Bayesian Flow Is All You Need to Sample Out-of-Distribution Chemical Spaces}, 
+      author={Nianze Tao},
+      year={2024},
+      eprint={2412.11439},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2412.11439}, 
+}
+```

bayesianflow_for_chem-1.2.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+bayesianflow_for_chem/__init__.py,sha256=9XG7bc3Fup_O4QjFVcUF-_ZkVioDlvmK1GHs2Caxo50,293
+bayesianflow_for_chem/data.py,sha256=9tpRba40lxwrB6aPSJMkxUglEVC3VEQC9wWxhDuz3Q8,7760
+bayesianflow_for_chem/model.py,sha256=xD4ZpJSFMCdfh5k4--8Lpa34nn_sUvgklK6d9CPjz10,35434
+bayesianflow_for_chem/scorer.py,sha256=mV1vX8aBGFra2BE7N8WHihVIo3dXmUdPQIGfSaiuNdk,4084
+bayesianflow_for_chem/tool.py,sha256=a8AnH3geBpNpF6SFiTqKBrWlwDcQzii1zSQiUoyiLgY,17009
+bayesianflow_for_chem/train.py,sha256=kj6icGqymUUYopDtpre1oE_wpvpeNilbpzgffBsd1tk,9589
+bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
+bayesianflow_for_chem-1.2.0.dist-info/METADATA,sha256=WXd5mvyBbHX5Ba8uF6jkys8SCfdurkMRkxg-AedcBlc,5743
+bayesianflow_for_chem-1.2.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+bayesianflow_for_chem-1.2.0.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
+bayesianflow_for_chem-1.2.0.dist-info/RECORD,,

bayesianflow_for_chem-1.2.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

bayesianflow_for_chem-1.2.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+bayesianflow_for_chem