From 3b92fd5b419b850a8248bd81faaa0e3d9c43cc4a Mon Sep 17 00:00:00 2001 From: Gao Enhao Date: Thu, 16 Nov 2023 15:53:53 +0800 Subject: [PATCH] [MNT] resolve some comments --- abl/learning/abl_model.py | 6 ++--- abl/learning/basic_nn.py | 50 +++++++++++++-------------------------- abl/reasoning/kb.py | 6 ++--- abl/reasoning/reasoner.py | 4 ++-- abl/utils/utils.py | 2 +- 5 files changed, 26 insertions(+), 42 deletions(-) diff --git a/abl/learning/abl_model.py b/abl/learning/abl_model.py index ab7bfb7..97775c0 100644 --- a/abl/learning/abl_model.py +++ b/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: diff --git a/abl/learning/basic_nn.py b/abl/learning/basic_nn.py index b1da93c..0b43fcb 100644 --- a/abl/learning/basic_nn.py +++ b/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, diff --git a/abl/reasoning/kb.py b/abl/reasoning/kb.py index 1bca1b3..10aa559 100644 --- a/abl/reasoning/kb.py +++ b/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 diff --git a/abl/reasoning/reasoner.py b/abl/reasoning/reasoner.py index 686e9dd..2e57570 100644 --- a/abl/reasoning/reasoner.py +++ b/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: diff --git a/abl/utils/utils.py b/abl/utils/utils.py index 1480045..0b5dff4 100644 --- a/abl/utils/utils.py +++ b/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.