[MNT] resolve all comments in abl_model.py

abl/learning/abl_model.py

@@ -9,7 +9,8 @@
 #   Description   ：
 #
 # ================================================================#
-from itertools import chain
+import pickle
+from utils import flatten, reform_idx
 from typing import List, Any, Optional
 
 
@@ -29,18 +30,31 @@ class ABLModel:
 
     Methods
     -------
-    predict(X: List[List[Any]], mapping: Optional[dict]) -> dict
+    predict(X: List[List[Any]], mapping: Optional[dict] = None) -> dict
         Predict the labels and probabilities for the given data.
     valid(X: List[List[Any]], Y: List[Any]) -> float
         Calculate the accuracy score for the given data.
-    train(X: List[List[Any]], Y: List[Any])
+    train(X: List[List[Any]], Y: List[Any]) -> float
         Train the model on the given data.
+    save(*args, **kwargs) -> None
+        Save the model to a file.
+    load(*args, **kwargs) -> None
+        Load the model from a file.
     """
 
     def __init__(self, base_model) -> None:
         self.classifier_list = []
         self.classifier_list.append(base_model)
 
+        if not (
+            hasattr(base_model, "fit")
+            and hasattr(base_model, "predict")
+            and hasattr(base_model, "score")
+        ):
+            raise NotImplementedError(
+                "base_model should have fit, predict and score methods."
+            )
+
     def predict(self, X: List[List[Any]], mapping: Optional[dict] = None) -> dict:
         """
         Predict the labels and probabilities for the given data.
@@ -49,20 +63,28 @@ class ABLModel:
         ----------
         X : List[List[Any]]
             The data to predict on.
+        mapping : Optional[dict], optional
+            A mapping dictionary to map labels to their original values, by default None.
 
         Returns
         -------
         dict
             A dictionary containing the predicted labels and probabilities.
         """
-        data_X, marks = self.merge_data(X)
-        prob = self.classifier_list[0].predict_proba(X=data_X)
-        label = prob.argmax(axis=1)
+        model = self.classifier_list[0]
+        data_X = flatten(X)
+        if hasattr(model, "predict_proba"):
+            prob = model.predict_proba(X=data_X)
+            label = prob.argmax(axis=1)
+            prob = reform_idx(prob, X)
+        else:
+            prob = None
+            label = model.predict(X=data_X)
+
         if mapping is not None:
-            label = [mapping[x] for x in label]
+            label = [mapping[y] for y in label]
 
-        prob = self.reshape_data(prob, marks)
-        label = self.reshape_data(label, marks)
+        label = reform_idx(label, X)
 
         return {"label": label, "prob": prob}
 
@@ -82,8 +104,8 @@ class ABLModel:
         float
             The accuracy score for the given data.
         """
-        data_X, _ = self.merge_data(X)
-        data_Y, _ = self.merge_data(Y)
+        data_X = flatten(X)
+        data_Y = flatten(Y)
         score = self.classifier_list[0].score(X=data_X, y=data_Y)
         return score
 
@@ -97,37 +119,44 @@ class ABLModel:
             The data to train on.
         Y : List[Any]
             The true labels for the given data.
+
+        Returns
+        -------
+        float
+            The loss value of the trained model.
         """
-        data_X, _ = self.merge_data(X)
-        data_Y, _ = self.merge_data(Y)
+        data_X = flatten(X)
+        data_Y = flatten(Y)
         return self.classifier_list[0].fit(X=data_X, y=data_Y)
-    
-    def save(self, *args, **kwargs) -> None:
-        _model = self.classifier_list[0]
-        if hasattr(_model, "save"):
-            _model.save(*args, **kwargs)
-        else:
-            raise NotImplementedError(f"{type(_model).__name__} object dosen't have the save method")
-        
-    def load(self, *args, **kwargs):
-        _model = self.classifier_list[0]
-        if hasattr(_model, "load"):
-            _model.load(*args, **kwargs)
+
+    def _model_operation(self, operation: str, *args, **kwargs):
+        model = self.classifier_list[0]
+        if hasattr(model, operation):
+            method = getattr(model, operation)
+            method(*args, **kwargs)
         else:
-            raise NotImplementedError(f"{type(_model).__name__} object dosen't have the load method")
-
-    @staticmethod
-    def merge_data(X):
-        ret_mark = list(map(lambda x: len(x), X))
-        ret_X = list(chain(*X))
-        return ret_X, ret_mark
-
-    @staticmethod
-    def reshape_data(Y, marks):
-        begin_mark = 0
-        ret_Y = []
-        for mark in marks:
-            end_mark = begin_mark + mark
-            ret_Y.append(list(Y[begin_mark:end_mark]))
-            begin_mark = end_mark
-        return ret_Y
+            try:
+                if not f"{operation}_path" in kwargs.keys():
+                    raise ValueError(f"'{operation}_path' should not be None")
+                if operation == "save":
+                    with open(kwargs["save_path"], 'wb') as file:
+                        pickle.dump(model, file, protocol=pickle.HIGHEST_PROTOCOL)
+                elif operation == "load":
+                    with open(kwargs["load_path"], 'rb') as file:
+                        self.classifier_list[0] = pickle.load(file)
+            except:
+                raise NotImplementedError(
+                    f"{type(model).__name__} object doesn't have the {operation} method"
+                )
+
+    def save(self, *args, **kwargs) -> None:
+        """
+        Save the model to a file.
+        """
+        self._model_operation("save", *args, **kwargs)
+
+    def load(self, *args, **kwargs) -> None:
+        """
+        Load the model from a file.
+        """
+        self._model_operation("load", *args, **kwargs)

abl/utils/utils.py

@@ -30,33 +30,36 @@ def flatten(nested_list):
     return list(chain.from_iterable(nested_list))
 
 
-def reform_idx(flattened_pred, saved_pred):
+def reform_idx(flattened_list, structured_list):
     """
-    Reform the index based on saved_pred structure.
+    Reform the index based on structured_list structure.
 
     Parameters
     ----------
-    flattened_pred : list
+    flattened_list : list
         A flattened list of predictions.
-    saved_pred : list
+    structured_list : list
         A list containing saved predictions, which could be nested lists or tuples.
 
     Returns
     -------
     list
-        A reformed list that mimics the structure of saved_pred.
+        A reformed list that mimics the structure of structured_list.
     """
-    if not isinstance(saved_pred[0], (list, tuple)):
-        return flattened_pred
+    if not isinstance(flattened_list, list):
+        raise TypeError("Input must be of type list.")
+    
+    if not isinstance(structured_list[0], (list, tuple)):
+        return flattened_list
 
-    reformed_pred = []
+    reformed_list = []
     idx_start = 0
-    for elem in saved_pred:
+    for elem in structured_list:
         idx_end = idx_start + len(elem)
-        reformed_pred.append(flattened_pred[idx_start:idx_end])
+        reformed_list.append(flattened_list[idx_start:idx_end])
         idx_start = idx_end
 
-    return reformed_pred
+    return reformed_list
 
 
 def hamming_dist(pred_pseudo_label, candidates):