bayesianflow-for-chem 1.2.6__tar.gz → 1.3.0__tar.gz

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/PKG-INFO

@@ -1,16 +1,15 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 1.2.6
+Version: 1.3.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
+License: AGPL-3.0-or-later
 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
@@ -39,6 +38,7 @@ Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: keywords
 Dynamic: license
+Dynamic: license-file
 Dynamic: project-url
 Dynamic: provides-extra
 Dynamic: requires-dist
@@ -86,7 +86,7 @@ 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).
+You can find pretrained models on our [🤗Hugging Face model page](https://huggingface.co/suenoomozawa/ChemBFN).
 
 ## Dataset Handling
 

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/README.md

@@ -39,7 +39,7 @@ 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).
+You can find pretrained models on our [🤗Hugging Face model page](https://huggingface.co/suenoomozawa/ChemBFN).
 
 ## Dataset Handling
 

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/bayesianflow_for_chem/__init__.py

@@ -4,8 +4,8 @@
 ChemBFN package.
 """
 from . import data, tool, train, scorer
-from .model import ChemBFN, MLP
+from .model import ChemBFN, MLP, EnsembleChemBFN
 
-__all__ = ["data", "tool", "train", "scorer", "ChemBFN", "MLP"]
-__version__ = "1.2.6"
+__all__ = ["data", "tool", "train", "scorer", "ChemBFN", "MLP", "EnsembleChemBFN"]
+__version__ = "1.3.0"
 __author__ = "Nianze A. Tao (Omozawa Sueno)"

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/bayesianflow_for_chem/data.py

@@ -55,7 +55,7 @@ aa_regex = re.compile(AA_REGEX_PATTERN)
 
 
 def load_vocab(
-    vocab_file: Union[str, Path]
+    vocab_file: Union[str, Path],
 ) -> Dict[str, Union[int, List[str], Dict[str, int]]]:
     """
     Load vocabulary from source file.
@@ -108,7 +108,7 @@ def geo2vec(geo2seq: str) -> List[int]:
     """
     Geo2Seq tokenisation using a dataset-independent regex pattern.
 
-    :param geo2seq: Geo2Seq string
+    :param geo2seq: `GEO2SEQ` string
     :type geo2seq: str
     :return: tokens w/o `<start>` and `<end>`
     :rtype: list

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/bayesianflow_for_chem/model.py

@@ -4,7 +4,8 @@
 Define Bayesian Flow Network for Chemistry (ChemBFN) model.
 """
 from pathlib import Path
-from typing import List, Tuple, Optional, Union
+from copy import deepcopy
+from typing import List, Tuple, Dict, Optional, Union, Callable
 import torch
 import torch.nn as nn
 from torch import Tensor
@@ -592,6 +593,13 @@ class ChemBFN(nn.Module):
         x, logits = torch.broadcast_tensors(x[..., None], logits)
         return (-logits.gather(-1, x[..., :1]).squeeze(-1)).mean()
 
+    @staticmethod
+    def reshape_y(y: Tensor) -> Tensor:
+        assert y.dim() <= 3  # this doesn't work if the model is frezen in JIT.
+        if y.dim() == 2:
+            return y[:, None, :]
+        return y
+
     @torch.jit.export
     def sample(
         self,
@@ -607,7 +615,7 @@ class ChemBFN(nn.Module):
 
         :param batch_size: batch size
         :param sequence_size: max sequence length
-        :param y: conditioning vector;      shape: (n_b, 1, n_f)
+        :param y: conditioning vector;      shape: (n_b, 1, n_f) or (n_b, n_f)
         :param sample_step: number of sampling steps
         :param guidance_strength: strength of conditional generation. It is not used if y is null.
         :param token_mask: token mask;      shape: (1, 1, n_vocab)
@@ -626,9 +634,7 @@ class ChemBFN(nn.Module):
             / self.K
         )
         if y is not None:
-            assert y.dim() == 3  # this doesn't work if the model is frezen in JIT.
-            if y.shape[0] == 1:
-                y = y.repeat(batch_size, 1, 1)
+            y = self.reshape_y(y)
         for i in torch.linspace(1, sample_step, sample_step, device=self.beta.device):
             t = (i - 1).view(1, 1, 1).repeat(batch_size, 1, 1) / sample_step
             p = self.discrete_output_distribution(theta, t, y, guidance_strength)
@@ -663,7 +669,7 @@ class ChemBFN(nn.Module):
 
         :param batch_size: batch size
         :param sequence_size: max sequence length
-        :param y: conditioning vector;      shape: (n_b, 1, n_f)
+        :param y: conditioning vector;      shape: (n_b, 1, n_f) or (n_b, n_f)
         :param sample_step: number of sampling steps
         :param guidance_strength: strength of conditional generation. It is not used if y is null.
         :param token_mask: token mask;      shape: (1, 1, n_vocab)
@@ -681,9 +687,7 @@ class ChemBFN(nn.Module):
         """
         z = torch.zeros((batch_size, sequence_size, self.K), device=self.beta.device)
         if y is not None:
-            assert y.dim() == 3  # this doesn't work if the model is frezen in JIT.
-            if y.shape[0] == 1:
-                y = y.repeat(batch_size, 1, 1)
+            y = self.reshape_y(y)
         for i in torch.linspace(1, sample_step, sample_step, device=self.beta.device):
             t = (i - 1).view(1, 1, 1).repeat(batch_size, 1, 1) / sample_step
             theta = torch.softmax(z, -1)
@@ -714,7 +718,7 @@ class ChemBFN(nn.Module):
         Molecule inpaint functionality.
 
         :param x: categorical indices of scaffold;  shape: (n_b, n_t)
-        :param y: conditioning vector;              shape: (n_b, 1, n_f)
+        :param y: conditioning vector;              shape: (n_b, 1, n_f) or (n_b, n_f)
         :param sample_step: number of sampling steps
         :param guidance_strength: strength of conditional generation. It is not used if y is null.
         :param token_mask: token mask;              shape: (1, 1, n_vocab)
@@ -733,9 +737,7 @@ class ChemBFN(nn.Module):
         x_onehot = nn.functional.one_hot(x, self.K) * mask
         theta = x_onehot + (1 - mask) * theta
         if y is not None:
-            assert y.dim() == 3  # this doesn't work if the model is frezen in JIT.
-            if y.shape[0] == 1:
-                y = y.repeat(n_b, 1, 1)
+            y = self.reshape_y(y)
         for i in torch.linspace(1, sample_step, sample_step, device=x.device):
             t = (i - 1).view(1, 1, 1).repeat(n_b, 1, 1) / sample_step
             p = self.discrete_output_distribution(theta, t, y, guidance_strength)
@@ -769,7 +771,7 @@ class ChemBFN(nn.Module):
         ODE inpainting.
 
         :param x: categorical indices of scaffold;  shape: (n_b, n_t)
-        :param y: conditioning vector;              shape: (n_b, 1, n_f)
+        :param y: conditioning vector;              shape: (n_b, 1, n_f) or (n_b, n_f)
         :param sample_step: number of sampling steps
         :param guidance_strength: strength of conditional generation. It is not used if y is null.
         :param token_mask: token mask;              shape: (1, 1, n_vocab)
@@ -789,9 +791,7 @@ class ChemBFN(nn.Module):
         x_onehot = nn.functional.one_hot(x, self.K) * mask
         z = torch.zeros((n_b, n_t, self.K), device=self.beta.device)
         if y is not None:
-            assert y.dim() == 3  # this doesn't work if the model is frezen in JIT.
-            if y.shape[0] == 1:
-                y = y.repeat(n_b, 1, 1)
+            y = self.reshape_y(y)
         for i in torch.linspace(1, sample_step, sample_step, device=self.beta.device):
             t = (i - 1).view(1, 1, 1).repeat(n_b, 1, 1) / sample_step
             theta = torch.softmax(z, -1)
@@ -847,13 +847,7 @@ class ChemBFN(nn.Module):
         with open(ckpt, "rb") as f:
             state = torch.load(f, "cpu", weights_only=True)
         nn, hparam = state["nn"], state["hparam"]
-        model = cls(
-            hparam["num_vocab"],
-            hparam["channel"],
-            hparam["num_layer"],
-            hparam["num_head"],
-            hparam["dropout"],
-        )
+        model = cls(**hparam)
         model.load_state_dict(nn, False)
         if ckpt_lora:
             with open(ckpt_lora, "rb") as g:
@@ -908,7 +902,7 @@ class MLP(nn.Module):
         if self.class_input:
             x = x.to(dtype=torch.long)
         for layer in self.layers[:-1]:
-            x = torch.selu(layer(x))
+            x = torch.selu(layer.forward(x))
         return self.layers[-1](x)
 
     @classmethod
@@ -926,10 +920,382 @@ class MLP(nn.Module):
         with open(ckpt, "rb") as f:
             state = torch.load(f, "cpu", weights_only=True)
         nn, hparam = state["nn"], state["hparam"]
-        model = cls(hparam["size"], hparam["class_input"], hparam["dropout"])
+        model = cls(**hparam)
         model.load_state_dict(nn, strict)
         return model
 
 
+class EnsembleChemBFN(ChemBFN):
+    """
+    This module does not fully support `torch.jit.script`. We have `EnsembleChemBFN.jit()`
+    method to JIT compile the submodels.
+    `torch.compile()` is a better choice to compiling the whole model.
+    """
+
+    def __init__(
+        self,
+        base_model_path: Union[str, Path],
+        lora_paths: Union[List[Union[str, Path]], Dict[str, Union[str, Path]]],
+        cond_heads: Union[List[nn.Module], Dict[str, nn.Module]],
+        adapter_weights: Optional[Union[List[float], Dict[str, float]]] = None,
+        semi_autoregressive_flags: Optional[Union[List[bool], Dict[str, bool]]] = None,
+    ) -> None:
+        """
+        Ensemble of ChemBFN models from LoRA checkpoints.
+
+        :param base_model_path: base model checkpoint file
+        :param lora_paths: a list of LoRA checkpoint files or a `dict` instance of these files
+        :param cond_heads: a list of conditioning network heads or a `dict` instance of these networks
+        :param adapter_weights: a list of weights of each LoRA finetuned model or a 'dict` instance of these weights; default is equally weighted
+        :param semi_autoregressive_flags: a list of the semi-autoregressive behaviour states of each LoRA finetuned model or a `dict` instance of these states; default is all `False`
+        :type base_model_path: str | pathlib.Path
+        :type lora_paths: list | dict
+        :type cond_heads: list | dict
+        :type adapter_weights: list | dict | None
+        :type semi_autoregressive_flags: list | dict | None
+        """
+        n = len(lora_paths)
+        assert type(lora_paths) == type(
+            cond_heads
+        ), "`lora_paths` and `cond_heads` should have the same type!"
+        assert n == len(
+            cond_heads
+        ), "`lora_paths` and `cond_heads` should have the same length!"
+        if adapter_weights:
+            assert type(lora_paths) == type(
+                adapter_weights
+            ), "`lora_paths` and `adapter_weights` should have the same type!"
+            assert n == len(
+                adapter_weights
+            ), "`lora_paths` and `adapter_weights` should have the same length!"
+        if semi_autoregressive_flags:
+            assert type(lora_paths) == type(
+                semi_autoregressive_flags
+            ), "`lora_paths` and `semi_autoregressive_flags` should have the same type!"
+            assert n == len(
+                semi_autoregressive_flags
+            ), "`lora_paths` and `semi_autoregressive_flags` should have the same length!"
+        _label_is_dict = isinstance(lora_paths, dict)
+        if isinstance(lora_paths, list):
+            names = tuple(f"val_{i}" for i in range(n))
+            lora_paths = dict(zip(names, lora_paths))
+            cond_heads = dict(zip(names, cond_heads))
+            if not adapter_weights:
+                adapter_weights = (1 / n for _ in names)
+            if not semi_autoregressive_flags:
+                semi_autoregressive_flags = (False for _ in names)
+            adapter_weights = dict(zip(names, adapter_weights))
+            semi_autoregressive_flags = dict(zip(names, semi_autoregressive_flags))
+        else:
+            names = tuple(lora_paths.keys())
+            if not adapter_weights:
+                adapter_weights = dict(zip(names, (1 / n for _ in names)))
+            if not semi_autoregressive_flags:
+                semi_autoregressive_flags = dict(zip(names, (False for _ in names)))
+        base_model = ChemBFN.from_checkpoint(base_model_path)
+        models = dict(zip(names, (deepcopy(base_model.eval()) for _ in names)))
+        for k in names:
+            with open(lora_paths[k], "rb") as f:
+                state = torch.load(f, "cpu", weights_only=True)
+            lora_nn, lora_param = state["lora_nn"], state["lora_param"]
+            models[k].enable_lora(**lora_param)
+            models[k].load_state_dict(lora_nn, False)
+            models[k].semi_autoregressive = semi_autoregressive_flags[k]
+        super().__init__(**base_model.hparam)
+        self.cond_heads = nn.ModuleDict(cond_heads)
+        self.models = nn.ModuleDict(models)
+        self.adapter_weights = adapter_weights
+        self._label_is_dict = _label_is_dict  # flag
+        # ------- remove unnecessary submodules -------
+        self.embedding = None
+        self.time_embed = None
+        self.position = None
+        self.encoder_layers = None
+        self.final_layer = None
+        self.__delattr__("embedding")
+        self.__delattr__("time_embed")
+        self.__delattr__("position")
+        self.__delattr__("encoder_layers")
+        self.__delattr__("final_layer")
+        # ------- remove unused attributes -------
+        self.__delattr__("semi_autoregressive")
+        self.__delattr__("lora_enabled")
+        self.__delattr__("lora_param")
+        self.__delattr__("hparam")
+
+    def construct_y(
+        self, c: Union[List[Tensor], Dict[str, Tensor]]
+    ) -> Dict[str, Tensor]:
+        assert (
+            isinstance(c, dict) is self._label_is_dict
+        ), f"`c` should be a {'`dict` instance' if self._label_is_dict else '`list` instance'} but got {type(c)} instand."
+        out: Dict[str, Tensor] = {}
+        if isinstance(c, list):
+            c = dict(zip([f"val_{i}" for i in range(len(c))], c))
+        for name, model in self.cond_heads.items():
+            y = model.forward(c[name])
+            if y.dim() == 2:
+                y = y[:, None, :]
+            out[name] = y
+        return out
+
+    def discrete_output_distribution(
+        self, theta: Tensor, t: Tensor, y: Dict[str, Tensor], w: float
+    ) -> Tensor:
+        """
+        :param theta: input distribution;          shape: (n_b, n_t, n_vocab)
+        :param t: continuous time in [0, 1];       shape: (n_b, 1, 1)
+        :param y: a dict of conditioning vectors;  shape: (n_b, 1, n_f) * n_h
+        :param w: guidance strength controlling the conditional generation
+        :type theta: torch.Tensor
+        :type t: torch.Tensor
+        :type y: dict
+        :type w: float
+        :return: output distribution;              shape: (n_b, n_t, n_vocab)
+        :rtype: torch.Tensor
+        """
+        theta = 2 * theta - 1  # rescale to [-1, 1]
+        p_uncond, p_cond = torch.zeros_like(theta), torch.zeros_like(theta)
+        # Q: Why not use `torch.vmap`? It's faster than doing the loop, isn't it?
+        #
+        # A: We have quite a few reasons to avoid using `vmap`:
+        #    1. JIT doesn't support vmap;
+        #    2. It's harder to switch on/off semi-autroregssive behaviours for individual
+        #       models when all models are stacked into one (we have a solution but it's not
+        #       that elegant);
+        #    3. We just found that the result from vmap was not identical to doing the loop;
+        #    4. vmap requires all models have the same size but it's not always that case
+        #       since we sometimes use different ranks of LoRA in finetuning.
+        for name, model in self.models.items():
+            p_uncond_ = model.forward(theta, t, None, None)
+            p_uncond += p_uncond_ * self.adapter_weights[name]
+            p_cond_ = model.forward(theta, t, None, y[name])
+            p_cond += p_cond_ * self.adapter_weights[name]
+        return softmax((1 + w) * p_cond - w * p_uncond, -1)
+
+    @staticmethod
+    def reshape_y(y: Dict[str, Tensor]) -> Dict[str, Tensor]:
+        for k in y:
+            assert y[k].dim() <= 3
+            if y[k].dim() == 2:
+                y[k] = y[k][:, None, :]
+        return y
+
+    @torch.inference_mode()
+    def sample(
+        self,
+        batch_size: int,
+        sequence_size: int,
+        conditions: Union[List[Tensor], Dict[str, Tensor]],
+        sample_step: int = 100,
+        guidance_strength: float = 4.0,
+        token_mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Sample from a piror distribution.
+
+        :param batch_size: batch size
+        :param sequence_size: max sequence length
+        :param conditions: guidance conditions;  shape: (n_b, n_c) * n_h
+        :param sample_step: number of sampling steps
+        :param guidance_strength: strength of conditional generation. It is not used if y is null.
+        :param token_mask: token mask;           shape: (1, 1, n_vocab)
+        :type batch_size: int
+        :type sequence_size: int
+        :type conditions: list | dict
+        :type sample_step: int
+        :type guidance_strength: float
+        :type token_mask: torch.Tensor | None
+        :return: sampled token indices;          shape: (n_b, n_t) \n
+                 entropy of the tokens;          shape: (n_b)
+        :rtype: tuple
+        """
+        y = self.construct_y(conditions)
+        return super().sample(
+            batch_size, sequence_size, y, sample_step, guidance_strength, token_mask
+        )
+
+    @torch.inference_mode()
+    def ode_sample(
+        self,
+        batch_size: int,
+        sequence_size: int,
+        conditions: Union[List[Tensor], Dict[str, Tensor]],
+        sample_step: int = 100,
+        guidance_strength: float = 4.0,
+        token_mask: Optional[Tensor] = None,
+        temperature: float = 0.5,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        ODE-based sampling.
+
+        :param batch_size: batch size
+        :param sequence_size: max sequence length
+        :param conditions: conditioning vector;  shape: (n_b, n_c) * n_h
+        :param sample_step: number of sampling steps
+        :param guidance_strength: strength of conditional generation. It is not used if y is null.
+        :param token_mask: token mask;           shape: (1, 1, n_vocab)
+        :param temperature: sampling temperature
+        :type batch_size: int
+        :type sequence_size: int
+        :type conditions: list | dict
+        :type sample_step: int
+        :type guidance_strength: float
+        :type token_mask: torch.Tensor | None
+        :type temperature: float
+        :return: sampled token indices;          shape: (n_b, n_t) \n
+                 entropy of the tokens;          shape: (n_b)
+        :rtype: tuple
+        """
+        y = self.construct_y(conditions)
+        return super().ode_sample(
+            batch_size,
+            sequence_size,
+            y,
+            sample_step,
+            guidance_strength,
+            token_mask,
+            temperature,
+        )
+
+    @torch.inference_mode()
+    def inpaint(
+        self,
+        x: Tensor,
+        conditions: Union[List[Tensor], Dict[str, Tensor]],
+        sample_step: int = 100,
+        guidance_strength: float = 4.0,
+        token_mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Molecule inpaint functionality.
+
+        :param x: categorical indices of scaffold;  shape: (n_b, n_t)
+        :param conditions: conditioning vector;     shape: (n_b, n_c) * n_h
+        :param sample_step: number of sampling steps
+        :param guidance_strength: strength of conditional generation. It is not used if y is null.
+        :param token_mask: token mask;              shape: (1, 1, n_vocab)
+        :type x: torch.Tensor
+        :type conditions: list | dict
+        :type sample_step: int
+        :type guidance_strength: float
+        :type token_mask: torch.Tensor | None
+        :return: sampled token indices;             shape: (n_b, n_t) \n
+                 entropy of the tokens;             shape: (n_b)
+        :rtype: tuple
+        """
+        y = self.construct_y(conditions)
+        return super().inpaint(x, y, sample_step, guidance_strength, token_mask)
+
+    @torch.inference_mode()
+    def ode_inpaint(
+        self,
+        x: Tensor,
+        conditions: Union[List[Tensor], Dict[str, Tensor]],
+        sample_step: int = 100,
+        guidance_strength: float = 4.0,
+        token_mask: Optional[Tensor] = None,
+        temperature: float = 0.5,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        ODE inpainting.
+
+        :param x: categorical indices of scaffold;  shape: (n_b, n_t)
+        :param conditions: conditioning vector;     shape: (n_b, n_c) * n_h
+        :param sample_step: number of sampling steps
+        :param guidance_strength: strength of conditional generation. It is not used if y is null.
+        :param token_mask: token mask;              shape: (1, 1, n_vocab)
+        :param temperature: sampling temperature
+        :type x: torch.Tensor
+        :type conditions: list | dict
+        :type sample_step: int
+        :type guidance_strength: float
+        :type token_mask: torch.Tensor | None
+        :type temperature: float
+        :return: sampled token indices;             shape: (n_b, n_t) \n
+                 entropy of the tokens;             shape: (n_b)
+        :rtype: tuple
+        """
+        y = self.construct_y(conditions)
+        return super().ode_inpaint(
+            x, y, sample_step, guidance_strength, token_mask, temperature
+        )
+
+    def quantise(
+        self, quantise_method: Optional[Callable[[ChemBFN], nn.Module]] = None
+    ) -> None:
+        """
+        Quantise the submodels. \n
+        This method should be called, if necessary, before `torch.compile()`.
+
+        :param quantise_method: quantisation method; default is `bayesianflow_for_chem.tool.quantise_model`
+        :type quantise_method: callable | None
+        :return:
+        :rtype: None
+        """
+        if quantise_method is None:
+            from bayesianflow_for_chem.tool import quantise_model
+
+            quantise_method = quantise_model
+        for k, v in self.models.items():
+            self.models[k] = quantise_method(v)
+
+    def jit(self, freeze: bool = False) -> None:
+        """
+        JIT compile the submodels. \n
+        This method should be called, if necessary, before `quantise()` method is called if applied.
+
+        :param freeze: whether to freeze the submodels; default is `False`. If set to `True` this
+                       method should be called before moving the model to a different device.
+        :type freeze: bool
+        :return:
+        :rtype: None
+        """
+        for k, v in self.models.items():
+            self.models[k] = torch.jit.script(v)
+            if freeze:
+                self.models[k] = torch.jit.freeze(
+                    self.models[k], ["semi_autoregressive"]
+                )
+
+    @torch.jit.ignore
+    def forward(self, *_, **__) -> None:
+        """
+        Don't use this method!
+        """
+        raise NotImplementedError("There's nothing here!")
+
+    def cts_loss(self, *_, **__) -> None:
+        """
+        Don't use this method!
+        """
+        raise NotImplementedError("There's nothing here!")
+
+    def reconstruction_loss(self, *_, **__) -> None:
+        """
+        Don't use this method!
+        """
+        raise NotImplementedError("There's nothing here!")
+
+    def enable_lora(self, *_, **__) -> None:
+        """
+        Don't use this method!
+        """
+        raise NotImplementedError("There's nothing here!")
+
+    def inference(self, *_, **__) -> None:
+        """
+        Don't use this method!
+        """
+        raise NotImplementedError("There's nothing here!")
+
+    @classmethod
+    def from_checkpoint(cls, *_, **__) -> None:
+        """
+        Don't use this method!
+        """
+        raise NotImplementedError("There's nothing here!")
+
+
 if __name__ == "__main__":
     ...

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/bayesianflow_for_chem/scorer.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 # Author: Nianze A. TAO (Omozawa SUENO)
 """
-Scorers.
+Define essential scorers.
 """
 from typing import List, Callable, Union, Optional
 import torch

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/bayesianflow_for_chem/tool.py

@@ -1,11 +1,12 @@
 # -*- coding: utf-8 -*-
 # Author: Nianze A. TAO (Omozawa SUENO)
 """
-Tools.
+Essential tools.
 """
 import re
 import csv
 import random
+import warnings
 from copy import deepcopy
 from pathlib import Path
 from typing import List, Dict, Tuple, Union, Optional
@@ -16,7 +17,8 @@ from torch import cuda, Tensor, softmax
 from torch.ao import quantization
 from torch.utils.data import DataLoader
 from typing_extensions import Self
-from rdkit.Chem import rdDetermineBonds, Bond, MolFromXYZBlock, CanonicalRankAtoms
+from rdkit.Chem.rdchem import Mol, Bond
+from rdkit.Chem import rdDetermineBonds, MolFromXYZBlock, MolToSmiles, CanonSmiles
 from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles  # type: ignore
 from sklearn.metrics import (
     roc_auc_score,
@@ -32,12 +34,11 @@ try:
 
     _use_pynauty = True
 except ImportError:
-    import warnings
+    import platform
 
     _use_pynauty = False
-
 from .data import VOCAB_KEYS
-from .model import ChemBFN, MLP, Linear
+from .model import ChemBFN, MLP, Linear, EnsembleChemBFN
 
 
 _atom_regex_pattern = (
@@ -161,6 +162,8 @@ def split_dataset(
     :return:
     :rtype: None
     """
+    if isinstance(file, Path):
+        file = file.__str__()
     assert file.endswith(".csv")
     assert len(split_ratio) == 3
     assert method in ("random", "scaffold")
@@ -170,7 +173,7 @@ def split_dataset(
     raw_data = data[1:]
     smiles_idx = []  # only first index will be used
     for key, h in enumerate(header):
-        if h.lower() == "smiles":
+        if "smiles" in h.lower():
             smiles_idx.append(key)
     assert len(smiles_idx) > 0
     data_len = len(raw_data)
@@ -186,11 +189,22 @@ def split_dataset(
         scaffolds: Dict[str, List] = {}
         for key, d in enumerate(raw_data):
             # compute Bemis-Murcko scaffold
-            scaffold = MurckoScaffoldSmiles(d[smiles_idx[0]])
-            if scaffold in scaffolds:
-                scaffolds[scaffold].append(key)
-            else:
-                scaffolds[scaffold] = [key]
+            if len(smiles_idx) > 1:
+                warnings.warn(
+                    "\033[32;1m"
+                    f"We found {len(smiles_idx)} SMILES strings in a row!"
+                    " Only the first SMILES will be used to compute the molecular scaffold."
+                    "\033[0m",
+                    stacklevel=2,
+                )
+            try:
+                scaffold = MurckoScaffoldSmiles(d[smiles_idx[0]])
+                if scaffold in scaffolds:
+                    scaffolds[scaffold].append(key)
+                else:
+                    scaffolds[scaffold] = [key]
+            except ValueError:  # do nothing when SMILES is not valid
+                ...
         scaffolds = {key: sorted(value) for key, value in scaffolds.items()}
         train_set, test_set, val_set = [], [], []
         for idxs in scaffolds.values():
@@ -212,137 +226,13 @@ def split_dataset(
         writer.writerows([header] + val_set)
 
 
-def geo2seq(
-    symbols: List[str],
-    coordinates: np.ndarray,
-    decimals: int = 2,
-    angle_unit: str = "degree",
-) -> str:
-    """
-    Geometry-to-sequence function.\n
-    The algorithm follows the descriptions in paper: https://arxiv.org/abs/2408.10120.
-
-    :param symbols: a list of atomic symbols
-    :param coordinates: Cartesian coordinates;  shape: (n_a, 3)
-    :param decimals: number of decimal places to round to
-    :param angle_unit: `'degree'` or `'radian'`
-    :type symbols: list
-    :type coordinates: numpy.ndarray
-    :type decimals: int
-    :type angle_unit: str
-    :return: `Geo2Seq` string
-    :rtype: str
-    """
-    assert angle_unit in ("degree", "radian")
-    angle_scale = 180 / np.pi if angle_unit == "degree" else 1.0
-    n = len(symbols)
-    if n == 1:
-        return f"{symbols[0]} {'0.0'} {'0.0'} {'0.0'}"
-    xyz_block = [str(n), ""]
-    for i, atom in enumerate(symbols):
-        xyz_block.append(
-            f"{atom} {'%.10f' % coordinates[i][0].item()} {'%.10f' % coordinates[i][1].item()} {'%.10f' % coordinates[i][2].item()}"
-        )
-    mol = MolFromXYZBlock("\n".join(xyz_block))
-    rdDetermineBonds.DetermineConnectivity(mol)
-    # ------- Canonicalization -------
-    if _use_pynauty:
-        pair_idx = np.array(_bond_pair_idx(mol.GetBonds())).T.tolist()
-        pair_dict: Dict[int, List[int]] = {}
-        for key, i in enumerate(pair_idx[0]):
-            if i not in pair_dict:
-                pair_dict[i] = [pair_idx[1][key]]
-            else:
-                pair_dict[i].append(pair_idx[1][key])
-        g = Graph(n, adjacency_dict=pair_dict)
-        cl = canon_label(g)  # type: list
-    else:
-        warnings.warn(
-            "\033[32;1m"
-            "`pynauty` is not installed."
-            " Switched to canonicalization function provided by `rdkit`."
-            " This is the expected behaviour only if you are working on Windows platform."
-            "\033[0m",
-            stacklevel=2,
-        )
-        cl = list(CanonicalRankAtoms(mol, breakTies=True))
-    symbols = np.array([[s] for s in symbols])[cl].flatten().tolist()
-    coordinates = coordinates[cl]
-    # ------- Find global coordinate frame -------
-    if n == 2:
-        d = np.round(np.linalg.norm(coordinates[0] - coordinates[1], 2), decimals)
-        return f"{symbols[0]} {'0.0'} {'0.0'} {'0.0'} {symbols[1]} {d} {'0.0'} {'0.0'}"
-    for idx_0 in range(n - 2):
-        _vec0 = coordinates[idx_0] - coordinates[idx_0 + 1]
-        _vec1 = coordinates[idx_0] - coordinates[idx_0 + 2]
-        _d1 = np.linalg.norm(_vec0, 2)
-        _d2 = np.linalg.norm(_vec1, 2)
-        if 1 - np.abs(np.dot(_vec0, _vec1) / (_d1 * _d2)) > 1e-6:
-            break
-    x = (coordinates[idx_0 + 1] - coordinates[idx_0]) / _d1
-    y = np.cross((coordinates[idx_0 + 2] - coordinates[idx_0]), x)
-    y_d = np.linalg.norm(y, 2)
-    y = y / np.ma.filled(np.ma.array(y_d, mask=y_d == 0), np.inf)
-    z = np.cross(x, y)
-    # ------- Build spherical coordinates -------
-    vec = coordinates - coordinates[idx_0]
-    d = np.linalg.norm(vec, 2, axis=-1)
-    _d = np.ma.filled(np.ma.array(d, mask=d == 0), np.inf)
-    theta = angle_scale * np.arccos(np.dot(vec, z) / _d)  # in [0, \pi]
-    phi = angle_scale * np.arctan2(np.dot(vec, y), np.dot(vec, x))  # in [-\pi, \pi]
-    info = np.vstack([d, theta, phi]).T
-    info[idx_0] = np.zeros(3)
-    info = [
-        f"{symbols[i]} {r[0]} {r[1]} {r[2]}"
-        for i, r in enumerate(np.round(info, decimals))
-    ]
-    return " ".join(info)
-
-
-def seq2geo(
-    seq: str, angle_unit: str = "degree"
-) -> Optional[Tuple[List[str], List[List[float]]]]:
-    """
-    Sequence-to-geometry function.\n
-    The method follows the descriptions in paper: https://arxiv.org/abs/2408.10120.
-
-    :param seq: `Geo2Seq` string
-    :param angle_unit: `'degree'` or `'radian'`
-    :type seq: str
-    :type angle_unit: str
-    :return: (symbols, coordinates) if `seq` is valid
-    :rtype: tuple | None
-    """
-    assert angle_unit in ("degree", "radian")
-    angle_scale = np.pi / 180 if angle_unit == "degree" else 1.0
-    tokens = seq.split()
-    if len(tokens) % 4 == 0:
-        tokens = np.array(tokens).reshape(-1, 4).tolist()
-        symbols, coordinates = [], []
-        for i in tokens:
-            symbol = i[0]
-            if len(_atom_regex.findall(symbol)) != 1:
-                return None
-            symbols.append(symbol)
-            try:
-                d, theta, phi = float(i[1]), float(i[2]), float(i[3])
-                x = d * np.sin(theta * angle_scale) * np.cos(phi * angle_scale)
-                y = d * np.sin(theta * angle_scale) * np.sin(phi * angle_scale)
-                z = d * np.cos(theta * angle_scale)
-                coordinates.append([x.item(), y.item(), z.item()])
-            except ValueError:
-                return None
-        return symbols, coordinates
-    return None
-
-
 @torch.no_grad()
 def sample(
-    model: ChemBFN,
+    model: Union[ChemBFN, EnsembleChemBFN],
     batch_size: int,
     sequence_size: int,
     sample_step: int = 100,
-    y: Optional[Tensor] = None,
+    y: Optional[Union[Tensor, Dict[str, Tensor], List[Tensor]]] = None,
     guidance_strength: float = 4.0,
     device: Union[str, torch.device, None] = None,
     vocab_keys: List[str] = VOCAB_KEYS,
@@ -358,7 +248,9 @@ def sample(
     :param batch_size: batch size
     :param sequence_size: max sequence length
     :param sample_step: number of sampling steps
-    :param y: conditioning vector;  shape: (n_b, 1, n_f)
+    :param y: conditioning vector;             shape: (n_b, 1, n_f) or (n_b, n_f) \n
+              or a list/`dict` of conditions;  shape: (n_b, n_c) * n_h
+
     :param guidance_strength: strength of conditional generation. It is not used if y is null.
     :param device: hardware accelerator
     :param vocab_keys: a list of (ordered) vocabulary
@@ -366,11 +258,11 @@ def sample(
     :param method: sampling method chosen from `"ODE:x"` or `"BFN"` where `x` is the value of sampling temperature; default is `"BFN"`
     :param allowed_tokens: a list of allowed tokens; default is `"all"`
     :param sort: whether to sort the samples according to entropy values; default is `False`
-    :type model: bayesianflow_for_chem.model.ChemBFN
+    :type model: bayesianflow_for_chem.model.ChemBFN | bayesianflow_for_chem.model.EnsembleChemBFN
     :type batch_size: int
     :type sequence_size: int
     :type sample_step: int
-    :type y: torch.Tensor | None
+    :type y: torch.Tensor | list | dict | None
     :type guidance_strength: float
     :type device: str | torch.device | None
     :type vocab_keys: list
@@ -382,11 +274,23 @@ def sample(
     :rtype: list
     """
     assert method.split(":")[0].lower() in ("ode", "bfn")
+    if isinstance(model, EnsembleChemBFN):
+        assert y is not None, "conditioning is required while using an ensemble model."
+        assert isinstance(y, list) or isinstance(y, dict)
+    else:
+        assert isinstance(y, Tensor) or y is None
     if device is None:
         device = _find_device()
     model.to(device).eval()
     if y is not None:
-        y = y.to(device)
+        if isinstance(y, Tensor):
+            y = y.to(device)
+        elif isinstance(y, list):
+            y = [i.to(device) for i in y]
+        elif isinstance(y, dict):
+            y = {k: v.to(device) for k, v in y.items()}
+        else:
+            raise NotImplementedError
     if isinstance(allowed_tokens, list):
         token_mask = [0 if i in allowed_tokens else 1 for i in vocab_keys]
         token_mask = torch.tensor([[token_mask]], dtype=torch.bool).to(device)
@@ -416,10 +320,10 @@ def sample(
 
 @torch.no_grad()
 def inpaint(
-    model: ChemBFN,
+    model: Union[ChemBFN, EnsembleChemBFN],
     x: Tensor,
     sample_step: int = 100,
-    y: Optional[Tensor] = None,
+    y: Optional[Union[Tensor, Dict[str, Tensor], List[Tensor]]] = None,
     guidance_strength: float = 4.0,
     device: Union[str, torch.device, None] = None,
     vocab_keys: List[str] = VOCAB_KEYS,
@@ -434,7 +338,9 @@ def inpaint(
     :param model: trained ChemBFN model
     :param x: categorical indices of scaffold;  shape: (n_b, n_t)
     :param sample_step: number of sampling steps
-    :param y: conditioning vector;              shape: (n_b, 1, n_f)
+    :param y: conditioning vector;              shape: (n_b, 1, n_f) or (n_b, n_f) \n
+              or a list/`dict` of conditions;   shape: (n_b, n_c) * n_h
+
     :param guidance_strength: strength of conditional generation. It is not used if y is null.
     :param device: hardware accelerator
     :param vocab_keys: a list of (ordered) vocabulary
@@ -442,10 +348,10 @@ def inpaint(
     :param method: sampling method chosen from `"ODE:x"` or `"BFN"` where `x` is the value of sampling temperature; default is `"BFN"`
     :param allowed_tokens: a list of allowed tokens; default is `"all"`
     :param sort: whether to sort the samples according to entropy values; default is `False`
-    :type model: bayesianflow_for_chem.model.ChemBFN
+    :type model: bayesianflow_for_chem.model.ChemBFN | bayesianflow_for_chem.model.EnsembleChemBFN
     :type x: torch.Tensor
     :type sample_step: int
-    :type y: torch.Tensor | None
+    :type y: torch.Tensor | list | dict | None
     :type guidance_strength: float
     :type device: str | torch.device | None
     :type vocab_keys: list
@@ -457,12 +363,24 @@ def inpaint(
     :rtype: list
     """
     assert method.split(":")[0].lower() in ("ode", "bfn")
+    if isinstance(model, EnsembleChemBFN):
+        assert y is not None, "conditioning is required while using an ensemble model."
+        assert isinstance(y, list) or isinstance(y, dict)
+    else:
+        assert isinstance(y, Tensor) or y is None
     if device is None:
         device = _find_device()
     model.to(device).eval()
     x = x.to(device)
     if y is not None:
-        y = y.to(device)
+        if isinstance(y, Tensor):
+            y = y.to(device)
+        elif isinstance(y, list):
+            y = [i.to(device) for i in y]
+        elif isinstance(y, dict):
+            y = {k: v.to(device) for k, v in y.items()}
+        else:
+            raise NotImplementedError
     if isinstance(allowed_tokens, list):
         token_mask = [0 if i in allowed_tokens else 1 for i in vocab_keys]
         token_mask = torch.tensor([[token_mask]], dtype=torch.bool).to(device)
@@ -627,3 +545,178 @@ def quantise_model(model: ChemBFN) -> nn.Module:
         model, {nn.Linear, Linear}, torch.qint8, mapping
     )
     return quantised_model
+
+
+class GeometryConverter:
+    """
+    Converting between different 2D/3D molecular representations.
+    """
+
+    @staticmethod
+    def _xyz2mol(symbols: List[str], coordinates: np.ndarray) -> Mol:
+        xyz_block = [str(len(symbols)), ""]
+        r = coordinates
+        for i, atom in enumerate(symbols):
+            xyz_block.append(f"{atom} {r[i][0]:.10f} {r[i][1]:.10f} {r[i][2]:.10f}")
+        return MolFromXYZBlock("\n".join(xyz_block))
+
+    def cartesian2smiles(
+        self,
+        symbols: List[str],
+        coordinates: np.ndarray,
+        charge: int = 0,
+        canonical: bool = True,
+    ) -> str:
+        """
+        Transform (guess out) molecular geometry to SMILES string.
+
+        :param symbols: a list of atomic symbols
+        :param coordinates: Cartesian coordinates;  shape: (n_a, 3)
+        :param charge: net charge
+        :param canonical: whether to canonicalise the SMILES
+        :type symbols: list
+        :type coordinates: numpy.ndarray
+        :type charge: int
+        :type canonical: bool
+        :return: SMILES string
+        :rtype: str
+        """
+        mol = self._xyz2mol(symbols, coordinates)
+        rdDetermineBonds.DetermineBonds(mol, charge=charge)
+        smiles = MolToSmiles(mol)
+        if canonical:
+            smiles = CanonSmiles(smiles)
+        return smiles
+
+    def canonicalise(
+        self, symbols: List[str], coordinates: np.ndarray
+    ) -> Tuple[List[str], np.ndarray]:
+        """
+        Canonicalising the 3D molecular graph.
+
+        :param symbols: a list of atomic symbols
+        :param coordinates: Cartesian coordinates;  shape: (n_a, 3)
+        :type symbols: list
+        :type coordinates: numpy.ndarray
+        :return: canonicalised symbols \n
+                 canonicalised coordinates;         shape: (n_a, 3)
+        :rtype: tuple
+        """
+        if not _use_pynauty:
+            if platform.system() == "Windows":
+                raise NotImplementedError(
+                    "This method is not implemented on Windows platform."
+                )
+            else:
+                raise ImportError("`pynauty` is not installed.")
+        n = len(symbols)
+        if n == 1:
+            return symbols, coordinates
+        mol = self._xyz2mol(symbols, coordinates)
+        rdDetermineBonds.DetermineConnectivity(mol)
+        # ------- Canonicalization -------
+        pair_idx = np.array(_bond_pair_idx(mol.GetBonds())).T.tolist()
+        pair_dict: Dict[int, List[int]] = {}
+        for key, i in enumerate(pair_idx[0]):
+            if i not in pair_dict:
+                pair_dict[i] = [pair_idx[1][key]]
+            else:
+                pair_dict[i].append(pair_idx[1][key])
+        g = Graph(n, adjacency_dict=pair_dict)
+        cl = canon_label(g)  # type: list
+        symbols = np.array([[s] for s in symbols])[cl].flatten().tolist()
+        coordinates = coordinates[cl]
+        return symbols, coordinates
+
+    @staticmethod
+    def cartesian2spherical(coordinates: np.ndarray) -> np.ndarray:
+        """
+        Transforming Cartesian coordinate to spherical form.\n
+        The method is adapted from the paper: https://arxiv.org/abs/2408.10120.
+
+        :param coordinates: Cartesian coordinates;  shape: (n_a, 3)
+        :type coordinates: numpy.ndarray
+        :return: spherical coordinates;             shape: (n_a, 3)
+        :rtype: numpy.ndarray
+        """
+        n = coordinates.shape[0]
+        if n == 1:
+            return np.array([[0.0, 0.0, 0.0]])
+        # ------- Find global coordinate frame -------
+        if n == 2:
+            d = np.linalg.norm(coordinates[0] - coordinates[1], 2)
+            return np.array([[0.0, 0.0, 0.0], [d, 0.0, 0.0]])
+        for idx_0 in range(n - 2):
+            _vec0 = coordinates[idx_0] - coordinates[idx_0 + 1]
+            _vec1 = coordinates[idx_0] - coordinates[idx_0 + 2]
+            _d1 = np.linalg.norm(_vec0, 2)
+            _d2 = np.linalg.norm(_vec1, 2)
+            if 1 - np.abs(np.dot(_vec0, _vec1) / (_d1 * _d2)) > 1e-6:
+                break
+        x = (coordinates[idx_0 + 1] - coordinates[idx_0]) / _d1
+        y = np.cross((coordinates[idx_0 + 2] - coordinates[idx_0]), x)
+        y_d = np.linalg.norm(y, 2)
+        y = y / np.ma.filled(np.ma.array(y_d, mask=y_d == 0), np.inf)
+        z = np.cross(x, y)
+        # ------- Build spherical coordinates -------
+        vec = coordinates - coordinates[idx_0]
+        d = np.linalg.norm(vec, 2, axis=-1)
+        _d = np.ma.filled(np.ma.array(d, mask=d == 0), np.inf)
+        theta = np.arccos(np.dot(vec, z) / _d)  # in [0, \pi]
+        phi = np.arctan2(np.dot(vec, y), np.dot(vec, x))  # in [-\pi, \pi]
+        info = np.vstack([d, theta, phi]).T
+        info[idx_0] = np.zeros_like(info[idx_0])
+        return info
+
+    def geo2seq(
+        self, symbols: List[str], coordinates: np.ndarray, decimals: int = 2
+    ) -> str:
+        """
+        Geometry-to-sequence function.\n
+        The algorithm follows the descriptions in paper: https://arxiv.org/abs/2408.10120.
+
+        :param symbols: a list of atomic symbols
+        :param coordinates: Cartesian coordinates;  shape: (n_a, 3)
+        :param decimals: the maxmium number of decimals to keep; default is 2
+        :type symbols: list
+        :type coordinates: numpy.ndarray
+        :type decimals: int
+        :return: `Geo2Seq` string
+        :rtype: str
+        """
+        symbols, coordinates = self.canonicalise(symbols, coordinates)
+        info = self.cartesian2spherical(coordinates)
+        info = [
+            f"{symbols[i]} {r[0]} {r[1]} {r[2]}"
+            for i, r in enumerate(np.round(info, decimals))
+        ]
+        return " ".join(info)
+
+    @staticmethod
+    def seq2geo(seq: str) -> Tuple[Optional[List[str]], Optional[np.ndarray]]:
+        """
+        Sequence-to-geometry function.\n
+        The method follows the descriptions in paper: https://arxiv.org/abs/2408.10120.
+
+        :param seq: `Geo2Seq` string
+        :type seq: str
+        :return: (symbols, coordinates) if `seq` is valid
+        :rtype: tuple
+        """
+        tokens = seq.split()
+        if len(tokens) % 4 != 0:
+            return None, None
+        tokens = np.array(tokens).reshape(-1, 4)
+        symbols, coordinates = tokens[::, 0], tokens[::, 1:]
+        if sum([len(_atom_regex.findall(sym)) for sym in symbols]) != len(symbols):
+            return None, None
+        try:
+            coord = [[float(i) for i in j] for j in coordinates]
+            coord = np.array(coord)
+        except ValueError:
+            return None, None
+        d, theta, phi = coord[::, 0, None], coord[::, 1, None], coord[::, 2, None]
+        x = d * np.sin(theta) * np.cos(phi)
+        y = d * np.sin(theta) * np.sin(phi)
+        z = d * np.cos(theta)
+        return symbols, np.concatenate([x, y, z], -1)

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/bayesianflow_for_chem/train.py

@@ -37,7 +37,8 @@ class Model(LightningModule):
         """
         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`.
+        the model(s) will be saved to `YOUR_WORK_DIR/model.pt` (if LoRA is enabled then `YOUR_WORK_DIR/lora.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`.
@@ -135,7 +136,8 @@ class Regressor(LightningModule):
         """
         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`.
+        the models will be saved to `YOUR_WORK_DIR/model_ft.pt` (if LoRA is enabled then `YOUR_WORK_DIR/lora.pt`)
+        and `YOUR_WORK_DIR/readout.pt`.
 
         :param model: `~bayesianflow_for_chem.model.ChemBFN` instance.
         :param mlp: `~bayesianflow_for_chem.model.MLP` instance.
@@ -218,7 +220,7 @@ class Regressor(LightningModule):
         """
         Save the trained model.
 
-        :param workdir: the directory to save the model
+        :param workdir: the directory to save the models
         :type workdir: pathlib.Path
         :return:
         :rtype: None

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/bayesianflow_for_chem.egg-info/PKG-INFO

@@ -1,16 +1,15 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 1.2.6
+Version: 1.3.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
+License: AGPL-3.0-or-later
 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
@@ -39,6 +38,7 @@ Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: keywords
 Dynamic: license
+Dynamic: license-file
 Dynamic: project-url
 Dynamic: provides-extra
 Dynamic: requires-dist
@@ -86,7 +86,7 @@ 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).
+You can find pretrained models on our [🤗Hugging Face model page](https://huggingface.co/suenoomozawa/ChemBFN).
 
 ## Dataset Handling
 

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/pyproject.toml

@@ -1,3 +1,3 @@
 [build-system]
-requires = ["setuptools >= 75.0"]
+requires = ["setuptools >= 77.0.3"]
 build-backend = "setuptools.build_meta"

{bayesianflow_for_chem-1.2.6 → bayesianflow_for_chem-1.3.0}/setup.py

@@ -28,7 +28,8 @@ setup(
     description="Bayesian flow network framework for Chemistry",
     long_description=long_description,
     long_description_content_type="text/markdown",
-    license="AGPL-3.0 licence",
+    license="AGPL-3.0-or-later",
+    license_files=["LICEN[CS]E*"],
     package_dir={"bayesianflow_for_chem": "bayesianflow_for_chem"},
     package_data={"bayesianflow_for_chem": ["./*.txt", "./*.py"]},
     include_package_data=True,
@@ -52,7 +53,6 @@ setup(
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: GNU Affero General Public License v3",
         "Natural Language :: English",
         "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.9",