bayesianflow-for-chem 2.0.5__py3-none-any.whl → 2.2.1__py3-none-any.whl

bayesianflow_for_chem/__init__.py

@@ -3,10 +3,8 @@
 """
 ChemBFN package.
 """
-import colorama
 from . import data, tool, train, scorer, spectra
 from .model import ChemBFN, MLP, EnsembleChemBFN
-from .cli import main_script
 
 __all__ = [
     "data",
@@ -18,7 +16,7 @@ __all__ = [
     "MLP",
     "EnsembleChemBFN",
 ]
-__version__ = "2.0.5"
+__version__ = "2.2.1"
 __author__ = "Nianze A. Tao (Omozawa Sueno)"
 
 
@@ -29,6 +27,9 @@ def main() -> None:
     :return:
     :rtype: None
     """
+    import colorama
+    from bayesianflow_for_chem.cli import main_script
+
     colorama.just_fix_windows_console()
     main_script(__version__)
     colorama.deinit()

bayesianflow_for_chem/cli.py

@@ -12,13 +12,8 @@ from pathlib import Path
 from functools import partial
 from typing import List, Tuple, Dict, Union, Callable
 import torch
-import lightning as L
 from rdkit.Chem import MolFromSmiles, CanonSmiles
-from torch.utils.data import DataLoader
-from lightning.pytorch import loggers
-from lightning.pytorch.callbacks import ModelCheckpoint
 from bayesianflow_for_chem import ChemBFN, MLP
-from bayesianflow_for_chem.train import Model
 from bayesianflow_for_chem.scorer import smiles_valid, Scorer
 from bayesianflow_for_chem.data import (
     VOCAB_COUNT,
@@ -32,7 +27,7 @@ from bayesianflow_for_chem.data import (
     collate,
     CSVData,
 )
-from bayesianflow_for_chem.tool import sample, inpaint
+from bayesianflow_for_chem.tool import sample, inpaint, optimise, adjust_lora_
 
 
 """
@@ -90,6 +85,7 @@ checkpoint_save_path = "home/user/project/ckpt"
 train_strategy = "auto"  # or any strategy supported by Lightning, e.g., "ddp"
 accumulate_grad_batches = 1
 enable_progress_bar = false
+plugin_script = ""  # define customised behaviours of dataset, datasetloader, etc in a python script
 
 # Remove this table if inference is unnecessary
 [inference]
@@ -99,9 +95,11 @@ sample_size = 1000  # the minimum number of samples you want
 sample_step = 100
 sample_method = "ODE:0.5"  # ODE-solver with temperature of 0.5; another choice is "BFN"
 semi_autoregressive = false
+lora_scaling = 1.0  # LoRA scaling if applied
 guidance_objective = [-0.023, 0.09, 0.113]  # if no objective is needed set it to empty array []
 guidance_objective_strength = 4.0  # unnecessary if guidance_objective = []
 guidance_scaffold = "c1ccccc1"  # if no scaffold is used set it to empty string ""
+sample_template = ""  # template for mol2mol task; leave it blank if scaffold is used
 unwanted_token = []
 exclude_invalid = true  # to only store valid samples
 exclude_duplicate = true  # to only store unique samples
@@ -118,6 +116,32 @@ madmadmadmadmadmadmadmadmadmadmadmadmadmadmad
 madmadmadmadmadmadmadmadmadmadmadmadmadmadmad
 """
 
+_ALLOWED_PLUGINS = [
+    "collate_fn",
+    "num_workers",
+    "max_sequence_length",
+    "shuffle",
+    "CustomData",
+]
+
+
+def _load_plugin(plugin_file: str) -> Dict[str, Union[int, Callable, object, None]]:
+    if not plugin_file:
+        return {n: None for n in _ALLOWED_PLUGINS}
+    from importlib import util as iutil
+
+    spec = iutil.spec_from_file_location(Path(plugin_file).stem, plugin_file)
+    plugins = iutil.module_from_spec(spec)
+    spec.loader.exec_module(plugins)
+    plugin_names: List[str] = plugins.__all__
+    plugin_dict = {}
+    for n in _ALLOWED_PLUGINS:
+        if n in plugin_names:
+            plugin_dict[n] = getattr(plugins, n)
+        else:
+            plugin_dict[n] = None
+    return plugin_dict
+
 
 def parse_cli(version: str) -> argparse.Namespace:
     """
@@ -130,7 +154,7 @@ def parse_cli(version: str) -> argparse.Namespace:
     """
     parser = argparse.ArgumentParser(
         description="Madmol: a CLI molecular design tool for "
-        "de novo design, R-group replacement, and sequence in-filling, "
+        "de novo design, R-group replacement, molecule optimisation, and sequence in-filling, "
         "based on generative route of ChemBFN method. "
         "Let's make some craziest molecules.",
         epilog=f"Madmol {version}, developed in Hiroshima University by chemists for chemists. "
@@ -157,7 +181,7 @@ def parse_cli(version: str) -> argparse.Namespace:
         "-D",
         "--dryrun",
         action="store_true",
-        help="dry-run to check the configurations",
+        help="dry-run to check the configurations and exit",
     )
     parser.add_argument("-V", "--version", action="version", version=version)
     return parser.parse_args()
@@ -265,6 +289,14 @@ def load_runtime_config(
                     f"\033[0;31mCritical\033[0;0m in {config_file}: Restart checkpoint file {ckpt_file} does not exist."
                 )
                 flag_critical += 1
+        # ↓ added in v2.2.0; need to be compatible with old versions.
+        plugin_script: str = config["train"].get("plugin_script", "")
+        if plugin_script:
+            if not os.path.exists(plugin_script):
+                print(
+                    f"\033[0;31mCritical\033[0;0m in {config_file}: Plugin script {plugin_script} does not exist."
+                )
+                flag_critical += 1
     if "inference" in config:
         if not "train" in config:
             if not isinstance(config["inference"]["sequence_length"], int):
@@ -284,6 +316,14 @@ def load_runtime_config(
                 f"\033[0;33mWarning\033[0;0m in {config_file}: Directory {result_dir} to save the result does not exist."
             )
             flag_warning += 1
+        if (
+            config["inference"]["guidance_scaffold"] != ""
+            and config["inference"]["sample_template"] != ""
+        ):
+            print(
+                f"\033[0;33mWarning\033[0;0m in {config_file}: Inpaint task or mol2mol task?"
+            )
+            flag_warning += 1
     return config, flag_critical, flag_warning
 
 
@@ -325,14 +365,15 @@ def main_script(version: str) -> None:
                 )
                 flag_warning += 1
         if not os.path.exists(runtime_config["train"]["checkpoint_save_path"]):
-            os.makedirs(runtime_config["train"]["checkpoint_save_path"])
+            if not parser.dryrun:  # only create it in real tasks
+                os.makedirs(runtime_config["train"]["checkpoint_save_path"])
     else:
         if not model_config["ChemBFN"]["base_model"]:
             print(
                 f"\033[0;33mWarning\033[0;0m in {parser.model_config}: You should load a pretrained ChemBFN model."
             )
             flag_warning += 1
-        if not model_config["MLP"]["base_model"]:
+        if "MLP" in model_config and not model_config["MLP"]["base_model"]:
             print(
                 f"\033[0;33mWarning\033[0;0m in {parser.model_config}: You should load a pretrained MLP."
             )
@@ -340,13 +381,17 @@ def main_script(version: str) -> None:
     if "inference" in runtime_config:
         if runtime_config["inference"]["guidance_objective"]:
             if not "MLP" in model_config:
-                print(f"Warning in {parser.model_config}: Oh no, you don't have a MLP.")
+                print(
+                    f"\033[0;33mWarning\033[0;0m in {parser.model_config}: Oh no, you don't have a MLP."
+                )
                 flag_warning += 1
     if parser.dryrun:
         if flag_critical != 0:
             print("Configuration check failed!")
         elif flag_warning != 0:
-            print("Your job will probably run, but it may not follow your expectation.")
+            print(
+                "Your job will probably run, but it may not follow your expectations."
+            )
         else:
             print("Configuration check passed.")
         return
@@ -405,6 +450,15 @@ def main_script(version: str) -> None:
         mlp = None
     # ------- train -------
     if "train" in runtime_config:
+        import lightning as L
+        from torch.utils.data import DataLoader
+        from lightning.pytorch import loggers
+        from lightning.pytorch.callbacks import ModelCheckpoint
+        from bayesianflow_for_chem.train import Model
+
+        # ####### get plugins #######
+        plugin_file = runtime_config["train"].get("plugin_script", "")
+        plugins = _load_plugin(plugin_file)
         # ####### build scorer #######
         if (tokeniser_name == "smiles" or tokeniser_name == "safe") and runtime_config[
             "train"
@@ -418,30 +472,43 @@ def main_script(version: str) -> None:
         mol_tag = runtime_config["train"]["molecule_tag"]
         obj_tag = runtime_config["train"]["objective_tag"]
         dataset_file = runtime_config["train"]["dataset"]
-        with open(dataset_file, "r") as db:
-            _data = db.readlines()
-        header = _data[0]
-        mol_idx = []
-        for i, tag in enumerate(header.replace("\n", "").split(",")):
-            if tag == mol_tag:
-                mol_idx.append(i)
-        _data_len = []
-        for i in _data[1:]:
-            i = i.replace("\n", "").split(",")
-            _mol = ".".join([i[j] for j in mol_idx])
-            _data_len.append(tokeniser(_mol).shape[-1])
-        lmax = max(_data_len)
-        dataset = CSVData(dataset_file)
+        if plugins["max_sequence_length"]:
+            lmax = plugins["max_sequence_length"]
+        else:
+            with open(dataset_file, "r") as db:
+                _data = db.readlines()
+            _header = _data[0]
+            _mol_idx = []
+            for i, tag in enumerate(_header.replace("\n", "").split(",")):
+                if tag == mol_tag:
+                    _mol_idx.append(i)
+            _data_len = []
+            for i in _data[1:]:
+                i = i.replace("\n", "").split(",")
+                _mol = ".".join([i[j] for j in _mol_idx])
+                _data_len.append(tokeniser(_mol).shape[-1])
+            lmax = max(_data_len)
+            del _data, _data_len, _header, _mol_idx  # clear memory
+        if plugins["CustomData"] is not None:
+            dataset = plugins["CustomData"](dataset_file)
+        else:
+            dataset = CSVData(dataset_file)
         dataset.map(
             partial(_encode, mol_tag=mol_tag, obj_tag=obj_tag, tokeniser=tokeniser)
         )
         dataloader = DataLoader(
             dataset,
             runtime_config["train"]["batch_size"],
-            True,
-            num_workers=4,
-            collate_fn=collate,
-            persistent_workers=True,
+            True if plugins["shuffle"] is None else plugins["shuffle"],
+            num_workers=4 if plugins["num_workers"] is None else plugins["num_workers"],
+            collate_fn=(
+                collate if plugins["collate_fn"] is None else plugins["collate_fn"]
+            ),
+            persistent_workers=(
+                True
+                if (plugins["num_workers"] is None or plugins["num_workers"] > 0)
+                else False
+            ),
         )
         # ####### build trainer #######
         logger_name = runtime_config["train"]["logger_name"].lower()
@@ -520,6 +587,8 @@ def main_script(version: str) -> None:
         if "train" in runtime_config:
             bfn = model.model
             mlp = model.mlp
+        # ↓ added in v2.1.0; need to be compatible with old versions
+        lora_scaling = runtime_config["inference"].get("lora_scaling", 1.0)
         # ####### strat inference #######
         bfn.semi_autoregressive = runtime_config["inference"]["semi_autoregressive"]
         _device = (
@@ -550,8 +619,16 @@ def main_script(version: str) -> None:
                 x[:-1], (0, sequence_length - x.shape[-1] + 1), value=0
             )
             x = x[None, :].repeat(batch_size, 1)
+            # then sample template will be ignored.
+        elif runtime_config["inference"]["sample_template"]:
+            template = runtime_config["inference"]["sample_template"]
+            x = tokeniser(template)
+            x = torch.nn.functional.pad(x, (0, sequence_length - x.shape[-1]), value=0)
+            x = x[None, :].repeat(batch_size, 1)
         else:
             x = None
+        if bfn.lora_enabled:
+            adjust_lora_(bfn, lora_scaling)
         mols = []
         while len(mols) < runtime_config["inference"]["sample_size"]:
             if x is None:
@@ -567,7 +644,7 @@ def main_script(version: str) -> None:
                     method=sample_method,
                     allowed_tokens=allowed_token,
                 )
-            else:
+            elif runtime_config["inference"]["guidance_scaffold"]:
                 s = inpaint(
                     bfn,
                     x,
@@ -579,6 +656,18 @@ def main_script(version: str) -> None:
                     method=sample_method,
                     allowed_tokens=allowed_token,
                 )
+            else:
+                s = optimise(
+                    bfn,
+                    x,
+                    sample_step,
+                    y,
+                    guidance_strength,
+                    _device,
+                    vocab_keys,
+                    method=sample_method,
+                    allowed_tokens=allowed_token,
+                )
             if runtime_config["inference"]["exclude_invalid"]:
                 s = [i for i in s if i]
                 if tokeniser_name == "smiles" or tokeniser_name == "safe":

bayesianflow_for_chem/data.py

@@ -61,7 +61,7 @@ def load_vocab(
     }
 
 
-_DEFUALT_VOCAB = load_vocab(__filedir__ / "vocab.txt")
+_DEFUALT_VOCAB = load_vocab(__filedir__ / "_data/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"]

bayesianflow_for_chem/model.py

@@ -659,12 +659,91 @@ class ChemBFN(nn.Module):
         return (-logits.gather(-1, x[..., :1]).squeeze(-1)).mean()
 
     @staticmethod
-    def reshape_y(y: Tensor) -> Tensor:
+    def _reshape(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
 
+    def _process(
+        self,
+        theta: Tensor,
+        mask: Optional[Tuple[Tensor, Tensor]],
+        y: Optional[Tensor],
+        sample_step: int,
+        guidance_strength: float,
+        token_mask: Optional[Tensor],
+    ) -> Tuple[Tensor, Tensor]:
+        # BFN inference process.
+        #
+        # theta: piror distribution;            shape: (n_b, n_t, n_vocab)
+        # mask: masked condition distribution;  shape: (n_b, n_t, n_vocab)
+        #       condition distribution mask;    shape: (n_b, n_t, 1)
+        n_b = theta.shape[0]
+        if y is not None:
+            y = self._reshape(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
+            p = self.discrete_output_distribution(theta, t, y, guidance_strength)
+            if token_mask is not None:
+                p = p.masked_fill_(token_mask, 0.0)
+            alpha = self.calc_discrete_alpha(t, t + 1 / sample_step)
+            e_k = nn.functional.one_hot(torch.argmax(p, -1), self.K).float()
+            mu = alpha * (self.K * e_k - 1)
+            sigma = (alpha * self.K).sqrt()
+            theta = (mu + sigma * torch.randn_like(mu)).exp() * theta
+            theta = theta / theta.sum(-1, True)
+            if mask is not None:
+                x_onehot, x_mask = mask
+                theta = x_onehot + (1 - x_mask) * theta
+        t_final = torch.ones((n_b, 1, 1), device=self.beta.device)
+        p = self.discrete_output_distribution(theta, t_final, y, guidance_strength)
+        entropy = -(p * p.log()).sum(-1).mean(-1)
+        if token_mask is not None:
+            p = p.masked_fill_(token_mask, 0.0)
+        return torch.argmax(p, -1), entropy
+
+    def _ode_process(
+        self,
+        z: Tensor,
+        mask: Optional[Tuple[Tensor, Tensor]],
+        y: Optional[Tensor],
+        sample_step: int,
+        guidance_strength: float,
+        token_mask: Optional[Tensor],
+        temperature: float,
+    ) -> Tuple[Tensor, Tensor]:
+        # ODE-solver engaged inference process.
+        #
+        # z: prior latent vector;               shape: (n_b, n_t, n_vocab)
+        # mask: masked condition distribution;  shape: (n_b, n_t, n_vocab)
+        #       condition distribution mask;    shape: (n_b, n_t, 1)
+        n_b = z.shape[0]
+        if y is not None:
+            y = self._reshape(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 = softmax(z, -1)
+            if mask is not None:
+                x_onehot, x_mask = mask
+                theta = x_onehot + (1 - x_mask) * theta
+            beta = self.calc_beta(t + 1 / sample_step)
+            p = self.discrete_output_distribution(theta, t, y, guidance_strength)
+            if token_mask is not None:
+                p = p.masked_fill_(token_mask, 0.0)
+            u = torch.randn_like(z)
+            z = (self.K * p - 1) * beta + (self.K * beta * temperature).sqrt() * u
+        t_final = torch.ones((n_b, 1, 1), device=self.beta.device)
+        theta = softmax(z, -1)
+        if mask is not None:
+            x_onehot, x_mask = mask
+            theta = x_onehot + (1 - x_mask) * theta
+        p = self.discrete_output_distribution(theta, t_final, y, guidance_strength)
+        entropy = -(p * p.log()).sum(-1).mean(-1)
+        if token_mask is not None:
+            p = p.masked_fill_(token_mask, 0.0)
+        return torch.argmax(p, -1), entropy
+
     @torch.jit.export
     def sample(
         self,
@@ -676,48 +755,30 @@ class ChemBFN(nn.Module):
         token_mask: Optional[Tensor] = None,
     ) -> Tuple[Tensor, Tensor]:
         """
-        Sample from a piror distribution.
+        Sample from a uniform piror distribution.
 
         :param batch_size: batch size
         :param sequence_size: max sequence length
-        :param y: conditioning vector;      shape: (n_b, 1, n_f) or (n_b, 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 assigning unwanted token(s) with `True`;
-                                            shape: (1, 1, n_vocab)
+                                         shape: (1, 1, n_vocab)
         :type batch_size: int
         :type sequence_size: int
         :type y: torch.Tensor | None
         :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)
+        :return: sampled token indices;  shape: (n_b, n_t) \n
+                 entropy of the tokens;  shape: (n_b)
         :rtype: tuple
         """
         theta = (
             torch.ones((batch_size, sequence_size, self.K), device=self.beta.device)
             / self.K
         )
-        if y is not None:
-            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)
-            if token_mask is not None:
-                p = p.masked_fill_(token_mask, 0.0)
-            alpha = self.calc_discrete_alpha(t, t + 1 / sample_step)
-            e_k = nn.functional.one_hot(torch.argmax(p, -1), self.K).float()
-            mu = alpha * (self.K * e_k - 1)
-            sigma = (alpha * self.K).sqrt()
-            theta = (mu + sigma * torch.randn_like(mu)).exp() * theta
-            theta = theta / theta.sum(-1, True)
-        t_final = torch.ones((batch_size, 1, 1), device=self.beta.device)
-        p = self.discrete_output_distribution(theta, t_final, y, guidance_strength)
-        entropy = -(p * p.log()).sum(-1).mean(-1)
-        if token_mask is not None:
-            p = p.masked_fill_(token_mask, 0.0)
-        return torch.argmax(p, -1), entropy
+        return self._process(theta, None, y, sample_step, guidance_strength, token_mask)
 
     @torch.jit.export
     def ode_sample(
@@ -735,11 +796,11 @@ 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) or (n_b, 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 assigning unwanted token(s) with `True`;
-                                            shape: (1, 1, n_vocab)
+                                         shape: (1, 1, n_vocab)
         :param temperature: sampling temperature
         :type batch_size: int
         :type sequence_size: int
@@ -748,29 +809,14 @@ class ChemBFN(nn.Module):
         :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)
+        :return: sampled token indices;  shape: (n_b, n_t) \n
+                 entropy of the tokens;  shape: (n_b)
         :rtype: tuple
         """
         z = torch.zeros((batch_size, sequence_size, self.K), device=self.beta.device)
-        if y is not None:
-            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)
-            beta = self.calc_beta(t + 1 / sample_step)
-            p = self.discrete_output_distribution(theta, t, y, guidance_strength)
-            if token_mask is not None:
-                p = p.masked_fill_(token_mask, 0.0)
-            u = torch.randn_like(z)
-            z = (self.K * p - 1) * beta + (self.K * beta * temperature).sqrt() * u
-        t_final = torch.ones((batch_size, 1, 1), device=self.beta.device)
-        theta = torch.softmax(z, -1)
-        p = self.discrete_output_distribution(theta, t_final, y, guidance_strength)
-        entropy = -(p * p.log()).sum(-1).mean(-1)
-        if token_mask is not None:
-            p = p.masked_fill_(token_mask, 0.0)
-        return torch.argmax(p, -1), entropy
+        return self._ode_process(
+            z, None, y, sample_step, guidance_strength, token_mask, temperature
+        )
 
     @torch.jit.export
     def inpaint(
@@ -800,30 +846,12 @@ class ChemBFN(nn.Module):
         :rtype: tuple
         """
         n_b, n_t = x.shape
-        mask = (x != 0).float()[..., None]
+        x_mask = (x != 0).float()[..., None]
         theta = torch.ones((n_b, n_t, self.K), device=x.device) / self.K
-        x_onehot = nn.functional.one_hot(x, self.K) * mask
-        theta = x_onehot + (1 - mask) * theta
-        if y is not None:
-            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)
-            if token_mask is not None:
-                p = p.masked_fill_(token_mask, 0.0)
-            alpha = self.calc_discrete_alpha(t, t + 1 / sample_step)
-            e_k = nn.functional.one_hot(torch.argmax(p, -1), self.K).float()
-            mu = alpha * (self.K * e_k - 1)
-            sigma = (alpha * self.K).sqrt()
-            theta = (mu + sigma * torch.randn_like(mu)).exp() * theta
-            theta = theta / theta.sum(-1, True)
-            theta = x_onehot + (1 - mask) * theta
-        t_final = torch.ones((n_b, 1, 1), device=x.device)
-        p = self.discrete_output_distribution(theta, t_final, y, guidance_strength)
-        entropy = -(p * p.log()).sum(-1).mean(-1)
-        if token_mask is not None:
-            p = p.masked_fill_(token_mask, 0.0)
-        return torch.argmax(p, -1), entropy
+        x_onehot = nn.functional.one_hot(x, self.K) * x_mask
+        theta = x_onehot + (1 - x_mask) * theta
+        mask = (x_onehot, x_mask)
+        return self._process(theta, mask, y, sample_step, guidance_strength, token_mask)
 
     @torch.jit.export
     def ode_inpaint(
@@ -856,29 +884,80 @@ class ChemBFN(nn.Module):
         :rtype: tuple
         """
         n_b, n_t = x.shape
-        mask = (x != 0).float()[..., None]
-        x_onehot = nn.functional.one_hot(x, self.K) * mask
+        x_mask = (x != 0).float()[..., None]
+        x_onehot = nn.functional.one_hot(x, self.K) * x_mask
         z = torch.zeros((n_b, n_t, self.K), device=self.beta.device)
-        if y is not None:
-            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)
-            theta = x_onehot + (1 - mask) * theta
-            beta = self.calc_beta(t + 1 / sample_step)
-            p = self.discrete_output_distribution(theta, t, y, guidance_strength)
-            if token_mask is not None:
-                p = p.masked_fill_(token_mask, 0.0)
-            u = torch.randn_like(z)
-            z = (self.K * p - 1) * beta + (self.K * beta * temperature).sqrt() * u
-        t_final = torch.ones((n_b, 1, 1), device=self.beta.device)
-        theta = torch.softmax(z, -1)
-        theta = x_onehot + (1 - mask) * theta
-        p = self.discrete_output_distribution(theta, t_final, y, guidance_strength)
-        entropy = -(p * p.log()).sum(-1).mean(-1)
-        if token_mask is not None:
-            p = p.masked_fill_(token_mask, 0.0)
-        return torch.argmax(p, -1), entropy
+        mask = (x_onehot, x_mask)
+        return self._ode_process(
+            z, mask, y, sample_step, guidance_strength, token_mask, temperature
+        )
+
+    @torch.jit.export
+    def optimise(
+        self,
+        x: Tensor,
+        y: Optional[Tensor] = None,
+        sample_step: int = 100,
+        guidance_strength: float = 4.0,
+        token_mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Optimise the template molecule (mol2mol). \n
+        This method is equivalent to sampling from a customised prior distribution.
+
+        :param x: categorical indices of template;  shape: (n_b, n_t)
+        :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 assigning unwanted token(s) with `True`;
+                                                    shape: (1, 1, n_vocab)
+        :type x: torch.Tensor
+        :type y: torch.Tensor | None
+        :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
+        """
+        x_onehot = nn.functional.one_hot(x, self.K).float()
+        theta = softmax(x_onehot, -1)
+        return self._process(theta, None, y, sample_step, guidance_strength, token_mask)
+
+    @torch.jit.export
+    def ode_optimise(
+        self,
+        x: Tensor,
+        y: Optional[Tensor] = None,
+        sample_step: int = 100,
+        guidance_strength: float = 4.0,
+        token_mask: Optional[Tensor] = None,
+        temperature: float = 0.5,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        ODE mol2mol.
+
+        :param x: categorical indices of template;  shape: (n_b, n_t)
+        :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 assigning unwanted token(s) with `True`;
+                                                    shape: (1, 1, n_vocab)
+        :param temperature: sampling temperature
+        :type x: torch.Tensor
+        :type y: torch.Tensor | None
+        :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
+        """
+        z = nn.functional.one_hot(x, self.K).float()
+        return self._ode_process(
+            z, None, y, sample_step, guidance_strength, token_mask, temperature
+        )
 
     def inference(
         self, x: Tensor, mlp: MLP, embed_fn: Optional[Callable[[Tensor], Tensor]] = None
@@ -1052,22 +1131,6 @@ class EnsembleChemBFN(ChemBFN):
                     module.lora_dropout = None
             v.lora_enabled = False
 
-    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:
@@ -1102,8 +1165,24 @@ class EnsembleChemBFN(ChemBFN):
             p_cond += p_cond_ * self.adapter_weights[name]
         return softmax((1 + w) * p_cond - w * p_uncond, -1)
 
+    def _map_to_dict(
+        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
+
     @staticmethod
-    def reshape_y(y: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    def _reshape(y: Dict[str, Tensor]) -> Dict[str, Tensor]:
         for k in y:
             assert y[k].dim() <= 3
             if y[k].dim() == 2:
@@ -1139,7 +1218,7 @@ class EnsembleChemBFN(ChemBFN):
                  entropy of the tokens;          shape: (n_b)
         :rtype: tuple
         """
-        y = self.construct_y(conditions)
+        y = self._map_to_dict(conditions)
         return super().sample(
             batch_size, sequence_size, y, sample_step, guidance_strength, token_mask
         )
@@ -1176,7 +1255,7 @@ class EnsembleChemBFN(ChemBFN):
                  entropy of the tokens;          shape: (n_b)
         :rtype: tuple
         """
-        y = self.construct_y(conditions)
+        y = self._map_to_dict(conditions)
         return super().ode_sample(
             batch_size,
             sequence_size,
@@ -1213,7 +1292,7 @@ class EnsembleChemBFN(ChemBFN):
                  entropy of the tokens;             shape: (n_b)
         :rtype: tuple
         """
-        y = self.construct_y(conditions)
+        y = self._map_to_dict(conditions)
         return super().inpaint(x, y, sample_step, guidance_strength, token_mask)
 
     @torch.inference_mode()
@@ -1245,11 +1324,76 @@ class EnsembleChemBFN(ChemBFN):
                  entropy of the tokens;             shape: (n_b)
         :rtype: tuple
         """
-        y = self.construct_y(conditions)
+        y = self._map_to_dict(conditions)
         return super().ode_inpaint(
             x, y, sample_step, guidance_strength, token_mask, temperature
         )
 
+    @torch.inference_mode()
+    def optimise(
+        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]:
+        """
+        Optimise the template molecule (mol2mol). \n
+        This method is equivalent to sampling from a customised prior distribution.
+
+        :param x: categorical indices of template;  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 assigning unwanted token(s) with `True`;
+                                                    shape: (1, 1, n_vocab)
+        :type x: torch.Tensor
+        :type y: torch.Tensor | None
+        :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._map_to_dict(conditions)
+        return super().optimise(x, y, sample_step, guidance_strength, token_mask)
+
+    @torch.inference_mode()
+    def ode_optimise(
+        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 template;  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._map_to_dict(conditions)
+        return super().ode_optimise(
+            x, y, sample_step, guidance_strength, token_mask, temperature
+        )
+
     def quantise(
         self, quantise_method: Optional[Callable[[ChemBFN], None]] = None
     ) -> None:

bayesianflow_for_chem/tool.py

@@ -9,7 +9,6 @@ import warnings
 from pathlib import Path
 from typing import List, Dict, Tuple, Union, Optional
 import torch
-import colorama
 import numpy as np
 from torch import cuda, Tensor, softmax
 from torch.utils.data import DataLoader
@@ -24,15 +23,7 @@ from rdkit.Chem import (
     AddHs,
     Mol,
 )
-from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles  # type: ignore
-from sklearn.metrics import (
-    roc_auc_score,
-    auc,
-    precision_recall_curve,
-    r2_score,
-    mean_absolute_error,
-    root_mean_squared_error,
-)
+from rdkit.Chem.Scaffolds.MurckoScaffold import MurckoScaffoldSmiles
 from .data import VOCAB_KEYS
 from .model import ChemBFN, MLP, EnsembleChemBFN
 
@@ -45,6 +36,68 @@ def _find_device() -> torch.device:
     return torch.device("cpu")
 
 
+def _parse_and_assert_param(
+    model: Union[ChemBFN, EnsembleChemBFN],
+    y: Optional[Union[Tensor, Dict[str, Tensor], List[Tensor]]],
+    method: str,
+) -> Optional[float]:
+    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 "ode" in method.lower():
+        tp = float(method.split(":")[-1])
+        assert tp > 0, "Sampling temperature should be higher than 0."
+        return tp
+    return None
+
+
+def _map_to_device(
+    y: Optional[Union[Tensor, Dict[str, Tensor], List[Tensor]]],
+    device: Union[str, torch.device],
+) -> Optional[Union[Tensor, Dict[str, Tensor], List[Tensor]]]:
+    if y is not None:
+        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
+    return y
+
+
+def _build_token_mask(
+    allowed_tokens: Union[str, List[str]],
+    vocab_keys: List[str],
+    device: Union[str, torch.tensor],
+) -> Optional[Tensor]:
+    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)
+    else:
+        token_mask = None
+    return token_mask
+
+
+def _token_to_seq(
+    tokens: Tensor, entropy: Tensor, vocab_keys: List[str], separator: str, sort: bool
+) -> List[str]:
+    if sort:
+        sorted_idx = entropy.argsort(stable=True)
+        tokens = tokens[sorted_idx]
+    return [
+        separator.join([vocab_keys[i] for i in j])
+        .split("<start>" + separator)[-1]
+        .split(separator + "<end>")[0]
+        .replace("<pad>", "")
+        for j in tokens
+    ]
+
+
 @torch.no_grad()
 def test(
     model: ChemBFN,
@@ -86,6 +139,12 @@ def test(
         predict_y.append(y_hat.detach().to("cpu"))
     predict_y, label_y = torch.cat(predict_y, 0), torch.cat(label_y, 0).split(1, -1)
     if mode == "regression":
+        from sklearn.metrics import (
+            r2_score,
+            mean_absolute_error,
+            root_mean_squared_error,
+        )
+
         predict_y = [
             predict[label_y[i] != torch.inf]
             for (i, predict) in enumerate(predict_y.split(1, -1))
@@ -99,6 +158,8 @@ def test(
         r2 = [r2_score(label, predict) for (label, predict) in y_zipped]
         return {"MAE": mae, "RMSE": rmse, "R^2": r2}
     if mode == "classification":
+        from sklearn.metrics import roc_auc_score, auc, precision_recall_curve
+
         n_c = len(label_y)
         predict_y = predict_y.chunk(n_c, -1)
         y_zipped = list(zip(label_y, predict_y))
@@ -142,7 +203,6 @@ def split_dataset(
     assert file.endswith(".csv")
     assert len(split_ratio) == 3
     assert method in ("random", "scaffold")
-    colorama.just_fix_windows_console()
     with open(file, "r") as f:
         data = list(csv.reader(f))
     header = data[0]
@@ -167,10 +227,8 @@ def split_dataset(
             # compute Bemis-Murcko scaffold
             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",
+                    " Only the first SMILES will be used to compute the molecular scaffold.",
                     stacklevel=2,
                 )
             try:
@@ -197,10 +255,10 @@ def split_dataset(
     with open(file.replace(".csv", "_test.csv"), "w", newline="") as fte:
         writer = csv.writer(fte)
         writer.writerows([header] + test_set)
-    with open(file.replace(".csv", "_val.csv"), "w", newline="") as fva:
-        writer = csv.writer(fva)
-        writer.writerows([header] + val_set)
-    colorama.deinit()
+    if val_set:
+        with open(file.replace(".csv", "_val.csv"), "w", newline="") as fva:
+            writer = csv.writer(fva)
+            writer.writerows([header] + val_set)
 
 
 @torch.no_grad()
@@ -219,7 +277,7 @@ def sample(
     sort: bool = False,
 ) -> List[str]:
     """
-    Sampling.
+    Sampling molecules.
 
     :param model: trained ChemBFN model
     :param batch_size: batch size
@@ -250,32 +308,12 @@ def sample(
     :return: a list of generated molecular strings
     :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()
+    tp = _parse_and_assert_param(model, y, method)
+    device = _find_device() if device is None else device
     model.to(device).eval()
-    if y is not None:
-        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)
-    else:
-        token_mask = None
-    if "ode" in method.lower():
-        tp = float(method.split(":")[-1])
-        assert tp > 0, "Sampling temperature should be higher than 0."
+    y = _map_to_device(y, device)
+    token_mask = _build_token_mask(allowed_tokens, vocab_keys, device)
+    if tp:
         tokens, entropy = model.ode_sample(
             batch_size, sequence_size, y, sample_step, guidance_strength, token_mask, tp
         )
@@ -283,16 +321,7 @@ def sample(
         tokens, entropy = model.sample(
             batch_size, sequence_size, y, sample_step, guidance_strength, token_mask
         )
-    if sort:
-        sorted_idx = entropy.argsort(stable=True)
-        tokens = tokens[sorted_idx]
-    return [
-        seperator.join([vocab_keys[i] for i in j])
-        .split("<start>" + seperator)[-1]
-        .split(seperator + "<end>")[0]
-        .replace("<pad>", "")
-        for j in tokens
-    ]
+    return _token_to_seq(tokens, entropy, vocab_keys, seperator, sort)
 
 
 @torch.no_grad()
@@ -339,33 +368,13 @@ def inpaint(
     :return: a list of generated molecular strings
     :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()
+    tp = _parse_and_assert_param(model, y, method)
+    device = _find_device() if device is None else device
     model.to(device).eval()
     x = x.to(device)
-    if y is not None:
-        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)
-    else:
-        token_mask = None
-    if "ode" in method.lower():
-        tp = float(method.split(":")[-1])
-        assert tp > 0, "Sampling temperature should be higher than 0."
+    y = _map_to_device(y, device)
+    token_mask = _build_token_mask(allowed_tokens, vocab_keys, device)
+    if tp:
         tokens, entropy = model.ode_inpaint(
             x, y, sample_step, guidance_strength, token_mask, tp
         )
@@ -373,16 +382,68 @@ def inpaint(
         tokens, entropy = model.inpaint(
             x, y, sample_step, guidance_strength, token_mask
         )
-    if sort:
-        sorted_idx = entropy.argsort(stable=True)
-        tokens = tokens[sorted_idx]
-    return [
-        separator.join([vocab_keys[i] for i in j])
-        .split("<start>" + separator)[-1]
-        .split(separator + "<end>")[0]
-        .replace("<pad>", "")
-        for j in tokens
-    ]
+    return _token_to_seq(tokens, entropy, vocab_keys, separator, sort)
+
+
+@torch.no_grad()
+def optimise(
+    model: Union[ChemBFN, EnsembleChemBFN],
+    x: Tensor,
+    sample_step: int = 100,
+    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,
+    separator: str = "",
+    method: str = "BFN",
+    allowed_tokens: Union[str, List[str]] = "all",
+    sort: bool = False,
+) -> List[str]:
+    """
+    Optimising template molecules (mol2mol).
+
+    :param model: trained ChemBFN model
+    :param x: categorical indices of template;  shape: (n_b, n_t)
+    :param sample_step: number of sampling steps
+    :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
+    :param separator: token separator; default is `""`
+    :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 | bayesianflow_for_chem.model.EnsembleChemBFN
+    :type x: torch.Tensor
+    :type sample_step: int
+    :type y: torch.Tensor | list | dict | None
+    :type guidance_strength: float
+    :type device: str | torch.device | None
+    :type vocab_keys: list
+    :type separator: str
+    :type method: str
+    :type allowed_tokens: str | list
+    :type sort: bool
+    :return: a list of generated molecular strings
+    :rtype: list
+    """
+    tp = _parse_and_assert_param(model, y, method)
+    device = _find_device() if device is None else device
+    model.to(device).eval()
+    x = x.to(device)
+    y = _map_to_device(y, device)
+    token_mask = _build_token_mask(allowed_tokens, vocab_keys, device)
+    if tp:
+        tokens, entropy = model.ode_optimise(
+            x, y, sample_step, guidance_strength, token_mask, tp
+        )
+    else:
+        tokens, entropy = model.optimise(
+            x, y, sample_step, guidance_strength, token_mask
+        )
+    return _token_to_seq(tokens, entropy, vocab_keys, separator, sort)
 
 
 def quantise_model_(model: ChemBFN) -> None:

bayesianflow_for_chem/train.py

@@ -8,7 +8,6 @@ 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
@@ -55,6 +54,8 @@ class Model(LightningModule):
         self.scorer = scorer
         self.save_hyperparameters(hparam, ignore=["model", "mlp", "scorer"])
         if model.lora_enabled:
+            from loralib import mark_only_lora_as_trainable
+
             mark_only_lora_as_trainable(self.model)
         self.use_scorer = self.scorer is not None
 
@@ -107,6 +108,8 @@ class Model(LightningModule):
         :rtype: None
         """
         if self.model.lora_enabled:
+            from loralib import lora_state_dict
+
             torch.save(
                 {
                     "lora_nn": lora_state_dict(self.model),
@@ -152,6 +155,8 @@ class Regressor(LightningModule):
         self.model.requires_grad_(not hparam["freeze"])
         self.save_hyperparameters(hparam, ignore=["model", "mlp"])
         if model.lora_enabled:
+            from loralib import mark_only_lora_as_trainable
+
             mark_only_lora_as_trainable(self.model)
         assert hparam["mode"] in ("regression", "classification")
 
@@ -231,6 +236,8 @@ class Regressor(LightningModule):
         )
         if not self.hparams.freeze:
             if self.model.lora_enabled:
+                from loralib import lora_state_dict
+
                 torch.save(
                     {
                         "lora_nn": lora_state_dict(self.model),

{bayesianflow_for_chem-2.0.5.dist-info → bayesianflow_for_chem-2.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 2.0.5
+Version: 2.2.1
 Summary: Bayesian flow network framework for Chemistry
 Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
 Author: Nianze A. Tao
@@ -54,12 +54,14 @@ This is the repository of the PyTorch implementation of ChemBFN model.
 
 [![PyPI](https://img.shields.io/pypi/v/bayesianflow-for-chem?color=ff69b4)](https://pypi.org/project/bayesianflow-for-chem/)
 ![pytest](https://github.com/Augus1999/bayesian-flow-network-for-chemistry/actions/workflows/pytest.yml/badge.svg)
+[![document](https://github.com/Augus1999/bayesian-flow-network-for-chemistry/actions/workflows/pages/pages-build-deployment/badge.svg)](https://augus1999.github.io/bayesian-flow-network-for-chemistry/)
 
 ## Features
 
 ChemBFN provides the state-of-the-art functionalities of
 * SMILES or SELFIES-based *de novo* molecule generation
 * Protein sequence *de novo* generation
+* Template optimisation (mol2mol)
 * Classifier-free guidance conditional generation (single or multi-objective optimisation)
 * Context-guided conditional generation (inpaint)
 * Outstanding out-of-distribution chemical space sampling
@@ -71,6 +73,7 @@ in an all-in-one-model style.
 
 ## News
 
+* [09/10/2025] A web app [`chembfn_webui`](https://github.com/Augus1999/ChemBFN-WebUI) for hosting ChemBFN models is available on [PyPI](https://pypi.org/project/chembfn-webui/).
 * [30/01/2025] The package `bayesianflow_for_chem` is available on [PyPI](https://pypi.org/project/bayesianflow-for-chem/).
 * [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).

bayesianflow_for_chem-2.2.1.dist-info/RECORD

@@ -0,0 +1,15 @@
+bayesianflow_for_chem/__init__.py,sha256=l6twtXTdnix9SC8oL0Aqr0dYJ1uBUmMtk2QbUaRrhkE,642
+bayesianflow_for_chem/cli.py,sha256=jZPBUVwOl4qnjUMb3Lavu6owv1RP9jGdrDlM7KDA99c,26600
+bayesianflow_for_chem/data.py,sha256=RM3YaZUkTJ4-RNLH37GwhM4sRddOOusrC996ZPAe1lI,6590
+bayesianflow_for_chem/model.py,sha256=M35G4u4mX4btl9vOK3Iqs6yOSuIKI_OoCTmLhmjbwNk,57559
+bayesianflow_for_chem/scorer.py,sha256=gQFUlkyxitch02ntqcRh1ZS8aondKLynW5U6NfTQTb4,4084
+bayesianflow_for_chem/spectra.py,sha256=Ba9ib1aDvTtDYbH3b4d-lIty3ZSQMu7jwehuV2KmhwA,1781
+bayesianflow_for_chem/tool.py,sha256=VQo6xShwwmJbBsC-xTRL5ysCI7OLoOpmwQ-A7AuAPI8,23587
+bayesianflow_for_chem/train.py,sha256=7AU0A-eZwzSYsLyIe3OxGTNWPnhGpHmVUaQLplV2Fn8,9886
+bayesianflow_for_chem/_data/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
+bayesianflow_for_chem-2.2.1.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+bayesianflow_for_chem-2.2.1.dist-info/METADATA,sha256=YKuXgdklakB9DH4cFpTy3EbYdmo8dz9BgQQ4bI9CVQs,6476
+bayesianflow_for_chem-2.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bayesianflow_for_chem-2.2.1.dist-info/entry_points.txt,sha256=N63RMoJsr8rxuKxc7Fj802SL8J5AlpCoPkS8E3IFPLI,54
+bayesianflow_for_chem-2.2.1.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
+bayesianflow_for_chem-2.2.1.dist-info/RECORD,,

bayesianflow_for_chem-2.0.5.dist-info/RECORD

@@ -1,15 +0,0 @@
-bayesianflow_for_chem/__init__.py,sha256=U8ExDm5IRa-OPOLgt8VfjMDAiCmARiJonIKnZ8AVDQ0,612
-bayesianflow_for_chem/cli.py,sha256=g5pjXymycpchlgRE6SKr0LjTBjUl6MFHnFdQiKjcE3Q,22803
-bayesianflow_for_chem/data.py,sha256=Pl0gGWHmMKTKHpsxznvLgYPCwwlLNL7nqH19Vipjkxs,6584
-bayesianflow_for_chem/model.py,sha256=QF15BLpUjEpUCneTOHoU6MswvxArPfbFMiOHjwJ9JrM,52230
-bayesianflow_for_chem/scorer.py,sha256=gQFUlkyxitch02ntqcRh1ZS8aondKLynW5U6NfTQTb4,4084
-bayesianflow_for_chem/spectra.py,sha256=Ba9ib1aDvTtDYbH3b4d-lIty3ZSQMu7jwehuV2KmhwA,1781
-bayesianflow_for_chem/tool.py,sha256=JSlKrar1vVouTRUJBYU6lc5AQmAIaexvpUA3W3oQcKs,21284
-bayesianflow_for_chem/train.py,sha256=jYkhSguW50lrcTEydCQ20yig_mmc1j7WH9KmVwBCTAo,9727
-bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
-bayesianflow_for_chem-2.0.5.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-bayesianflow_for_chem-2.0.5.dist-info/METADATA,sha256=qCd0LnD66eQbWlTvKDpHZHtJKmy9NQogm_k9_amEt2s,6057
-bayesianflow_for_chem-2.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bayesianflow_for_chem-2.0.5.dist-info/entry_points.txt,sha256=N63RMoJsr8rxuKxc7Fj802SL8J5AlpCoPkS8E3IFPLI,54
-bayesianflow_for_chem-2.0.5.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
-bayesianflow_for_chem-2.0.5.dist-info/RECORD,,