From d16270cdeaf46cd304280db56745893173718cc6 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 6 Aug 2026 10:53:30 +0200 Subject: [PATCH 1/2] compression changes --- chebai_graph/preprocessing/datasets/base.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 35b92fc..f2c60c1 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -143,7 +143,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 ], @@ -357,6 +362,9 @@ 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] @@ -505,6 +513,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] From 99571a6c429a9bd5e095fec3b3c8cb83bb9c4d52 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 6 Aug 2026 10:59:58 +0200 Subject: [PATCH 2/2] cast ids to string --- chebai_graph/preprocessing/datasets/base.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index f2c60c1..26bc507 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -357,6 +357,7 @@ def load_processed_data( """ base_data = super().load_processed_data(kind, filename) base_df = pd.DataFrame(base_data) + base_df["ident"] = base_df["ident"].astype(str) for property in self.properties: property_data = torch.load( @@ -374,6 +375,7 @@ def load_processed_data( property_df.rename( columns={property.name: f"{property.name}"}, inplace=True ) + property_df["ident"] = property_df["ident"].astype(str) base_df = base_df.merge(property_df, on="ident", how="left") base_df["features"] = base_df.apply( @@ -459,6 +461,7 @@ def load_processed_data( """ base_data = super().load_processed_data(kind, filename) base_df = pd.DataFrame(base_data) + base_df["ident"] = base_df["ident"].astype(str) props_categories = { "AllNodeTypeProperties": [], "FGNodeTypeProperties": [], @@ -521,6 +524,7 @@ def load_processed_data( ) property_df = pd.DataFrame(property_data) + property_df["ident"] = property_df["ident"].astype(str) property_df.rename( columns={property.name: f"{property.name}"}, inplace=True )