bayesianflow-for-chem 2.0.2__py3-none-any.whl → 2.0.3__py3-none-any.whl

bayesianflow_for_chem/__init__.py

@@ -3,6 +3,7 @@
 """
 ChemBFN package.
 """
+import colorama
 from . import data, tool, train, scorer, spectra
 from .model import ChemBFN, MLP, EnsembleChemBFN
 from .cli import main_script
@@ -17,7 +18,7 @@ __all__ = [
     "MLP",
     "EnsembleChemBFN",
 ]
-__version__ = "2.0.2"
+__version__ = "2.0.3"
 __author__ = "Nianze A. Tao (Omozawa Sueno)"
 
 
@@ -28,4 +29,6 @@ def main() -> None:
     :return:
     :rtype: None
     """
+    colorama.just_fix_windows_console()
     main_script(__version__)
+    colorama.deinit()

bayesianflow_for_chem/cli.py

@@ -180,14 +180,16 @@ def load_model_config(
         model_config = tomllib.load(f)
     if model_config["ChemBFN"]["num_vocab"] != "match vocabulary size":
         if not isinstance(model_config["ChemBFN"]["num_vocab"], int):
-            print(f"Critical in {config_file}: You must specify num_vocab.")
+            print(
+                f"\033[0;31mCritical\033[0;0m in {config_file}: You must specify num_vocab."
+            )
             flag_critical += 1
     if model_config["ChemBFN"]["base_model"]:
         model_file = model_config["ChemBFN"]["base_model"]
         for fn in model_file:
             if not os.path.exists(fn):
                 print(
-                    f"Critical in {config_file}: Base model file {fn} does not exist."
+                    f"\033[0;31mCritical\033[0;0m in {config_file}: Base model file {fn} does not exist."
                 )
                 flag_critical += 1
     if "MLP" in model_config:
@@ -195,14 +197,14 @@ def load_model_config(
         b = model_config["MLP"]["size"][-1]
         if a != b:
             print(
-                f"Critical in {config_file}: MLP hidden size {b} should match ChemBFN hidden size {a}."
+                f"\033[0;31mCritical\033[0;0m in {config_file}: MLP hidden size {b} should match ChemBFN hidden size {a}."
             )
             flag_critical += 1
         if model_config["MLP"]["base_model"]:
             model_file = model_config["MLP"]["base_model"]
             if not os.path.exists(model_file):
                 print(
-                    f"Critical in {config_file}: Base model file {fn} does not exist."
+                    f"\033[0;31mCritical\033[0;0m in {config_file}: Base model file {fn} does not exist."
                 )
                 flag_critical += 1
     return model_config, flag_critical, flag_warning
@@ -226,49 +228,61 @@ def load_runtime_config(
         config = tomllib.load(f)
     tokeniser_name = config["tokeniser"]["name"].lower()
     if not tokeniser_name in "smiles selfies safe fasta".split():
-        print(f"Critical in {config_file}: Unknown tokensier name: {tokeniser_name}.")
+        print(
+            f"\033[0;31mCritical\033[0;0m in {config_file}: Unknown tokensier name: {tokeniser_name}."
+        )
         flag_critical += 1
     if tokeniser_name == "selfies":
         vocab = config["tokeniser"]["vocab"]
         if vocab.lower() == "default":
-            print(f"Critical in {config_file}: You should specify a vocabulary file.")
+            print(
+                f"\033[0;31mCritical\033[0;0m in {config_file}: You should specify a vocabulary file."
+            )
             flag_critical += 1
         elif not os.path.exists(vocab):
-            print(f"Critical in {config_file}: Vocabulary file {vocab} does not exist.")
+            print(
+                f"\033[0;31mCritical\033[0;0m in {config_file}: Vocabulary file {vocab} does not exist."
+            )
             flag_critical += 1
     if "train" in config:
         dataset_file = config["train"]["dataset"]
         if not os.path.exists(dataset_file):
             print(
-                f"Critical in {config_file}: Dataset file {dataset_file} does not exist."
+                f"\033[0;31mCritical\033[0;0m in {config_file}: Dataset file {dataset_file} does not exist."
             )
             flag_critical += 1
         logger_name = config["train"]["logger_name"].lower()
         if not logger_name in "csv tensorboard wandb".split():
-            print(f"Critical in {config_file}: Unknown logger: {logger_name}.")
+            print(
+                f"\033[0;31mCritical\033[0;0m in {config_file}: Unknown logger: {logger_name}."
+            )
             flag_critical += 1
         if config["train"]["restart"]:
             ckpt_file = config["train"]["restart"]
             if not os.path.exists(ckpt_file):
                 print(
-                    f"Critical in {config_file}: Restart checkpoint file {ckpt_file} does not exist."
+                    f"\033[0;31mCritical\033[0;0m in {config_file}: Restart checkpoint file {ckpt_file} does not exist."
                 )
                 flag_critical += 1
     if "inference" in config:
         if not "train" in config:
             if not isinstance(config["inference"]["sequence_length"], int):
                 print(
-                    f"Critical in {config_file}: You must set an integer for sequence_length."
+                    f"\033[0;31mCritical\033[0;0m in {config_file}: You must set an integer for sequence_length."
                 )
                 flag_critical += 1
         if config["inference"]["guidance_objective"]:
             if not "guidance_objective_strength" in config["inference"]:
                 print(
-                    f"Critical in {config_file}: You need to add guidance_objective_strength."
+                    f"\033[0;31mCritical\033[0;0m in {config_file}: You need to add guidance_objective_strength."
                 )
                 flag_critical += 1
         result_dir = Path(config["inference"]["result_file"]).parent
-        assert os.path.exists(result_dir), f"directory {result_dir} does not exist."
+        if not os.path.exists(result_dir):
+            print(
+                f"\033[0;33mWarning\033[0;0m in {config_file}: Directory {result_dir} to save the result does not exist."
+            )
+            flag_warning += 1
     return config, flag_critical, flag_warning
 
 
@@ -306,7 +320,7 @@ def main_script(version: str) -> None:
         if runtime_config["train"]["enable_lora"]:
             if not model_config["ChemBFN"]["base_model"]:
                 print(
-                    f"Warning in {parser.model_config}: You should load a pretrained model first."
+                    f"\033[0;33mWarning\033[0;0m in {parser.model_config}: You should load a pretrained model first."
                 )
                 flag_warning += 1
         if not os.path.exists(runtime_config["train"]["checkpoint_save_path"]):
@@ -314,12 +328,12 @@ def main_script(version: str) -> None:
     else:
         if not model_config["ChemBFN"]["base_model"]:
             print(
-                f"Warning in {parser.model_config}: You should load a pretrained ChemBFN model."
+                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"]:
             print(
-                f"Warning in {parser.model_config}: You should load a pretrained MLP."
+                f"\033[0;33mWarning\033[0;0m in {parser.model_config}: You should load a pretrained MLP."
             )
             flag_warning += 1
     if "inference" in runtime_config:

bayesianflow_for_chem/tool.py

@@ -9,6 +9,7 @@ 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
@@ -141,6 +142,7 @@ 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]
@@ -198,6 +200,7 @@ def split_dataset(
     with open(file.replace(".csv", "_val.csv"), "w", newline="") as fva:
         writer = csv.writer(fva)
         writer.writerows([header] + val_set)
+    colorama.deinit()
 
 
 @torch.no_grad()

{bayesianflow_for_chem-2.0.2.dist-info → bayesianflow_for_chem-2.0.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 2.0.2
+Version: 2.0.3
 Summary: Bayesian flow network framework for Chemistry
 Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
 Author: Nianze A. Tao
@@ -23,6 +23,7 @@ License-File: LICENSE
 Requires-Dist: rdkit>=2025.3.5
 Requires-Dist: torch>=2.8.0
 Requires-Dist: torchao>=0.12
+Requires-Dist: colorama>=0.4.6
 Requires-Dist: numpy>=2.3.2
 Requires-Dist: scipy>=1.16.1
 Requires-Dist: loralib>=0.1.2
@@ -49,6 +50,11 @@ Dynamic: summary
 
 This is the repository of the PyTorch implementation of ChemBFN model.
 
+## Build State
+
+[![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)
+
 ## Features
 
 ChemBFN provides the state-of-the-art functionalities of

{bayesianflow_for_chem-2.0.2.dist-info → bayesianflow_for_chem-2.0.3.dist-info}/RECORD

@@ -1,15 +1,15 @@
-bayesianflow_for_chem/__init__.py,sha256=ZFYmYKG9x1nn1ootKEjcAM1e5_8OYhmJ2iajT3Zq1rw,534
-bayesianflow_for_chem/cli.py,sha256=iies_Dr_rdwfXw5wcNmX2T36eo1UlvAtIkkZVYlKASQ,22125
+bayesianflow_for_chem/__init__.py,sha256=D4bOlPLzpGhxTdr4Zb96Xtoc5y5BftiiuQT-SL86a7s,612
+bayesianflow_for_chem/cli.py,sha256=decINvgtoDNUZy-wJvUE41ggUb-7RKZZh3Vbu0v_nSk,22695
 bayesianflow_for_chem/data.py,sha256=Pl0gGWHmMKTKHpsxznvLgYPCwwlLNL7nqH19Vipjkxs,6584
 bayesianflow_for_chem/model.py,sha256=bswVv3DiQTF3u37A80lrj_UPzklYtRH06woFfMXy84k,51643
 bayesianflow_for_chem/scorer.py,sha256=gQFUlkyxitch02ntqcRh1ZS8aondKLynW5U6NfTQTb4,4084
 bayesianflow_for_chem/spectra.py,sha256=Ba9ib1aDvTtDYbH3b4d-lIty3ZSQMu7jwehuV2KmhwA,1781
-bayesianflow_for_chem/tool.py,sha256=-ldRAvGLJndsgGY_anHMrIw19ePk2dVcx5qt2DGWI3s,21211
+bayesianflow_for_chem/tool.py,sha256=8fNe8pL5veUHHnnqV02AQXDe76miKRrZzypjb9Nu_hA,21289
 bayesianflow_for_chem/train.py,sha256=hGKyhGhLch-exSYPZdLXrLn3gf39Q1VLSJs2qtuikQE,9709
 bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
-bayesianflow_for_chem-2.0.2.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-bayesianflow_for_chem-2.0.2.dist-info/METADATA,sha256=2uEwTlHBiFBKxeS4QOBGTYw6VKgnpp-TZCuiQ88itq4,5762
-bayesianflow_for_chem-2.0.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bayesianflow_for_chem-2.0.2.dist-info/entry_points.txt,sha256=N63RMoJsr8rxuKxc7Fj802SL8J5AlpCoPkS8E3IFPLI,54
-bayesianflow_for_chem-2.0.2.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
-bayesianflow_for_chem-2.0.2.dist-info/RECORD,,
+bayesianflow_for_chem-2.0.3.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+bayesianflow_for_chem-2.0.3.dist-info/METADATA,sha256=eR9Kn7lM9xsALxC5y2cThQQ5iiGX-atOl_sBaLdg_NM,6056
+bayesianflow_for_chem-2.0.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bayesianflow_for_chem-2.0.3.dist-info/entry_points.txt,sha256=N63RMoJsr8rxuKxc7Fj802SL8J5AlpCoPkS8E3IFPLI,54
+bayesianflow_for_chem-2.0.3.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
+bayesianflow_for_chem-2.0.3.dist-info/RECORD,,