From ae6b9b33c286bdf0ece9a6dc1d2f9b1f8f15773e Mon Sep 17 00:00:00 2001 From: sfluegel Date: Fri, 24 Jul 2026 14:18:37 +0200 Subject: [PATCH 1/4] 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/4] 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) From 58666ca2824fd309289d8ac47284b26a1209df1f Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 10:16:39 +0200 Subject: [PATCH 3/4] always convert indices to in16 (assumes that we usually have more than 255 and always less than 32768 classes) --- chebai_graph/preprocessing/property_encoder.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index a84a453..38cb279 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -296,8 +296,7 @@ def compress(self, tensor: torch.Tensor) -> torch.Tensor: 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) + return indices.to(torch.int16) def decompress(self, tensor: torch.Tensor) -> torch.Tensor: """ From feed1a696fc1f757d7768532c313c1520c79375f Mon Sep 17 00:00:00 2001 From: sfluegel Date: Mon, 3 Aug 2026 10:17:44 +0200 Subject: [PATCH 4/4] fix trailing whitespace --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 312f3df..b2a8ec0 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,7 @@ The list can be found in the `configs/data/chebi50_graph_properties.yml` file. python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/csv_logger.yml --model=../python-chebai-graph/configs/model/gnn_res_gated.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/chebi50_graph_properties.yml --data.init_args.batch_size=128 --trainer.accumulate_grad_batches=4 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=241 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_weighted.yml ``` -## Augmented Graphs +## Augmented Graphs _See thesis related to this work [here](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf)_. Graph Neural Networks (GNNs) often fail to explicitly leverage the chemically meaningful substructures present within molecules (i.e. **functional groups (FGs)**). To make this implicit information explicitly accessible to GNNs, we augment molecular graphs with **artificial nodes** that represent these substructures. The resulting graph are referred to as **augmented graphs**.