diff --git a/abl/learning/abl_model.py b/abl/learning/abl_model.py index 700976a..f1512fe 100644 --- a/abl/learning/abl_model.py +++ b/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) diff --git a/abl/utils/utils.py b/abl/utils/utils.py index 475e5de..c9d00ce 100644 --- a/abl/utils/utils.py +++ b/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):