bayesianflow-for-chem 1.2.4__py3-none-any.whl → 1.2.6__py3-none-any.whl

bayesianflow_for_chem/__init__.py

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

bayesianflow_for_chem/tool.py

@@ -492,15 +492,16 @@ def inpaint(
 
 def quantise_model(model: ChemBFN) -> nn.Module:
     """
-    Dynamic quantisation of the trained model.
+    Dynamic quantisation of the trained model to `torch.qint8` data type.
 
     :param model: trained ChemBFN model
     :type model: bayesianflow_for_chem.model.ChemBFN
     :return: quantised model
     :rtype: torch.nn.Module
     """
-    from torch.ao.nn.quantized.modules.utils import _quantize_weight
     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
 
     class QuantisedLinear(dynamic.Linear):
         # Modified from https://github.com/pytorch/pytorch/blob/main/torch/ao/nn/quantized/dynamic/modules/linear.py
@@ -521,6 +522,9 @@ def quantise_model(model: ChemBFN) -> nn.Module:
             self.scaling: Optional[float] = None
             self.lora_dropout: Optional[float] = None
 
+        def _get_name(self) -> str:
+            return "DynamicQuantizedLoRALinear"
+
         def enable_lora(
             self, r: int = 8, lora_alpha: int = 1, lora_dropout: float = 0.0
         ) -> None:
@@ -540,7 +544,6 @@ def quantise_model(model: ChemBFN) -> nn.Module:
             self._packed_params.requires_grad_(False)
 
         def forward(self, x: Tensor) -> Tensor:
-            # Note that we can handle self.bias == None case.
             if self._packed_params.dtype == torch.qint8:
                 if self.version is None or self.version < 4:
                     Y = torch.ops.quantized.linear_dynamic(
@@ -575,11 +578,6 @@ def quantise_model(model: ChemBFN) -> nn.Module:
             if mod.qconfig is not None and mod.qconfig.weight is not None:
                 weight_observer = mod.qconfig.weight()
             else:
-                # We have the circular import issues if we import the qconfig in the beginning of this file:
-                # https://github.com/pytorch/pytorch/pull/24231. The current workaround is to postpone the
-                # import until we need it.
-                from torch.ao.quantization.qconfig import default_dynamic_qconfig
-
                 weight_observer = default_dynamic_qconfig.weight()
             dtype = weight_observer.dtype
             assert dtype in [torch.qint8, torch.float16], (
@@ -599,10 +597,28 @@ def quantise_model(model: ChemBFN) -> nn.Module:
             qlinear.set_weight_bias(qweight, mod.bias)
             if mod.lora_enabled:
                 qlinear.lora_enabled = True
-                qlinear.lora_A = mod.lora_A
-                qlinear.lora_B = mod.lora_B
-                qlinear.scaling = mod.scaling
-                qlinear.lora_dropout = mod.lora_dropout
+                qlinear.lora_A = nn.Parameter(mod.lora_A.clone().detach_())
+                qlinear.lora_B = nn.Parameter(mod.lora_B.clone().detach_())
+                qlinear.scaling = deepcopy(mod.scaling)
+                qlinear.lora_dropout = deepcopy(mod.lora_dropout)
+            return qlinear
+
+        @classmethod
+        def from_reference(cls, ref_qlinear: Self) -> Self:
+            qlinear = cls(
+                ref_qlinear.in_features,
+                ref_qlinear.out_features,
+                dtype=ref_qlinear.weight_dtype,
+            )
+            qweight = ref_qlinear.get_quantized_weight()
+            bias = ref_qlinear.bias
+            qlinear.set_weight_bias(qweight, bias)
+            if ref_qlinear.lora_enabled:
+                qlinear.lora_enabled = True
+                qlinear.lora_A = nn.Parameter(ref_qlinear.lora_A.clone().detach_())
+                qlinear.lora_B = nn.Parameter(ref_qlinear.lora_B.clone().detach_())
+                qlinear.scaling = deepcopy(ref_qlinear.scaling)
+                qlinear.lora_dropout = deepcopy(ref_qlinear.lora_dropout)
             return qlinear
 
     mapping = deepcopy(quantization.DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS)

{bayesianflow_for_chem-1.2.4.dist-info → bayesianflow_for_chem-1.2.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: bayesianflow_for_chem
-Version: 1.2.4
+Version: 1.2.6
 Summary: Bayesian flow network framework for Chemistry
 Home-page: https://augus1999.github.io/bayesian-flow-network-for-chemistry/
 Author: Nianze A. Tao

{bayesianflow_for_chem-1.2.4.dist-info → bayesianflow_for_chem-1.2.6.dist-info}/RECORD

@@ -1,12 +1,12 @@
-bayesianflow_for_chem/__init__.py,sha256=-_0xD4lo_Vn2GrlXG-y13MCTwDfj391kzgTnyLplkNk,293
+bayesianflow_for_chem/__init__.py,sha256=sdyCK-Zd32-FNOcjuSB02ABx8vn53phorQeVqyWMWk4,293
 bayesianflow_for_chem/data.py,sha256=9tpRba40lxwrB6aPSJMkxUglEVC3VEQC9wWxhDuz3Q8,7760
 bayesianflow_for_chem/model.py,sha256=HvEvW_xRbkv4eSv5lhd72BJMZkg-ZACEi1DAW3p5Q1Y,35918
 bayesianflow_for_chem/scorer.py,sha256=mV1vX8aBGFra2BE7N8WHihVIo3dXmUdPQIGfSaiuNdk,4084
-bayesianflow_for_chem/tool.py,sha256=d-g47Ctn6qb_j1bWCWV99ytUxJ23zJ32SJacQ_WXONk,23028
+bayesianflow_for_chem/tool.py,sha256=VuEqbT7Qraa4vnKMHbToyAYIiRoQI7gEPLKEBCWGmVg,23706
 bayesianflow_for_chem/train.py,sha256=kj6icGqymUUYopDtpre1oE_wpvpeNilbpzgffBsd1tk,9589
 bayesianflow_for_chem/vocab.txt,sha256=HgtAZmpWYk4y8PqEVC4vqut1vE75DfRKE_10s2UW0rU,790
-bayesianflow_for_chem-1.2.4.dist-info/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-bayesianflow_for_chem-1.2.4.dist-info/METADATA,sha256=78FGoGjMsdwBavH4rSDtQ_psRYLSUdcg6cdR7KRmgVQ,5890
-bayesianflow_for_chem-1.2.4.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-bayesianflow_for_chem-1.2.4.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
-bayesianflow_for_chem-1.2.4.dist-info/RECORD,,
+bayesianflow_for_chem-1.2.6.dist-info/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+bayesianflow_for_chem-1.2.6.dist-info/METADATA,sha256=Akoh5dQW_0jeYuGC4ZKKYHS1WJn0xRwGDr7ut-Q-5sc,5890
+bayesianflow_for_chem-1.2.6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+bayesianflow_for_chem-1.2.6.dist-info/top_level.txt,sha256=KHsanI3BMCt8D9Qpze2ycrF6nMa3PyojgO6eS1c8kco,22
+bayesianflow_for_chem-1.2.6.dist-info/RECORD,,