From ae6b9b33c286bdf0ece9a6dc1d2f9b1f8f15773e Mon Sep 17 00:00:00 2001 From: sfluegel Date: Fri, 24 Jul 2026 14:18:37 +0200 Subject: [PATCH 1/2] add compression toproperty files --- chebai_graph/preprocessing/datasets/chebi.py | 11 +- .../preprocessing/property_encoder.py | 104 ++++++++++++++++++ 2 files changed, 114 insertions(+), 1 deletion(-) diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index a83b58e..d7cf10a 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -168,7 +168,12 @@ def enc_if_not_none(encode, value): assert len(encoded_values) == len(idents) == len(features) torch.save( [ - {property.name: torch.cat(feat), "ident": id} + { + property.name: property.encoder.compress( + torch.cat(feat) + ), + "ident": id, + } for feat, id in zip(encoded_values, idents) if feat is not None ], @@ -384,6 +389,8 @@ def load_processed_data( property_data = torch.load( self.get_property_path(property), weights_only=False ) + for entry in property_data: + entry[property.name] = property.encoder.decompress(entry[property.name]) if len(property_data[0][property.name].shape) > 1: property.encoder.set_encoding_length( property_data[0][property.name].shape[1] @@ -535,6 +542,8 @@ def load_processed_data( property_data = torch.load( self.get_property_path(property), weights_only=False ) + for entry in property_data: + entry[property.name] = property.encoder.decompress(entry[property.name]) if len(property_data[0][property.name].shape) > 1: property.encoder.set_encoding_length( property_data[0][property.name].shape[1] diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 3edd2d2..7592661 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -46,6 +46,39 @@ def encode(self, value) -> torch.Tensor: """ return value + def compress(self, tensor: torch.Tensor) -> torch.Tensor: + """ + Compress an encoded tensor into a more compact on-disk representation. + + Called just before caching property values to disk. The default + implementation is a no-op; subclasses override it to reduce file size + (e.g. by downcasting the dtype or storing indices instead of one-hot + vectors). Must be losslessly invertible by :meth:`decompress` (except + for deliberate float precision reductions). + + Args: + tensor: The encoded property tensor for a single molecule. + + Returns: + A compact tensor to store on disk. + """ + return tensor + + def decompress(self, tensor: torch.Tensor) -> torch.Tensor: + """ + Reconstruct the full encoded tensor from its compressed on-disk form. + + Inverse of :meth:`compress`, called right after loading cached property + values. The default implementation is a no-op. + + Args: + tensor: The compressed property tensor as loaded from disk. + + Returns: + The reconstructed encoded property tensor. + """ + return tensor + def on_start(self, **kwargs) -> None: """Hook called at the start of encoding process.""" pass @@ -238,6 +271,61 @@ def encode(self, token: str | None) -> torch.Tensor: self.tokens_dict[token], num_classes=self.get_encoding_length() ) + def compress(self, tensor: torch.Tensor) -> torch.Tensor: + """ + Store one index per node instead of the full one-hot matrix. + + A dense ``(N, n_classes)`` int64 one-hot matrix is reduced to an + ``(N,)`` vector of class indices. Index ``0`` is reserved for all-zero + rows (produced by :meth:`encode` for unknown tokens); real classes are + stored as ``argmax + 1``. The result uses ``uint8`` when it fits, else + ``int16``. + + Args: + tensor: One-hot tensor of shape ``(N, n_classes)``. + + Returns: + Index tensor of shape ``(N,)``. + """ + if tensor.dim() != 2: + # already compressed / unexpected shape - leave untouched + return tensor + has_class = tensor.any(dim=1) + indices = torch.where( + has_class, + tensor.argmax(dim=1) + 1, + torch.zeros_like(has_class, dtype=torch.long), + ) + dtype = torch.uint8 if tensor.shape[1] + 1 < 256 else torch.int16 + return indices.to(dtype) + + def decompress(self, tensor: torch.Tensor) -> torch.Tensor: + """ + Reconstruct the dense one-hot matrix from stored class indices. + + Inverse of :meth:`compress`. Index ``0`` maps back to an all-zero row; + index ``i > 0`` maps to a one-hot with class ``i - 1`` set. + + Args: + tensor: Index tensor of shape ``(N,)`` as produced by + :meth:`compress`. + + Returns: + One-hot tensor of shape ``(N, n_classes)`` with ``int64`` dtype. + """ + if tensor.dim() != 1: + # already expanded / unexpected shape - leave untouched + return tensor + n_classes = self.get_encoding_length() + indices = tensor.to(torch.int64) + out = torch.zeros((indices.shape[0], n_classes), dtype=torch.int64) + non_zero = indices > 0 + if non_zero.any(): + out[non_zero] = torch.nn.functional.one_hot( + indices[non_zero] - 1, num_classes=n_classes + ) + return out + class AsIsEncoder(PropertyEncoder): """ @@ -271,6 +359,14 @@ def encode(self, token: float | int | None) -> torch.Tensor: # ----- fix: for above warning return torch.tensor(token).unsqueeze(0) # shape: (1, len(token)) + def compress(self, tensor: torch.Tensor) -> torch.Tensor: + """Downcast float values to float32 to halve the on-disk size.""" + if tensor.is_floating_point(): + return tensor.to(torch.float32) + return tensor + + # decompress is a no-op: float32 values are used as-is at load time. + class BoolEncoder(PropertyEncoder): """ @@ -293,3 +389,11 @@ def encode(self, token: bool) -> torch.Tensor: Tensor with 1 if True else 0. """ return torch.tensor([1 if token else 0]) + + def compress(self, tensor: torch.Tensor) -> torch.Tensor: + """Store the 0/1 values as ``uint8`` instead of ``int64`` (8x smaller).""" + return tensor.to(torch.uint8) + + def decompress(self, tensor: torch.Tensor) -> torch.Tensor: + """Restore the original ``int64`` dtype of the boolean encoding.""" + return tensor.to(torch.int64) From 2ea3170c2b3c09fcc0dfa91fb49cc7401cf56368 Mon Sep 17 00:00:00 2001 From: sfluegel Date: Fri, 24 Jul 2026 15:21:23 +0200 Subject: [PATCH 2/2] don't decompress what doesn't need decrompessing --- .../preprocessing/property_encoder.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 7592661..a84a453 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -311,19 +311,20 @@ def decompress(self, tensor: torch.Tensor) -> torch.Tensor: :meth:`compress`. Returns: - One-hot tensor of shape ``(N, n_classes)`` with ``int64`` dtype. + One-hot tensor of shape ``(N, n_classes)`` keeping the compact + stored dtype (e.g. ``uint8``); it is promoted to float when merged + into the node/edge feature matrix at load time. """ if tensor.dim() != 1: # already expanded / unexpected shape - leave untouched return tensor n_classes = self.get_encoding_length() - indices = tensor.to(torch.int64) - out = torch.zeros((indices.shape[0], n_classes), dtype=torch.int64) - non_zero = indices > 0 + out = torch.zeros((tensor.shape[0], n_classes), dtype=tensor.dtype) + non_zero = tensor > 0 if non_zero.any(): out[non_zero] = torch.nn.functional.one_hot( - indices[non_zero] - 1, num_classes=n_classes - ) + tensor[non_zero].to(torch.int64) - 1, num_classes=n_classes + ).to(out.dtype) return out @@ -365,8 +366,6 @@ def compress(self, tensor: torch.Tensor) -> torch.Tensor: return tensor.to(torch.float32) return tensor - # decompress is a no-op: float32 values are used as-is at load time. - class BoolEncoder(PropertyEncoder): """ @@ -393,7 +392,3 @@ def encode(self, token: bool) -> torch.Tensor: def compress(self, tensor: torch.Tensor) -> torch.Tensor: """Store the 0/1 values as ``uint8`` instead of ``int64`` (8x smaller).""" return tensor.to(torch.uint8) - - def decompress(self, tensor: torch.Tensor) -> torch.Tensor: - """Restore the original ``int64`` dtype of the boolean encoding.""" - return tensor.to(torch.int64)