[MNT] resolve some comments

abl/learning/abl_model.py

@@ -13,7 +13,7 @@ import pickle
 from typing import Any, Dict
 
 from ..structures import ListData
-from ..utils import reform_idx
+from ..utils import reform_list
 
 
 class ABLModel:
@@ -69,11 +69,11 @@ class ABLModel:
         if hasattr(model, "predict_proba"):
             prob = model.predict_proba(X=data_X)
             label = prob.argmax(axis=1)
-            prob = reform_idx(prob, data_samples.X)
+            prob = reform_list(prob, data_samples.X)
         else:
             prob = None
             label = model.predict(X=data_X)
-        label = reform_idx(label, data_samples.X)
+        label = reform_list(label, data_samples.X)
 
         data_samples.pred_idx = label
         if prob is not None:

abl/learning/basic_nn.py

@@ -66,7 +66,8 @@ class BasicNN:
         num_workers: int = 0,
         save_interval: Optional[int] = None,
         save_dir: Optional[str] = None,
-        transform: Callable[..., Any] = None,
+        train_transform: Callable[..., Any] = None,
+        test_transform: Callable[..., Any] = None,
         collate_fn: Callable[[List[T]], Any] = None,
     ) -> None:
         self.model = model.to(device)
@@ -79,9 +80,18 @@ class BasicNN:
         self.num_workers = num_workers
         self.save_interval = save_interval
         self.save_dir = save_dir
-        self.transform = transform
+        self.train_transform = train_transform
+        self.test_transform = test_transform
         self.collate_fn = collate_fn
 
+        if self.train_transform is not None and self.test_transform is None:
+            print_log(
+                "Transform used in the training phase will be used in prediction.",
+                "current",
+                level=logging.WARNING,
+            )
+            self.test_transform = self.train_transform
+
     def _fit(self, data_loader) -> float:
         """
         Internal method to fit the model on data for n epochs, with early stopping.
@@ -198,12 +208,7 @@ class BasicNN:
 
         return torch.cat(results, axis=0)
 
-    def predict(
-        self,
-        data_loader: DataLoader = None,
-        X: List[Any] = None,
-        test_transform: Callable[..., Any] = None,
-    ) -> numpy.ndarray:
+    def predict(self, data_loader: DataLoader = None, X: List[Any] = None) -> numpy.ndarray:
         """
         Predict the class of the input data.
 
@@ -221,15 +226,7 @@ class BasicNN:
         """
 
         if data_loader is None:
-            if test_transform is None:
-                print_log(
-                    "Transform used in the training phase will be used in prediction.",
-                    "current",
-                    level=logging.WARNING,
-                )
-                dataset = PredictionDataset(X, self.transform)
-            else:
-                dataset = PredictionDataset(X, test_transform)
+            dataset = PredictionDataset(X, self.test_transform)
             data_loader = DataLoader(
                 dataset,
                 batch_size=self.batch_size,
@@ -238,12 +235,7 @@ class BasicNN:
             )
         return self._predict(data_loader).argmax(axis=1).cpu().numpy()
 
-    def predict_proba(
-        self,
-        data_loader: DataLoader = None,
-        X: List[Any] = None,
-        test_transform: Callable[..., Any] = None,
-    ) -> numpy.ndarray:
+    def predict_proba(self, data_loader: DataLoader = None, X: List[Any] = None) -> numpy.ndarray:
         """
         Predict the probability of each class for the input data.
 
@@ -261,15 +253,7 @@ class BasicNN:
         """
 
         if data_loader is None:
-            if test_transform is None:
-                print_log(
-                    "Transform used in the training phase will be used in prediction.",
-                    "current",
-                    level=logging.WARNING,
-                )
-                dataset = PredictionDataset(X, self.transform)
-            else:
-                dataset = PredictionDataset(X, test_transform)
+            dataset = PredictionDataset(X, self.test_transform)
             data_loader = DataLoader(
                 dataset,
                 batch_size=self.batch_size,
@@ -379,7 +363,7 @@ class BasicNN:
         if not (len(y) == len(X)):
             raise ValueError("X and y should have equal length.")
 
-        dataset = ClassificationDataset(X, y, transform=self.transform)
+        dataset = ClassificationDataset(X, y, transform=self.train_transform)
         data_loader = DataLoader(
             dataset,
             batch_size=self.batch_size,

abl/reasoning/kb.py

@@ -9,7 +9,7 @@ from functools import lru_cache
 import numpy as np
 import pyswip
 
-from ..utils.utils import flatten, reform_idx, hamming_dist, to_hashable, restore_from_hashable
+from ..utils.utils import flatten, reform_list, hamming_dist, to_hashable, restore_from_hashable
 from ..utils.cache import abl_cache
 
 
@@ -390,7 +390,7 @@ class PrologKB(KBBase):
 
         for idx in revision_idx:
             revision_pred_pseudo_label[idx] = "P" + str(idx)
-        revision_pred_pseudo_label = reform_idx(revision_pred_pseudo_label, pred_pseudo_label)
+        revision_pred_pseudo_label = reform_list(revision_pred_pseudo_label, pred_pseudo_label)
 
         regex = r"'P\d+'"
         return re.sub(regex, lambda x: x.group().replace("'", ""), str(revision_pred_pseudo_label))
@@ -423,7 +423,7 @@ class PrologKB(KBBase):
             candidate = pred_pseudo_label.copy()
             for i, idx in enumerate(revision_idx):
                 candidate[idx] = c[i]
-            candidate = reform_idx(candidate, save_pred_pseudo_label)
+            candidate = reform_list(candidate, save_pred_pseudo_label)
             candidates.append(candidate)
         return candidates
 

abl/reasoning/reasoner.py

@@ -3,7 +3,7 @@ from zoopt import Dimension, Objective, Parameter, Opt
 from ..utils.utils import (
     confidence_dist,
     flatten,
-    reform_idx,
+    reform_list,
     hamming_dist,
 )
 
@@ -542,7 +542,7 @@ if __name__ == "__main__":
             return candidate
 
         def zoopt_revision_score(self, symbol_num, pred_res, pred_prob, y, sol):
-            all_revision_flag = reform_idx(sol.get_x(), pred_res)
+            all_revision_flag = reform_list(sol.get_x(), pred_res)
             lefted_idxs = [i for i in range(len(pred_res))]
             candidate_size = []
             while lefted_idxs:

abl/utils/utils.py

@@ -31,7 +31,7 @@ def flatten(nested_list):
     return list(chain.from_iterable(nested_list))
 
 
-def reform_idx(flattened_list, structured_list):
+def reform_list(flattened_list, structured_list):
     """
     Reform the index based on structured_list structure.