bayesianflow-for-chem 1.2.0__tar.gz

bayesianflow_for_chem-1.2.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,116 @@
+# 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/bayesianflow_for_chem/__init__.py

@@ -0,0 +1,11 @@
+# -*- coding: utf-8 -*-
+# Author: Nianze A. Tao (Omozawa Sueno)
+"""
+ChemBFN package.
+"""
+from . import data, tool, train, scorer
+from .model import ChemBFN, MLP
+
+__all__ = ["data", "tool", "train", "scorer", "ChemBFN", "MLP"]
+__version__ = "1.2.0"
+__author__ = "Nianze A. Tao (Omozawa Sueno)"

bayesianflow_for_chem-1.2.0/bayesianflow_for_chem/data.py

@@ -0,0 +1,250 @@
+# -*- coding: utf-8 -*-
+# Author: Nianze A. TAO (Omozawa SUENO)
+"""
+Tokenise SMILES/SAFE/SELFIES/GEO2SEQ/protein-sequence strings.
+"""
+import os
+import re
+from pathlib import Path
+from typing import Any, List, Dict, Union, Callable
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.utils.data import Dataset
+
+__filedir__ = Path(__file__).parent
+
+SMI_REGEX_PATTERN = (
+    r"(\[|\]|H[e,f,g,s,o]?|"
+    r"L[i,v,a,r,u]|"
+    r"B[e,r,a,i,h,k]?|"
+    r"C[l,a,r,o,u,d,s,n,e,m,f]?|"
+    r"N[e,a,i,b,h,d,o,p]?|"
+    r"O[s,g]?|S[i,c,e,r,n,m,b,g]?|"
+    r"K[r]?|T[i,c,e,a,l,b,h,m,s]|"
+    r"G[a,e,d]|R[b,u,h,e,n,a,f,g]|"
+    r"Yb?|Z[n,r]|P[t,o,d,r,a,u,b,m]?|"
+    r"F[e,r,l,m]?|M[g,n,o,t,c,d]|"
+    r"A[l,r,s,g,u,t,c,m]|I[n,r]?|"
+    r"W|X[e]|E[u,r,s]|U|D[b,s,y]|"
+    r"b|c|n|o|s|p|"
+    r"\(|\)|\.|=|#|-|\+|\\|\/|:|"
+    r"~|@|\?|>>?|\*|\$|\%[0-9]{2}|[0-9])"
+)
+SEL_REGEX_PATTERN = r"(\[[^\]]+]|\.)"
+GEO_REGEX_PATTERN = (
+    r"(H[e,f,g,s,o]?|"
+    r"L[i,v,a,r,u]|"
+    r"B[e,r,a,i,h,k]?|"
+    r"C[l,a,r,o,u,d,s,n,e,m,f]?|"
+    r"N[e,a,i,b,h,d,o,p]?|"
+    r"O[s,g]?|S[i,c,e,r,n,m,b,g]?|"
+    r"K[r]?|T[i,c,e,a,l,b,h,m,s]|"
+    r"G[a,e,d]|R[b,u,h,e,n,a,f,g]|"
+    r"Yb?|Z[n,r]|P[t,o,d,r,a,u,b,m]?|"
+    r"F[e,r,l,m]?|M[g,n,o,t,c,d]|"
+    r"A[l,r,s,g,u,t,c,m]|I[n,r]?|"
+    r"W|X[e]|E[u,r,s]|U|D[b,s,y]|"
+    r"-|.| |[0-9])"
+)
+AA_REGEX_PATTERN = r"(A|B|C|D|E|F|G|H|I|K|L|M|N|P|Q|R|S|T|V|W|Y|Z|-|.)"
+smi_regex = re.compile(SMI_REGEX_PATTERN)
+sel_regex = re.compile(SEL_REGEX_PATTERN)
+geo_regex = re.compile(GEO_REGEX_PATTERN)
+aa_regex = re.compile(AA_REGEX_PATTERN)
+
+
+def load_vocab(
+    vocab_file: Union[str, Path]
+) -> Dict[str, Union[int, List[str], Dict[str, int]]]:
+    """
+    Load vocabulary from source file.
+
+    :param vocab_file: file that contains vocabulary
+    :type vocab_file: str | pathlib.Path
+    :return: {"vocab_keys": vocab_keys, "vocab_count": vocab_count, "vocab_dict": vocab_dict}
+    :rtype: dict
+    """
+    with open(vocab_file, "r", encoding="utf-8") as f:
+        lines = f.read().strip()
+    vocab_keys = lines.split("\n")
+    vocab_count = len(vocab_keys)
+    vocab_dict = dict(zip(vocab_keys, range(vocab_count)))
+    return {
+        "vocab_keys": vocab_keys,
+        "vocab_count": vocab_count,
+        "vocab_dict": vocab_dict,
+    }
+
+
+_DEFUALT_VOCAB = load_vocab(__filedir__ / "vocab.txt")
+VOCAB_KEYS: List[str] = _DEFUALT_VOCAB["vocab_keys"]
+VOCAB_DICT: Dict[str, int] = _DEFUALT_VOCAB["vocab_dict"]
+VOCAB_COUNT: int = _DEFUALT_VOCAB["vocab_count"]
+AA_VOCAB_KEYS = (
+    VOCAB_KEYS[0:3] + "A B C D E F G H I K L M N P Q R S T V W Y Z - .".split()
+)
+AA_VOCAB_COUNT = len(AA_VOCAB_KEYS)
+AA_VOCAB_DICT = dict(zip(AA_VOCAB_KEYS, range(AA_VOCAB_COUNT)))
+GEO_VOCAB_KEYS = VOCAB_KEYS[0:3] + [" "] + VOCAB_KEYS[22:150] + [".", "-"]
+GEO_VOCAB_COUNT = len(GEO_VOCAB_KEYS)
+GEO_VOCAB_DICT = dict(zip(GEO_VOCAB_KEYS, range(GEO_VOCAB_COUNT)))
+
+
+def smiles2vec(smiles: str) -> List[int]:
+    """
+    SMILES tokenisation using a dataset-independent regex pattern.
+
+    :param smiles: SMILES string
+    :type smiles: str
+    :return: tokens w/o `<start>` and `<end>`
+    :rtype: list
+    """
+    tokens = [token for token in smi_regex.findall(smiles)]
+    return [VOCAB_DICT[token] for token in tokens]
+
+
+def geo2vec(geo2seq: str) -> List[int]:
+    """
+    Geo2Seq tokenisation using a dataset-independent regex pattern.
+
+    :param geo2seq: Geo2Seq string
+    :type geo2seq: str
+    :return: tokens w/o `<start>` and `<end>`
+    :rtype: list
+    """
+    tokens = [token for token in geo_regex.findall(geo2seq)]
+    return [GEO_VOCAB_DICT[token] for token in tokens]
+
+
+def aa2vec(aa_seq: str) -> List[int]:
+    """
+    Protein sequence tokenisation using a dataset-independent regex pattern.
+
+    :param aa_seq: protein (amino acid) sequence
+    :type aa_seq: str
+    :return: tokens w/o `<start>` and `<end>`
+    :rtype: list
+    """
+    tokens = [token for token in aa_regex.findall(aa_seq)]
+    return [AA_VOCAB_DICT[token] for token in tokens]
+
+
+def split_selfies(selfies: str) -> List[str]:
+    """
+    SELFIES tokenisation.
+
+    :param selfies: SELFIES string
+    :type selfies: str
+    :return: SELFIES vocab
+    :rtype: list
+    """
+    return [token for token in sel_regex.findall(selfies)]
+
+
+def smiles2token(smiles: str) -> Tensor:
+    # start token: <start> = 1; end token: <esc> = 2
+    return torch.tensor([1] + smiles2vec(smiles) + [2], dtype=torch.long)
+
+
+def geo2token(geo2seq: str) -> Tensor:
+    # start token: <start> = 1; end token: <esc> = 2
+    return torch.tensor([1] + geo2vec(geo2seq) + [2], dtype=torch.long)
+
+
+def aa2token(aa_seq: str) -> Tensor:
+    # start token: <start> = 1; end token: <end> = 2
+    return torch.tensor([1] + aa2vec(aa_seq) + [2], dtype=torch.long)
+
+
+def collate(batch: List[Dict[str, Tensor]]) -> Dict[str, Tensor]:
+    """
+    Padding the data in one batch into the same size.\n
+    Should be passed to `~torch.utils.data.DataLoader` as `DataLoader(collate_fn=collate, ...)`.
+
+    :param batch: a list of data (one batch)
+    :type batch: list
+    :return: batched {"token": token} or {"token": token, "value": value}
+    :rtype: dict
+    """
+    token = [i["token"] for i in batch]
+    if "MAX_PADDING_LENGTH" in os.environ:
+        lmax = int(os.environ["MAX_PADDING_LENGTH"])
+    else:
+        lmax = max([len(w) for w in token])
+    token = torch.cat(
+        [F.pad(i, (0, lmax - len(i)), value=0)[None, :] for i in token], 0
+    )
+    out_dict = {"token": token}
+    if "value" in batch[0]:
+        out_dict["value"] = torch.cat([i["value"][None, :] for i in batch], 0)
+    if "mask" in batch[0]:
+        mask = [i["mask"] for i in batch]
+        out_dict["mask"] = torch.cat(
+            [F.pad(i, (0, lmax - len(i)), value=0)[None, :] for i in mask], 0
+        )
+    return out_dict
+
+
+class CSVData(Dataset):
+    def __init__(self, file: Union[str, Path]):
+        """
+        Define dataset stored in CSV file.
+
+        :param file: dataset file name <file>
+        :type file: str | pathlib.Path
+        """
+        super().__init__()
+        with open(file, "r") as db:
+            self.data = db.readlines()
+        self.header_idx_dict: Dict[str, List[int]] = {}
+        for key, i in enumerate(self.data[0].replace("\n", "").split(",")):
+            if i in self.header_idx_dict:
+                self.header_idx_dict[i].append(key)
+            else:
+                self.header_idx_dict[i] = [key]
+        self.mapping = lambda x: x
+
+    def __len__(self) -> int:
+        return len(self.data) - 1
+
+    def __getitem__(self, idx: Union[int, Tensor]) -> Dict[str, Tensor]:
+        if torch.is_tensor(idx):
+            idx = idx.tolist()
+        # valid `idx` should start from 1 instead of 0
+        data: List[str] = self.data[idx + 1].replace("\n", "").split(",")
+        data_dict: Dict[str, List[str]] = {}
+        for key in self.header_idx_dict:
+            data_dict[key] = [data[i] for i in self.header_idx_dict[key]]
+        return self.mapping(data_dict)
+
+    def map(self, mapping: Callable[[Dict[str, List[str]]], Any]) -> None:
+        """
+        Pass a customised mapping function to transform the data entities to tensors.
+
+        e.g.
+        ```python
+        import torch
+        from bayesianflow_for_chem.data import smiles2token, CSVData
+
+
+        def encode(x):
+            return {
+                "token": smiles2token(".".join(x["smiles"])),
+                "value": torch.tensor([float(i) if i != "" else torch.inf for i in x["value"]]),
+            }
+
+        dataset = CSVData(...)
+        dataset.map(encode)
+        ```
+
+        :param mapping: customised mapping function
+        :type mapping: callable
+        :return:
+        :rtype: None
+        """
+        self.mapping = mapping
+
+
+if __name__ == "__main__":
+    ...