Skip to content

add compression to property files - #42

Draft
sfluegel05 wants to merge 2 commits into
devfrom
feature/property-compression
Draft

add compression to property files#42
sfluegel05 wants to merge 2 commits into
devfrom
feature/property-compression

Conversation

@sfluegel05

Copy link
Copy Markdown
Contributor

Property files take up an absurd amount of disk space.

For instance, for a ChEBI50 dataset (v244), the graph_properties\properties folder has nearly 10GB. AtomType_one_hot.pt alone has 6.4GB.

This PR compresses data before saving it to a file (and decompresses when loading). Especially the one-hot encodings don't need to be saved as one-hot tensors. So they are now stored as 1-d and re-converted to one-hot when loading.

File sizes are reduced by a factor of up to 109:

DONE  AtomAromaticity_bool.pt: 106.9 MB -> 59.6 MB (1.8x)
DONE  AtomCharge_one_hot.pt: 743.5 MB -> 59.0 MB (12.6x)
DONE  AtomHybridization_one_hot.pt: 427.6 MB -> 61.3 MB (7.0x)
DONE  AtomNumHs_one_hot.pt: 424.4 MB -> 58.7 MB (7.2x)
DONE  AtomType_one_hot.pt: 6376.0 MB -> 58.5 MB (109.0x)
DONE  BondAromaticity_bool.pt: 109.5 MB -> 59.8 MB (1.8x)
DONE  BondInRing_bool.pt: 107.1 MB -> 58.2 MB (1.8x)
DONE  BondType_one_hot.pt: 331.8 MB -> 58.5 MB (5.7x)
DONE  NumAtomBonds_one_hot.pt: 639.1 MB -> 59.6 MB (10.7x)
DONE  RDKit2DNormalized_asis.pt: 360.7 MB -> 203.5 MB (1.8x)

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR introduces a compression/decompression layer for cached molecular property tensors to substantially reduce on-disk storage, with special handling to avoid storing large dense one-hot matrices.

Changes:

  • Add compress() / decompress() hooks to PropertyEncoder, with concrete implementations for one-hot, boolean, and float “as-is” encoders.
  • Apply compression when saving processed property caches and decompression immediately after loading them in the ChEBI dataset pipeline.
  • Store one-hot encodings as compact class-index vectors on disk and reconstruct one-hot tensors at load time.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File Description
chebai_graph/preprocessing/property_encoder.py Adds encoder-level compression/decompression APIs and implements compact storage for one-hot/bool/float encodings.
chebai_graph/preprocessing/datasets/chebi.py Integrates compression into property caching (torch.save) and decompression into property loading (torch.load).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment on lines +299 to +300
dtype = torch.uint8 if tensor.shape[1] + 1 < 256 else torch.int16
return indices.to(dtype)
Comment on lines 168 to +176
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,
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants