bayesianflow-for-chem 1.4.1__py3-none-any.whl → 1.4.2__py3-none-any.whl

bayesianflow_for_chem/__init__.py

@@ -7,5 +7,5 @@ from . import data, tool, train, scorer
 from .model import ChemBFN, MLP, EnsembleChemBFN
 
 __all__ = ["data", "tool", "train", "scorer", "ChemBFN", "MLP", "EnsembleChemBFN"]
-__version__ = "1.4.1"
+__version__ = "1.4.2"
 __author__ = "Nianze A. Tao (Omozawa Sueno)"

bayesianflow_for_chem/model.py

@@ -54,9 +54,19 @@ class Linear(nn.Linear):
         :return:
         :rtype: None
         """
+        from torchao.dtypes.affine_quantized_tensor import AffineQuantizedTensor
+
         assert r > 0, "Rank should be larger than 0."
-        self.lora_A = nn.Parameter(self.weight.new_zeros((r, self.in_features)))
-        self.lora_B = nn.Parameter(self.weight.new_zeros((self.out_features, r)))
+        if isinstance(self.weight, AffineQuantizedTensor):
+            self.lora_A = nn.Parameter(
+                torch.zeros((r, self.in_features), device=self.weight.device)
+            )
+            self.lora_B = nn.Parameter(
+                torch.zeros((self.out_features, r), device=self.weight.device)
+            )
+        else:
+            self.lora_A = nn.Parameter(self.weight.new_zeros((r, self.in_features)))
+            self.lora_B = nn.Parameter(self.weight.new_zeros((self.out_features, r)))
         self.scaling = lora_alpha / r
         self.lora_dropout = lora_dropout
         self.lora_enabled = True
@@ -1207,23 +1217,23 @@ class EnsembleChemBFN(ChemBFN):
         )
 
     def quantise(
-        self, quantise_method: Optional[Callable[[ChemBFN], nn.Module]] = None
+        self, quantise_method: Optional[Callable[[ChemBFN], None]] = 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`
+        :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
+            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)
+            quantise_method = quantise_model_
+        for _, v in self.models.items():
+            quantise_method(v)
 
     def jit(self, freeze: bool = False) -> None:
         """

bayesianflow_for_chem/tool.py

@@ -13,9 +13,8 @@ import torch
 import numpy as np
 import torch.nn as nn
 from torch import cuda, Tensor, softmax
-from torch.ao import quantization
 from torch.utils.data import DataLoader
-from typing_extensions import Self
+from typing_extensions import Self, deprecated
 from rdkit.Chem import (
     rdDetermineBonds,
     GetFormalCharge,
@@ -386,6 +385,11 @@ def inpaint(
     ]
 
 
+@deprecated(
+    "Eager mode quantization from `torch.ao` is deprecated and will be remove in version 2.10, "
+    "so this fuction will stop working since that time. "
+    "Please use `quantise_model_` instead."
+)
 def quantise_model(model: ChemBFN) -> nn.Module:
     """
     Dynamic quantisation of the trained model to `torch.qint8` data type.
@@ -395,6 +399,7 @@ def quantise_model(model: ChemBFN) -> nn.Module:
     :return: quantised model
     :rtype: torch.nn.Module
     """
+    from torch.ao import quantization
     from torch.ao.nn.quantized import dynamic
     from torch.ao.nn.quantized.modules.utils import _quantize_weight
     from torch.ao.quantization.qconfig import default_dynamic_qconfig
@@ -527,6 +532,24 @@ def quantise_model(model: ChemBFN) -> nn.Module:
     return quantised_model
 
 
+def quantise_model_(model: ChemBFN) -> None:
+    """
+    In-place dynamic quantisation of the trained model to `int8` data type. \n
+    Due to some limitations of `torchao` module, it is slower than method previded by `torch.ao`.
+
+    :param model: trained ChemBFN model
+    :type model: bayesianflow_for_chem.model.ChemBFN
+    :return:
+    :rtype: None
+    """
+    from torchao.quantization.quant_api import (
+        quantize_,
+        Int8DynamicActivationInt8WeightConfig,
+    )
+
+    quantize_(model, Int8DynamicActivationInt8WeightConfig())
+
+
 class GeometryConverter:
     """
     Converting between different 2D/3D molecular representations.

{bayesianflow_for_chem-1.4.1.dist-info → bayesianflow_for_chem-1.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bayesianflow_for_chem
-Version: 1.4.1
+Version: 1.4.2
 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 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: rdkit>=2023.9.6
 Requires-Dist: torch>=2.3.1
+Requires-Dist: torchao>=0.12
 Requires-Dist: numpy>=1.26.4
 Requires-Dist: loralib>=0.1.2
 Requires-Dist: lightning>=2.2.0

bayesianflow_for_chem-1.4.2.dist-info/RECORD

@@ -0,0 +1,12 @@
+bayesianflow_for_chem/__init__.py,sha256=IeIasLe6wLuGbH7DIlB38ehDPqvlMBT388hf58I3J30,329
+bayesianflow_for_chem/data.py,sha256=WoOCOVmJX4WeHa2WeO4i66J2FS8rvRaYRCdlBN7ZeOM,6576
+bayesianflow_for_chem/model.py,sha256=6pxGuIM7rKyawcz2hI8dT88rv3qFsnCvlLhDj1CB9YU,50595
+bayesianflow_for_chem/scorer.py,sha256=7G1TVSwC0qONtNm6kiDZUWwvuFPzasNSjp4eJAk5TL0,4101
+bayesianflow_for_chem/tool.py,sha256=Ne_ew1P8r6KWOqUZpb-BL_q7Dm6fnSTtxhJvgV1JHHs,26264
+bayesianflow_for_chem/train.py,sha256=hGKyhGhLch-exSYPZdLXrLn3gf39Q1VLSJs2qtuikQE,9709
+bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
+bayesianflow_for_chem-1.4.2.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+bayesianflow_for_chem-1.4.2.dist-info/METADATA,sha256=s6k85HFXvasxvZBJD3Rj8cFNJXehS-utcMeKC6tP8F8,5673
+bayesianflow_for_chem-1.4.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bayesianflow_for_chem-1.4.2.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
+bayesianflow_for_chem-1.4.2.dist-info/RECORD,,

bayesianflow_for_chem-1.4.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-bayesianflow_for_chem/__init__.py,sha256=N_7P9Ea0eUmdC0wQKXIHiuMzPK4p9_cBF_YOexjo5yo,329
-bayesianflow_for_chem/data.py,sha256=WoOCOVmJX4WeHa2WeO4i66J2FS8rvRaYRCdlBN7ZeOM,6576
-bayesianflow_for_chem/model.py,sha256=zJkcUnZcxFa4iTo9_-BHzAM1MkJm1pbEGiczVgyUcxo,50186
-bayesianflow_for_chem/scorer.py,sha256=7G1TVSwC0qONtNm6kiDZUWwvuFPzasNSjp4eJAk5TL0,4101
-bayesianflow_for_chem/tool.py,sha256=Ma4dEBWP5GFKk-Tj5vBJgxGs_WAp4F87-b1UcsqUAn4,25486
-bayesianflow_for_chem/train.py,sha256=hGKyhGhLch-exSYPZdLXrLn3gf39Q1VLSJs2qtuikQE,9709
-bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
-bayesianflow_for_chem-1.4.1.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-bayesianflow_for_chem-1.4.1.dist-info/METADATA,sha256=460yUOjHG9PTavIddJJ2Ufdq0bkLBZqbmMugyq6LVPQ,5643
-bayesianflow_for_chem-1.4.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bayesianflow_for_chem-1.4.1.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
-bayesianflow_for_chem-1.4.1.dist-info/RECORD,,