bayesianflow-for-chem 1.4.1__tar.gz → 1.4.2__tar.gz

{bayesianflow_for_chem-1.4.1 → bayesianflow_for_chem-1.4.2}/PKG-INFO

@@ -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.1 → bayesianflow_for_chem-1.4.2}/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-1.4.1 → bayesianflow_for_chem-1.4.2}/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-1.4.1 → bayesianflow_for_chem-1.4.2}/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 → bayesianflow_for_chem-1.4.2}/bayesianflow_for_chem.egg-info/PKG-INFO

@@ -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.1 → bayesianflow_for_chem-1.4.2}/bayesianflow_for_chem.egg-info/requires.txt

@@ -1,5 +1,6 @@
 rdkit>=2023.9.6
 torch>=2.3.1
+torchao>=0.12
 numpy>=1.26.4
 loralib>=0.1.2
 lightning>=2.2.0

{bayesianflow_for_chem-1.4.1 → bayesianflow_for_chem-1.4.2}/setup.py

@@ -40,6 +40,7 @@ setup(
     install_requires=[
         "rdkit>=2023.9.6",
         "torch>=2.3.1",
+        "torchao>=0.12",
         "numpy>=1.26.4",
         "loralib>=0.1.2",
         "lightning>=2.2.0",