dataset changes

src/preprocess.py

@@ -80,6 +80,14 @@ class Tox21Dataset(InMemoryDataset):
 
         # Clean molecules & filter dataframe
         mols, clean_mask = create_clean_mol_objects(dataframe["smiles"].tolist())
+        self.clean_mask = torch.tensor(clean_mask, dtype=torch.bool)
+        
+        drop_cols = ["ID","smiles","inchikey","sdftitle","order","set","CVfold"]
+        labels_df = dataframe.drop(columns=drop_cols)
+        numeric_labels = labels_df.apply(pd.to_numeric, errors="coerce").fillna(0.0)
+        self.all_labels = torch.tensor(numeric_labels.values, dtype=torch.float)
+        self.all_label_masks = torch.tensor(~labels_df.isna().values, dtype=torch.bool)
+        
         dataframe = dataframe[clean_mask].reset_index(drop=True)
 
         # Now mols and dataframe are aligned, so we can zip
@@ -88,7 +96,6 @@ class Tox21Dataset(InMemoryDataset):
                 data = from_rdmol(mol)
 
                 # Extract labels as a pandas Series
-                drop_cols = ["ID","smiles","inchikey","sdftitle","order","set","CVfold"]
                 labels = row.drop(drop_cols)
 
                 # Mask for valid labels