[MNT] modify details about class name and logger

learnware/market/heterogeneous/organizer/__init__.py

@@ -15,7 +15,7 @@ import torch.multiprocessing as mp
 
 from ....learnware import Learnware, get_learnware_from_dirpath
 from ....logger import get_module_logger
-from ....specification.system import HeteroSpecification
+from ....specification.system import HeteroMapTableSpecification
 from ...base import BaseChecker, BaseUserInfo
 from ...easy import EasyOrganizer
 from ...easy.database_ops import DatabaseOperations
@@ -71,15 +71,15 @@ class HeteroMapTableOrganizer(EasyOrganizer):
             if not rebuild:
                 if os.path.exists(self.hetero_mappings_path):
                     for hetero_json_path in os.listdir(self.hetero_mappings_path):
-                        idx = hetero_json_path.split('.')[0]
-                        hetero_spec = HeteroSpecification()
+                        idx = hetero_json_path.split(".")[0]
+                        hetero_spec = HeteroMapTableSpecification()
                         hetero_spec.load(os.path.join(self.hetero_mappings_path, f"{idx}.json"))
                         try:
-                            self.learnware_list[idx].update_stat_spec("HeteroSpecification", hetero_spec)
+                            self.learnware_list[idx].update_stat_spec("HeteroMapTableSpecification", hetero_spec)
                         except:
                             logger.warning(f"Learnware ID {idx} NOT Found!")
                 else:
-                    logger.info("No HeteroSpecifications to reload. Use loaded market mapping to regenerate.")
+                    logger.info("No HeteroMapTableSpecification to reload. Use loaded market mapping to regenerate.")
                     self._update_learnware_by_ids(self.learnware_list.keys())
         else:
             logger.warning(f"No market mapping to reload!!")
@@ -90,7 +90,8 @@ class HeteroMapTableOrganizer(EasyOrganizer):
         self.auto_update = auto_update
         self.market_id = market_id
         self.training_args = kwargs
-        if auto_update_limit is not None: self.auto_update_limit = auto_update_limit
+        if auto_update_limit is not None:
+            self.auto_update_limit = auto_update_limit
 
     def add_learnware(
         self, zip_path: str, semantic_spec: dict, check_status: int, learnware_id: str = None
@@ -98,7 +99,7 @@ class HeteroMapTableOrganizer(EasyOrganizer):
         if check_status == BaseChecker.INVALID_LEARNWARE:
             logger.warning("Learnware is invalid!")
             return None, BaseChecker.INVALID_LEARNWARE
-        
+
         semantic_spec = copy.deepcopy(semantic_spec)
         logger.info("Get new learnware from %s" % (zip_path))
 
@@ -123,7 +124,7 @@ class HeteroMapTableOrganizer(EasyOrganizer):
             except:
                 pass
             return None, BaseChecker.INVALID_LEARNWARE
-        
+
         if new_learnware is None:
             return None, BaseChecker.INVALID_LEARNWARE
 
@@ -143,7 +144,7 @@ class HeteroMapTableOrganizer(EasyOrganizer):
         self.use_flags[learnware_id] = learnwere_status
         self._update_learnware_by_ids([learnware_id])
         self.count += 1
-        self.training_count += ([learnware_id] == self._get_table_type_learnware_ids([learnware_id]))
+        self.training_count += [learnware_id] == self._get_table_type_learnware_ids([learnware_id])
 
         if self.auto_update and self.training_count - self.last_training_count == self.auto_update_limit + 1:
             training_learnware_ids = self._get_table_type_learnware_ids(self.get_learnware_ids())
@@ -151,16 +152,16 @@ class HeteroMapTableOrganizer(EasyOrganizer):
             logger.warning(f"Leanwares for training: {training_learnware_ids}")
 
             updated_market_mapping = self.train(
-                learnware_list=training_learnwares,
-                save_dir=self.market_store_path, 
-                **self.training_args
+                learnware_list=training_learnwares, save_dir=self.market_store_path, **self.training_args
+            )
+
+            logger.warning(
+                f"Market mapping train completed. Now update HeteroMapTableSpecification for {training_learnware_ids}"
             )
-            
-            logger.warning(f"Market mapping train completed. Now update HeteroSpecification for {training_learnware_ids}")
             self.market_mapping = updated_market_mapping
             self._update_learnware_by_ids(training_learnware_ids)
             self.last_training_count = len(training_learnware_ids)
-        
+
         return learnware_id, learnwere_status
 
     @staticmethod
@@ -178,7 +179,7 @@ class HeteroMapTableOrganizer(EasyOrganizer):
         market_mapping_trainer.save_model(output_dir=save_dir)
 
         return market_mapping
-    
+
     def _update_learnware_by_ids(self, ids: List[str]):
         ids = self._get_table_type_learnware_ids(ids)
         for id in ids:
@@ -187,14 +188,14 @@ class HeteroMapTableOrganizer(EasyOrganizer):
                 semantic_spec, stat_spec = spec.get_semantic_spec(), spec.get_stat_spec()["RKMETableSpecification"]
                 features = semantic_spec["Input"]["Description"].values()
                 hetero_spec = self.market_mapping.hetero_mapping(stat_spec, features)
-                self.learnware_list[id].update_stat_spec("HeteroSpecification", hetero_spec)
-                
+                self.learnware_list[id].update_stat_spec("HeteroMapTableSpecification", hetero_spec)
+
                 save_path = os.path.join(self.hetero_mappings_path, f"{id}.json")
                 hetero_spec.save(save_path)
             except Exception as err:
-                logger.warning(f"Learnware {id} generate HeteroSpecification failed! Due to {err}")
-            
-    def generate_hetero_map_spec(self, user_info: BaseUserInfo) -> HeteroSpecification:
+                logger.warning(f"Learnware {id} generate HeteroMapTableSpecification failed! Due to {err}")
+
+    def generate_hetero_map_spec(self, user_info: BaseUserInfo) -> HeteroMapTableSpecification:
         user_stat_spec = user_info.stat_info["RKMETableSpecification"]
         user_features = user_info.get_semantic_spec()["Input"]["Description"].values()
 
@@ -210,13 +211,13 @@ class HeteroMapTableOrganizer(EasyOrganizer):
             features = spec.get_semantic_spec()["Input"]["Description"]
             learnware_df = pd.DataFrame(data=stat_spec.get_z(), columns=features.values())
             learnware_df_dict[tuple(sorted(features))].append(learnware_df)
-        
+
         return [pd.concat(dfs) for dfs in learnware_df_dict.values()]
-    
+
     def _get_table_type_learnware_ids(self, ids: List[str]) -> List[str]:
         ret = []
         for id in ids:
             semantic_spec = self.learnware_list[id].get_specification().get_semantic_spec()
             if semantic_spec["Data"]["Values"][0] == "Table":
                 ret.append(id)
-        return ret
+        return ret

learnware/market/heterogeneous/organizer/hetero_mapping/__init__.py

@@ -7,7 +7,7 @@ import torch
 import torch.nn.functional as F
 from torch import Tensor, nn
 
-from .....specification import HeteroSpecification, RKMETableSpecification
+from .....specification import HeteroMapTableSpecification, RKMETableSpecification
 from .feature_extractor import *
 from .trainer import Trainer, TransTabCollatorForCL
 
@@ -147,8 +147,8 @@ class HeteroMapping(nn.Module):
         loss = self._self_supervised_contrastive_loss(feat_x_multiview)
         return loss
 
-    def hetero_mapping(self, rkme_spec: RKMETableSpecification, cols: List[str]) -> HeteroSpecification:
-        hetero_spec = HeteroSpecification()
+    def hetero_mapping(self, rkme_spec: RKMETableSpecification, cols: List[str]) -> HeteroMapTableSpecification:
+        hetero_spec = HeteroMapTableSpecification()
         hetero_input_df = pd.DataFrame(data=rkme_spec.get_z(), columns=cols)
         hetero_embedding = self._extract_batch_features(hetero_input_df)
         hetero_spec.generate_stat_spec_from_system(hetero_embedding, rkme_spec)

learnware/market/heterogeneous/organizer/hetero_mapping/feature_extractor.py

@@ -6,7 +6,6 @@ from typing import Dict
 import numpy as np
 import torch
 import torch.nn.init as nn_init
-from loguru import logger
 from torch import Tensor, nn
 from transformers import BertTokenizerFast
 

learnware/market/heterogeneous/organizer/hetero_mapping/trainer.py

@@ -6,12 +6,14 @@ import time
 import numpy as np
 import pandas as pd
 import torch
-from loguru import logger
 from torch import nn
 from torch.utils.data import DataLoader, Dataset
 from tqdm.autonotebook import trange
 
 from .feature_extractor import FeatureTokenizer
+from .....logger import get_module_logger
+
+logger = get_module_logger("hetero_mapping_trainer")
 
 
 class Trainer:

learnware/market/heterogeneous/searcher.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from ...learnware import Learnware
 from ...logger import get_module_logger
-from ...specification import HeteroSpecification
+from ...specification import HeteroMapTableSpecification
 from ..base import BaseSearcher, BaseUserInfo
 from ..easy import EasySearcher
 from ..utils import parse_specification_type
@@ -34,28 +34,28 @@ class HeteroMapTableSearcher(EasySearcher):
             return [(max_dist - dist) / (max_dist - dist_epsilon) for dist in dist_list]
 
     def _search_by_hetero_spec_single(
-        self, 
-        learnware_list: List[Learnware], 
-        user_hetero_spec: HeteroSpecification
+        self, learnware_list: List[Learnware], user_hetero_spec: HeteroMapTableSpecification
     ) -> Tuple[List[float], List[Learnware]]:
-        hetero_spec_list = [learnware.specification.get_stat_spec_by_name("HeteroSpecification") for learnware in learnware_list]
+        hetero_spec_list = [
+            learnware.specification.get_stat_spec_by_name("HeteroMapTableSpecification") for learnware in learnware_list
+        ]
         mmd_dist_list = []
         for idx, hetero_spec in enumerate(hetero_spec_list):
             mmd_dist = hetero_spec.dist(user_hetero_spec)
             mmd_dist_list.append(mmd_dist)
-        
+
         sorted_idx_list = sorted(range(len(learnware_list)), key=lambda k: mmd_dist_list[k])
         sorted_dist_list = [mmd_dist_list[idx] for idx in sorted_idx_list]
         sorted_learnware_list = [learnware_list[idx] for idx in sorted_idx_list]
 
         return sorted_dist_list, sorted_learnware_list
-    
+
     def _filter_by_hetero_spec_single(
         self,
         sorted_score_list: List[float],
         learnware_list: List[Learnware],
         filter_score: float = 0.5,
-        min_num: int = 5
+        min_num: int = 5,
     ) -> Tuple[List[float], List[Learnware]]:
         idx = min(min_num, len(learnware_list))
         while idx < len(learnware_list):
@@ -64,11 +64,10 @@ class HeteroMapTableSearcher(EasySearcher):
             idx += 1
         return sorted_score_list[:idx], learnware_list[:idx]
 
-
     def __call__(
-        self, 
-        learnware_list: List[Learnware], 
-        user_info: BaseUserInfo, 
+        self,
+        learnware_list: List[Learnware],
+        user_info: BaseUserInfo,
     ) -> Tuple[List[float], List[Learnware], float, List[Learnware]]:
         # todo: use specially assigned search_gamma for calculating mmd dist
         user_hetero_spec = self.learnware_oganizer.generate_hetero_map_spec(user_info)
@@ -88,6 +87,7 @@ class HeteroMapTableSearcher(EasySearcher):
     def reset(self, organizer):
         self.learnware_oganizer = organizer
 
+
 class HeteroSearcher(EasySearcher):
     def __init__(self, organizer: HeteroMapTableOrganizer = None):
         super(HeteroSearcher, self).__init__(organizer)
@@ -96,7 +96,7 @@ class HeteroSearcher(EasySearcher):
     def reset(self, organizer):
         super().reset(organizer)
         self.hetero_stat_searcher.reset(organizer)
-    
+
     @staticmethod
     def check_user_info(user_info: BaseUserInfo):
         try:
@@ -105,7 +105,9 @@ class HeteroSearcher(EasySearcher):
 
             user_task_type = user_info.get_semantic_spec()["Task"]["Values"]
             if user_task_type not in [["Classification"], ["Regression"]]:
-                logger.warning("User doesn't provide correct task type, it must be either Classification or Regression.")
+                logger.warning(
+                    "User doesn't provide correct task type, it must be either Classification or Regression."
+                )
                 return False
 
             user_input_description = user_info.get_semantic_spec()["Input"]
@@ -115,10 +117,12 @@ class HeteroSearcher(EasySearcher):
             if user_input_shape != user_description_dim or user_input_shape != user_description_feature_num:
                 logger.warning("User data feature dimensions mismatch with semantic specification.")
                 return False
-            
+
             return True
         except Exception as e:
-            logger.info(f"Invalid heterogeneous search information provided. Use homogeneous search instead. Error: {e}")
+            logger.info(
+                f"Invalid heterogeneous search information provided. Use homogeneous search instead. Error: {e}"
+            )
             return False
 
     def __call__(
@@ -136,4 +140,4 @@ class HeteroSearcher(EasySearcher):
             else:
                 return self.stat_searcher(learnware_list, user_info, max_search_num, search_method)
         else:
-            return None, learnware_list, 0.0, None
+            return None, learnware_list, 0.0, None

learnware/reuse/hetero_reuser/feature_alignment.py

@@ -6,13 +6,15 @@ import torch.nn.functional as F
 import torch
 import time
 from tqdm import trange
-from loguru import logger
 
 from learnware.learnware import Learnware
 from learnware.specification import RKMETableSpecification
 from learnware.specification.regular.table.rkme import choose_device
 
 from ..base import BaseReuser
+from ...logger import get_module_logger
+
+logger = get_module_logger("hetero_feature_alignment")
 
 
 class FeatureAligner(BaseReuser):
@@ -66,7 +68,9 @@ class FeatureAligner(BaseReuser):
             The RKME specification from the user dataset.
         """
         target_rkme = self.learnware.specification.get_stat_spec()["RKMETableSpecification"]
-        trainer = FeatureAlignmentTrainer(target_rkme=target_rkme, user_rkme=user_rkme, cuda_idx=self.cuda_idx, **self.align_arguments)
+        trainer = FeatureAlignmentTrainer(
+            target_rkme=target_rkme, user_rkme=user_rkme, cuda_idx=self.cuda_idx, **self.align_arguments
+        )
         self.align_model = trainer.model
         self.align_model.eval()
 
@@ -85,7 +89,9 @@ class FeatureAligner(BaseReuser):
             Predicted output from the learnware model after alignment.
         """
         user_data = self._fill_data(user_data)
-        transformed_user_data = self.align_model(torch.tensor(user_data, device=self.device).float()).detach().cpu().numpy()
+        transformed_user_data = (
+            self.align_model(torch.tensor(user_data, device=self.device).float()).detach().cpu().numpy()
+        )
         y_pred = self.learnware.predict(transformed_user_data)
         return y_pred
 
@@ -120,7 +126,6 @@ class FeatureAligner(BaseReuser):
         return X
 
 
-
 class FeatureAlignmentModel(nn.Module):
     """
     FeatureAlignmentModel is a neural network module designed for feature alignment tasks.
@@ -128,7 +133,15 @@ class FeatureAlignmentModel(nn.Module):
     and supports different activation functions.
     """
 
-    def __init__(self, input_dim: int, output_dim: int, hidden_dims: list = [1024], activation: str = "relu", dropout_ratio: float = 0, use_bn: bool = False):
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        hidden_dims: list = [1024],
+        activation: str = "relu",
+        dropout_ratio: float = 0,
+        use_bn: bool = False,
+    ):
         """
         Initialize the FeatureAlignmentModel.
 
@@ -187,13 +200,13 @@ class FeatureAlignmentModel(nn.Module):
         """
         if len(self.fc_list) > 0:
             for fc, drop in zip(self.fc_list, self.drop_list):
-                x = fc(x)      # Apply fully connected layer
+                x = fc(x)  # Apply fully connected layer
                 x = self.activation(x)  # Apply activation function
-                x = drop(x)    # Apply dropout
+                x = drop(x)  # Apply dropout
         return self.final_fc(x)  # Return output from final fully connected layer
-    
 
-class FeatureAlignmentTrainer():
+
+class FeatureAlignmentTrainer:
     """
     FeatureAlignmentTrainer is a class designed to train a neural network for aligning features from a user dataset
     to a target dataset. It utilizes Maximum Mean Discrepancy (MMD) as the loss function for training.
@@ -248,7 +261,7 @@ class FeatureAlignmentTrainer():
         dropout_ratio: float = 0,
         use_bn: bool = False,
         const: float = 1e1,
-        cuda_idx: int = 0
+        cuda_idx: int = 0,
     ):
         """
         Initialize the FeatureAlignmentTrainer with the specified parameters.
@@ -266,7 +279,7 @@ class FeatureAlignmentTrainer():
         }
         self.network_type = network_type
         self.optimizer_type = optimizer_type
-        self.const=const
+        self.const = const
         self.device = choose_device(cuda_idx=cuda_idx)
         if extra_labeled_data is not None and target_learnware is not None:
             self.train_with_labeled_data(extra_labeled_data[0], extra_labeled_data[1], target_learnware)
@@ -294,7 +307,9 @@ class FeatureAlignmentTrainer():
         X12norm = torch.sum(x1**2, 1, keepdim=True) - 2 * x1 @ x2.T + torch.sum(x2**2, 1, keepdim=True).T
         return torch.exp(-X12norm * self.args["gamma"])
 
-    def compute_mmd(self, user_X: torch.Tensor, user_weight: torch.Tensor, target_X: torch.Tensor, target_weight: torch.Tensor) -> torch.Tensor:
+    def compute_mmd(
+        self, user_X: torch.Tensor, user_weight: torch.Tensor, target_X: torch.Tensor, target_weight: torch.Tensor
+    ) -> torch.Tensor:
         """
         Compute the Maximum Mean Discrepancy (MMD) between the user and target datasets.
 
@@ -327,7 +342,9 @@ class FeatureAlignmentTrainer():
         input_dim = self.user_rkme.get_z().shape[1]
         output_dim = self.target_rkme.get_z().shape[1]
 
-        user_model=FeatureAlignmentModel(input_dim, output_dim, args["hidden_dims"], args["activation"], args["dropout_ratio"], args["use_bn"])
+        user_model = FeatureAlignmentModel(
+            input_dim, output_dim, args["hidden_dims"], args["activation"], args["dropout_ratio"], args["use_bn"]
+        )
 
         # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         user_model.to(self.device)
@@ -355,4 +372,4 @@ class FeatureAlignmentTrainer():
             )
 
         self.model = user_model
-        logger.info("training complete, cost {:.1f} secs.".format(time.time() - start_time))
+        logger.info("training complete, cost {:.1f} secs.".format(time.time() - start_time))

learnware/specification/__init__.py

@@ -1,13 +1,13 @@
 from .base import Specification, BaseStatSpecification
 from .regular import (
-    RegularStatsSpecification,
+    RegularStatSpecification,
     RKMEStatSpecification,
     RKMETableSpecification,
     RKMEImageSpecification,
     RKMETextSpecification,
 )
 
-from .system import HeteroSpecification
+from .system import HeteroMapTableSpecification
 
 from ..utils import is_torch_avaliable
 

learnware/specification/regular/__init__.py

@@ -1,4 +1,4 @@
-from .base import RegularStatsSpecification
+from .base import RegularStatSpecification
 from ...utils import is_torch_avaliable
 
 from .text import RKMETextSpecification

learnware/specification/regular/base.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from ..base import BaseStatSpecification
 
 
-class RegularStatsSpecification(BaseStatSpecification):
+class RegularStatSpecification(BaseStatSpecification):
     def generate_stat_spec(self, **kwargs):
         self.generate_stat_spec_from_data(**kwargs)
 

learnware/specification/regular/image/rkme.py

@@ -17,11 +17,11 @@ from torchvision.transforms import Resize
 from tqdm import tqdm
 
 from . import cnn_gp
-from ..base import RegularStatsSpecification
+from ..base import RegularStatSpecification
 from ..table.rkme import solve_qp, choose_device, setup_seed
 
 
-class RKMEImageSpecification(RegularStatsSpecification):
+class RKMEImageSpecification(RegularStatSpecification):
     # INNER_PRODUCT_COUNT = 0
     IMAGE_WIDTH = 32
 

learnware/specification/regular/table/rkme.py

@@ -20,7 +20,7 @@ try:
 except ImportError:
     _FAISS_INSTALLED = False
 
-from ..base import RegularStatsSpecification
+from ..base import RegularStatSpecification
 from ....logger import get_module_logger
 
 logger = get_module_logger("rkme")
@@ -31,7 +31,7 @@ if not _FAISS_INSTALLED:
     )
 
 
-class RKMETableSpecification(RegularStatsSpecification):
+class RKMETableSpecification(RegularStatSpecification):
     """Reduced Kernel Mean Embedding (RKME) Specification"""
 
     def __init__(self, gamma: float = 0.1, cuda_idx: int = -1):

learnware/specification/system/__init__.py

@@ -1 +1 @@
-from .heter_table import HeteroSpecification
+from .heter_table import HeteroMapTableSpecification

learnware/specification/system/base.py

@@ -1,11 +1,7 @@
-from __future__ import annotations
-
-from loguru import logger
-
 from ..base import BaseStatSpecification
 
 
-class SystemStatsSpecification(BaseStatSpecification):
+class SystemStatSpecification(BaseStatSpecification):
     def generate_stat_spec(self, **kwargs):
         self.generate_stat_spec_from_system(**kwargs)
 

learnware/specification/system/heter_table.py

@@ -10,10 +10,10 @@ import torch
 
 from ..regular import RKMETableSpecification
 from ..regular.table.rkme import choose_device, setup_seed, torch_rbf_kernel
-from .base import SystemStatsSpecification
+from .base import SystemStatSpecification
 
 
-class HeteroSpecification(SystemStatsSpecification):
+class HeteroMapTableSpecification(SystemStatSpecification):
     """Heterogeneous Embedding Specification"""
 
     def __init__(self, gamma: float = 0.1, cuda_idx: int = -1):
@@ -26,7 +26,7 @@ class HeteroSpecification(SystemStatsSpecification):
         torch.cuda.empty_cache()
         self.device = choose_device(cuda_idx=cuda_idx)
         setup_seed(0)
-        super(HeteroSpecification, self).__init__(type=self.__class__.__name__)
+        super(HeteroMapTableSpecification, self).__init__(type=self.__class__.__name__)
 
     def get_z(self) -> np.ndarray:
         return self.z.detach().cpu().numpy()
@@ -38,7 +38,7 @@ class HeteroSpecification(SystemStatsSpecification):
         self.beta = rkme_spec.beta.to(self.device)
         self.z = torch.from_numpy(heter_embedding).double().to(self.device)
 
-    def inner_prod(self, Embed2: HeteroSpecification) -> float:
+    def inner_prod(self, Embed2: HeteroMapTableSpecification) -> float:
         beta_1 = self.beta.reshape(1, -1).double().to(self.device)
         beta_2 = Embed2.beta.reshape(1, -1).double().to(self.device)
         Z1 = self.z.double().reshape(self.z.shape[0], -1).to(self.device)
@@ -47,7 +47,7 @@ class HeteroSpecification(SystemStatsSpecification):
 
         return float(v)
 
-    def dist(self, Embed2: HeteroSpecification, omit_term1: bool = False) -> float:
+    def dist(self, Embed2: HeteroMapTableSpecification, omit_term1: bool = False) -> float:
         term1 = 0 if omit_term1 else self.inner_prod(self)
         term2 = self.inner_prod(Embed2)
         term3 = Embed2.inner_prod(Embed2)