Feature/data augmentation

chebai/models/base.py

@@ -46,7 +46,7 @@ def __init__(
         super().__init__()
         self.criterion = criterion
         self.save_hyperparameters(
-            ignore=["criterion", "train_metrics", "val_metrics", "test_metrics"]
+            ignore=["criterion", "train_metrics", "val_metrics", "test_metrics","_class_path"]
         )
         self.out_dim = out_dim
         if optimizer_kwargs:

chebai/molecule.py

@@ -66,7 +66,7 @@ class Molecule:
     max_number_of_parents = 7
 
     def __init__(
-        self, smile: str, logp: Optional[float] = None, contract_rings: bool = False
+            self, smile: str, logp: Optional[float] = None, contract_rings: bool = False
     ):
         """
         Initializes a Molecule object.
@@ -400,8 +400,8 @@ def num_of_features() -> int:
             int: Total number of features.
         """
         return (
-            Molecule.max_number_of_parents * Molecule.num_bond_features()
-            + Molecule.num_atom_features()
+                Molecule.max_number_of_parents * Molecule.num_bond_features()
+                + Molecule.num_atom_features()
         )
 
     @staticmethod

chebai/result/utils.py

@@ -94,6 +94,9 @@ def evaluate_model(
     Returns:
         Tensors with predictions and labels.
     """
+    print("Start of evaluate_model")
+    batch_size=5
+    print("batch_size: ", batch_size)
     model.eval()
     collate = data_module.reader.COLLATOR()
 
@@ -157,6 +160,7 @@ def evaluate_model(
                 torch.cat(labels_list),
                 os.path.join(buffer_dir, f"labels{save_ind:03d}.pt"),
             )
+    print("End of evaluate_model")
 
 
 def load_results_from_buffer(
@@ -172,6 +176,7 @@ def load_results_from_buffer(
     Returns:
         Tensors with predictions and labels.
     """
+    print("Start of load_results_from_buffer")
     preds_list = []
     labels_list = []
 
@@ -208,6 +213,7 @@ def load_results_from_buffer(
     else:
         test_labels = None
 
+    print("End of load_results_from_buffer")
     return test_preds, test_labels
 
 

configs/data/chebi50.yml

@@ -1 +1,5 @@
 class_path: chebai.preprocessing.datasets.chebi.ChEBIOver50
+init_args:
+  aug_data: True
+  augment_data_batch_size: 5000
+  num_smiles_variations: 5

configs/data/chebi_augmentation.yml

@@ -0,0 +1,5 @@
+class_path: chebai.preprocessing.datasets.chebi.ChEBIOver100
+init_args:
+  aug_data: True
+  augment_data_batch_size: 5000
+  num_smiles_variations: 5

eval.py

@@ -0,0 +1,74 @@
+import pandas as pd
+
+from chebai.result.utils import (
+    evaluate_model,
+    load_results_from_buffer,
+)
+from chebai.result.classification import print_metrics
+from chebai.models.electra import Electra
+from chebai.preprocessing.datasets.chebi import ChEBIOver50, ChEBIOver100
+import os
+import tqdm
+import torch
+import pickle
+
+DEVICE = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
+print(DEVICE)
+
+
+# Specify paths and parameters
+checkpoint_name = "best_epoch=31_val_loss=0.0204_val_macro-f1=0.7655_val_micro-f1=0.9246"
+print("checkpoint_name : ",checkpoint_name)
+checkpoint_path = os.path.join("logs/wandb/run-20241212_003611-8yohluv6/files/checkpoints", f"{checkpoint_name}.ckpt")
+print("checkpoint_path : ",checkpoint_path)
+kind = "test"  # Change to "train" or "validation" as needed
+buffer_dir = os.path.join("results_buffer", checkpoint_name, kind)
+print("buffer_dir : ",buffer_dir)
+batch_size = 10  # Set batch size
+
+# Load data module
+data_module = ChEBIOver100(chebi_version=231)
+
+data_module.splits_file_path="data/chebi_v231/ChEBI100/processed/augmented_splits.csv"
+model_class = Electra
+
+# evaluates model, stores results in buffer_dir
+model = model_class.load_from_checkpoint(checkpoint_path)
+if buffer_dir is None:
+    preds, labels = evaluate_model(
+        model,
+        data_module,
+        buffer_dir=buffer_dir,
+        # No need to provide this parameter for Chebi dataset, "kind" parameter should be provided
+        # filename=data_module.processed_file_names_dict[kind],
+        batch_size=10,
+        kind=kind,
+    )
+else:
+    evaluate_model(
+        model,
+        data_module,
+        buffer_dir=buffer_dir,
+        # No need to provide this parameter for Chebi dataset, "kind" parameter should be provided
+        # filename=data_module.processed_file_names_dict[kind],
+        batch_size=10,
+        kind=kind,
+    )
+    # load data from buffer_dir
+    preds, labels = load_results_from_buffer(buffer_dir, device=DEVICE)
+
+
+# Load classes from the classes.txt
+with open(os.path.join(data_module.processed_dir_main, "classes.txt"), "r") as f:
+    classes = [line.strip() for line in f.readlines()]
+
+
+# output relevant metrics
+print_metrics(
+    preds,
+    labels.to(torch.int),
+    DEVICE,
+    classes=classes,
+    markdown_output=False,
+    top_k=10,
+)