fusion-bench 0.2.12__py3-none-any.whl → 0.2.14__py3-none-any.whl

fusion_bench/method/smile_upscaling/smile_upscaling.py

@@ -13,348 +13,18 @@ from fusion_bench.method import BaseAlgorithm
 from fusion_bench.method.simple_average import simple_average
 from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
 from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.models.smile_moe.linear_from_module import (
+    ExpertNotTrainedError,
+    SmileCompressedLinear,
+    SmileGate,
+    SmileMoELinear,
+)
 from fusion_bench.models.utils import get_attr, set_attr
 from fusion_bench.utils.parameters import print_parameters
 
 log = logging.getLogger(__name__)
 
 
-class ExpertNotTrainedError(Exception):
-    pass
-
-
-def _is_all_zeros(tensor: Tensor | List[Tensor]) -> bool:
-    """
-    Check if a tensor or a list of tensors are all zeros.
-
-    Args:
-        tensor (Tensor | List[Tensor]): A tensor or a list of tensors.
-
-    Returns:
-        bool: True if all elements are zeros, False otherwise.
-    """
-    if isinstance(tensor, Tensor):
-        return torch.allclose(tensor, torch.zeros_like(tensor))
-    else:
-        return all(_is_all_zeros(t) for t in tensor)
-
-
-def _svd(w: Tensor, full_matrices=True) -> Tuple[Tensor, Tensor, Tensor]:
-    """
-    Perform Singular Value Decomposition (SVD) on a tensor.
-
-    Args:
-        w (Tensor): The input tensor.
-        full_matrices (bool): Whether to compute the full-sized U and V matrices.
-
-    Returns:
-        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
-    """
-    u, s, vh = torch.linalg.svd(
-        w, full_matrices=full_matrices, driver="gesvd" if w.is_cuda else None
-    )
-    v = vh.T
-    return u, s, v
-
-
-def svd(
-    w: Tensor, full_matrices=True, accelerator=None
-) -> Tuple[Tensor, Tensor, Tensor]:
-    """
-    Perform SVD on a tensor, optionally using a specified accelerator.
-
-    Args:
-        w (Tensor): The input tensor.
-        full_matrices (bool): Whether to compute the full-sized U and V matrices.
-        accelerator (str): The device to perform the computation on.
-
-    Returns:
-        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
-    """
-    if accelerator is None:
-        return _svd(w, full_matrices=full_matrices)
-    original_device = w.device
-    w = w.to(accelerator)
-    u, s, v = _svd(w)
-    return u.to(original_device), s.to(original_device), v.to(original_device)
-
-
-class SmileGate(nn.Module):
-    def __init__(
-        self,
-        input_features: int,
-        w_diff_list: List[Tensor],
-        k: int,
-        svd_list=None,  # cached `svd_list`, pass it to avoid recomputing
-        upscaling_accelerator=None,
-    ):
-        """
-        Initialize the SmileGate module.
-
-        Args:
-            input_features (int): The number of input features.
-            w_diff_list (List[Tensor]): A list of weight difference tensors.
-            k (int): The number of singular values to keep.
-            svd_list (List[Tuple[Tensor, Tensor, Tensor]]): Cached SVD results.
-            upscaling_accelerator (str): The device to perform the computation on.
-        """
-        super().__init__()
-        self.input_features = input_features
-        self.num_experts = len(w_diff_list)
-        weights = []
-        for i, w_diff in enumerate(w_diff_list):
-            if svd_list is None:
-                u, s, v = svd(w_diff, accelerator=upscaling_accelerator)
-            else:
-                u, s, v = svd_list[i]
-            u = u[:, :k]
-            s = s[:k]
-            v = v[:, :k]
-
-            # weights.append((s * v).T)
-            weights.append(v.T)
-        self.k = s.size(0)  # k is the actual k after truncation
-
-        weights = (
-            torch.stack(weights, dim=0)
-            .reshape(self.num_experts * self.k, -1)
-            .contiguous()
-        )
-        self.weights = nn.Parameter(
-            weights
-        )  # weights should be a tensor of shape (num_experts * k, n)
-
-    def forward(self, x: Tensor):
-        """
-        Forward pass of the SmileGate module.
-
-        Args:
-            x (Tensor): The input tensor.
-
-        Returns:
-            Tensor: The routing weights.
-        """
-        batch_size = x.size(0)
-        if self.num_experts == 1:
-            return torch.ones(batch_size, 1, device=x.device, dtype=x.dtype)
-
-        routing_weights = F.linear(x, self.weights).view(
-            batch_size, self.num_experts, self.k
-        )
-        routing_weights = routing_weights.norm(p=2, dim=2)
-        return routing_weights
-
-
-class SmileCompressedLinear(nn.Module):
-    def __init__(self, model: nn.Linear, k: int, svd_cache=None):
-        """
-        Initialize the SmileCompressedLinear module.
-
-        Args:
-            model (nn.Linear): The linear model to compress.
-            k (int): The number of singular values to keep.
-            svd_cache (Tuple[Tensor, Tensor, Tensor]): Cached SVD results.
-        """
-        super().__init__()
-        if svd_cache is None:
-            u, s, v = svd(model.weight)
-        else:
-            u, s, v = svd_cache
-        if k > 0:
-            u = u[:, :k]
-            s = s[:k]
-            v = v[:, :k]
-
-        self.u = nn.Parameter(u)
-        self.svh = nn.Parameter((s * v).T)
-
-        if model.bias is not None:
-            self.bias = nn.Parameter(model.bias.data, requires_grad=True)
-        else:
-            self.register_parameter("bias", None)
-
-    def forward(self, x):
-        """
-        Forward pass of the SmileCompressedLinear module.
-
-        Args:
-            x (Tensor): The input tensor.
-
-        Returns:
-            Tensor: The output tensor.
-        """
-        x = F.linear(x, self.svh)
-        x = F.linear(x, self.u, self.bias)
-        return x
-
-
-class SmileMoELinear(nn.Module):
-    @torch.no_grad()
-    def __init__(
-        self,
-        pretrained_model: nn.Linear,
-        finetuned_models: List[nn.Linear],
-        gate_k: int,
-        k: int,
-        top_k: int = 1,
-        full_matrices=True,
-        upscaling_accelerator=None,
-        routing_use_diff=True,
-    ):
-        """
-        Initialize the SmileMoELinear module.
-
-        Args:
-            pretrained_model (nn.Linear): The pretrained linear model.
-            finetuned_models (List[nn.Linear]): A list of fine-tuned linear models.
-            gate_k (int): The number of singular values to keep for the gate.
-            k (int): The number of singular values to keep for the experts.
-            top_k (int): The number of top experts to select.
-            full_matrices (bool): Whether to compute the full-sized U and V matrices.
-            upscaling_accelerator (str): The device to perform the computation on.
-            routing_use_diff (bool): Whether to use weight differences for routing.
-        """
-        super().__init__()
-        self.num_experts = len(finetuned_models)
-        self.top_k = top_k
-        self.k = k
-        self.gate_k = gate_k
-        self.in_features = pretrained_model.in_features
-        self.out_features = pretrained_model.out_features
-
-        w_diff_list = [m.weight - pretrained_model.weight for m in finetuned_models]
-        if _is_all_zeros(w_diff_list):
-            # All fine-tuned models are identical to the pretrained model
-            raise ExpertNotTrainedError()
-
-        if routing_use_diff or k > 0:
-            svd_cache_list = [
-                svd(w, full_matrices=full_matrices, accelerator=upscaling_accelerator)
-                for w in w_diff_list
-            ]  # the svd cache list to avoid recomputing
-
-        # construct the gate network
-        if routing_use_diff:
-            self.gate = SmileGate(
-                input_features=self.in_features,
-                w_diff_list=w_diff_list,
-                k=gate_k,
-                svd_list=svd_cache_list,
-                upscaling_accelerator=upscaling_accelerator,
-            )
-        else:
-            self.gate = SmileGate(
-                input_features=self.in_features,
-                w_diff_list=[m.weight for m in finetuned_models],
-                k=gate_k,
-                svd_list=None,
-                upscaling_accelerator=upscaling_accelerator,
-            )
-
-        # construct experts
-        for m, w_diff in zip(finetuned_models, w_diff_list):
-            m.weight.data = w_diff
-        if k > 0:
-            experts = [
-                SmileCompressedLinear(m, k, svd_cache=svd_cache)
-                for m, svd_cache in zip(finetuned_models, svd_cache_list)
-            ]
-        else:
-            # if k is not set (<0), we use the full fine-tuned model
-            experts = finetuned_models
-        self.experts = nn.ModuleList(experts)
-
-        if pretrained_model.bias is not None:
-            for m in experts:
-                m.bias.data = m.bias.data - pretrained_model.bias
-        # assign the pretrained model (the shared part)
-        self.pretrained_model = pretrained_model
-
-    def forward(self, hidden_states: Tensor):
-        """
-        Forward pass of the SmileMoELinear module.
-
-        Args:
-            hidden_states (Tensor): The input tensor.
-
-        Returns:
-            Tensor: The output tensor.
-        """
-        pretrained_out = self.pretrained_model(hidden_states)
-
-        input_shape = hidden_states.size()
-        hidden_states = hidden_states.view(-1, self.in_features)
-
-        router_logits = self.gate(hidden_states)
-        routing_weights = F.softmax(router_logits, dim=1)
-        # sample the expert according to the routing weights
-        routing_weights, selected_experts = torch.topk(
-            routing_weights, self.top_k, dim=-1
-        )
-        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
-
-        final_hidden_states = torch.zeros(
-            (hidden_states.size(0), self.out_features),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
-
-        # One hot encode the selected experts to create an expert mask
-        # this will be used to easily index which expert is going to be sollicitated
-        expert_mask = torch.nn.functional.one_hot(
-            selected_experts, num_classes=self.num_experts
-        ).permute(2, 1, 0)
-
-        # Loop over all available experts in the model and perform the computation on each expert
-        for expert_idx in range(self.num_experts):
-            expert_layer = self.experts[expert_idx]
-            idx, top_x = torch.where(expert_mask[expert_idx])
-
-            # Index the correct hidden states and compute the expert hidden state for
-            # the current expert. We need to make sure to multiply the output hidden
-            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
-            current_state = hidden_states[None, top_x].reshape(-1, self.in_features)
-            if current_state.numel() == 0:
-                continue
-            current_hidden_states = (
-                expert_layer(current_state) * routing_weights[top_x, idx, None]
-            )
-
-            # However `index_add_` only support torch tensors for indexing so we'll use
-            # the `top_x` tensor here.
-            final_hidden_states.index_add_(
-                0, top_x, current_hidden_states.to(hidden_states.dtype)
-            )
-        final_hidden_states = final_hidden_states.reshape(
-            *input_shape[:-1], self.out_features
-        )
-        final_hidden_states = pretrained_out + final_hidden_states
-        return final_hidden_states
-
-    @property
-    def weight(self):
-        """
-        Mimic linear layer. Bacause in some cases, user might indicate the device (or dtype of parameters) of the linear layer using `linear_layer.weight.device`
-        """
-        return self.pretrained_model.weight
-
-    @property
-    def bias(self):
-        return self.pretrained_model.bias
-
-    def __repr__(self):
-        return (
-            f"SingularMoELinear("
-            f"in_features={self.pretrained_model.in_features}, "
-            f"out_features={self.pretrained_model.out_features}, "
-            f"num_experts={self.num_experts}, "
-            f"top_k={self.top_k}, "
-            f"gate_k={self.gate_k}, "
-            f"k={self.k}"
-            f")"
-        )
-
-
 class SmileUpscalingAlgorithm(
     SimpleProfilerMixin,
     BaseAlgorithm,
@@ -442,16 +112,19 @@ class SmileUpscalingAlgorithm(
             print_parameters(model)
             return model
 
-        with self.profile("load pretrained model"):
-            pretrained_model = modelpool.load_model("_pretrained_")
-        with self.profile("load fine-tuned model"):
-            finetuned_models = [
-                m for m in tqdm(modelpool.models(), total=len(modelpool.model_names))
-            ]
-
-        if self.config.device == "cuda" and torch.cuda.is_available():
-            pretrained_model = pretrained_model.cuda()
-            finetuned_models = [m.cuda() for m in finetuned_models]
+        with self.profile("loading model"):
+            # load models and move to GPU if available
+            with self.profile("load pretrained model"):
+                pretrained_model = modelpool.load_model("_pretrained_")
+            with self.profile("load fine-tuned model"):
+                finetuned_models = [
+                    m
+                    for m in tqdm(modelpool.models(), total=len(modelpool.model_names))
+                ]
+
+            if self.config.device == "cuda" and torch.cuda.is_available():
+                pretrained_model = pretrained_model.cuda()
+                finetuned_models = [m.cuda() for m in finetuned_models]
 
         with self.profile("merge model"):
             model = self.merge(pretrained_model, finetuned_models)

fusion_bench/method/surgery/clip_layer_wise_adamerging_surgery.py

@@ -85,7 +85,14 @@ class CLIPLayerWiseAdaMergingSurgeryAlgorithm(
 
         if self.config.weights is not None:
             # skip the test-time adaptation
+            merge_weight: torch.Tensor = torch.load(self.config.weights)
+            module.merge_weight.data = merge_weight.to(
+                device=module.merge_weight.device
+            )
             merged_model = copy.deepcopy(module.merge_and_unload())
+            # setup the zero-shot classification head
+            self.on_test_time_adaptation_start()
+
         else:
             with self.profile("test-time adaptation"):
                 module = self.test_time_adaptation(module)

fusion_bench/method/task_arithmetic/task_arithmetic.py

@@ -6,7 +6,7 @@ http://arxiv.org/abs/2212.04089
 
 import logging
 from copy import deepcopy
-from typing import Dict, List, Mapping, TypeVar, Union  # noqa: F401
+from typing import Dict, List, Mapping, Optional, TypeVar, Union  # noqa: F401
 
 import torch
 from torch import nn
@@ -19,18 +19,18 @@ from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_mul,
     state_dict_sub,
 )
-from fusion_bench.utils.type import StateDictType
+from fusion_bench.utils.type import StateDictType, TorchModelType
 
 log = logging.getLogger(__name__)
 
 
 @torch.no_grad()
 def task_arithmetic_merge(
-    pretrained_model: nn.Module,
-    finetuned_models: List[nn.Module],
+    pretrained_model: TorchModelType,
+    finetuned_models: List[TorchModelType],
     scaling_factor: float,
     inplace: bool = True,
-) -> nn.Module:
+) -> TorchModelType:
     """
     Merges the task vectors from multiple fine-tuned models into a single pre-trained model.
 
@@ -46,15 +46,17 @@ def task_arithmetic_merge(
     """
     if not inplace:
         pretrained_model = deepcopy(pretrained_model)
-    task_vector: StateDictType = None
+    task_vector: Optional[StateDictType] = None
     # Calculate the total task vector
     for model in finetuned_models:
         if task_vector is None:
+            # calculate the task vector for the first model
             task_vector = state_dict_sub(
                 model.state_dict(keep_vars=True),
                 pretrained_model.state_dict(keep_vars=True),
             )
         else:
+            # calculate the task vector for the remaining models
             task_vector = state_dict_add(
                 task_vector,
                 state_dict_sub(

fusion_bench/method/ties_merging/ties_merging.py

@@ -16,6 +16,7 @@ from torch import Tensor, nn
 
 from fusion_bench.compat.modelpool import to_modelpool
 from fusion_bench.method import BaseAlgorithm
+from fusion_bench.mixins import SimpleProfilerMixin
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.type import StateDictType
 
@@ -24,7 +25,7 @@ from .ties_merging_utils import state_dict_to_vector, ties_merging, vector_to_st
 log = logging.getLogger(__name__)
 
 
-class TiesMergingAlgorithm(BaseAlgorithm):
+class TiesMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
     """
     TiesMergingAlgorithm is a class for fusing multiple models using the TIES merging technique.
 
@@ -84,34 +85,38 @@ class TiesMergingAlgorithm(BaseAlgorithm):
         scaling_factor = self.scaling_factor
         threshold = self.threshold
 
-        # Load the pretrained model
-        pretrained_model = modelpool.load_model("_pretrained_")
-
-        # Load the state dicts of the models
-        ft_checks: List[StateDictType] = [
-            modelpool.load_model(model_name).state_dict(keep_vars=True)
-            for model_name in modelpool.model_names
-        ]
-        ptm_check: StateDictType = pretrained_model.state_dict(keep_vars=True)
-
-        # Compute the task vectors
-        flat_ft: Tensor = torch.vstack(
-            [state_dict_to_vector(check, remove_keys) for check in ft_checks]
-        )
-        flat_ptm: Tensor = state_dict_to_vector(ptm_check, remove_keys)
-        tv_flat_checks = flat_ft - flat_ptm
-
-        # Perform TIES Merging
-        merged_tv = ties_merging(
-            tv_flat_checks,
-            reset_thresh=threshold,
-            merge_func=merge_func,
-        )
-        merged_check = flat_ptm + scaling_factor * merged_tv
-        merged_state_dict = vector_to_state_dict(
-            merged_check, ptm_check, remove_keys=remove_keys
-        )
-
-        # Load the merged state dict into the pretrained model
-        pretrained_model.load_state_dict(merged_state_dict)
+        with self.profile("loading models"):
+            # Load the pretrained model
+            pretrained_model = modelpool.load_model("_pretrained_")
+
+            # Load the state dicts of the models
+            ft_checks: List[StateDictType] = [
+                modelpool.load_model(model_name).state_dict(keep_vars=True)
+                for model_name in modelpool.model_names
+            ]
+            ptm_check: StateDictType = pretrained_model.state_dict(keep_vars=True)
+
+        with self.profile("merging models"):
+            # Compute the task vectors
+            flat_ft: Tensor = torch.vstack(
+                [state_dict_to_vector(check, remove_keys) for check in ft_checks]
+            )
+            flat_ptm: Tensor = state_dict_to_vector(ptm_check, remove_keys)
+            tv_flat_checks = flat_ft - flat_ptm
+
+            # Perform TIES Merging
+            merged_tv = ties_merging(
+                tv_flat_checks,
+                reset_thresh=threshold,
+                merge_func=merge_func,
+            )
+            merged_check = flat_ptm + scaling_factor * merged_tv
+            merged_state_dict = vector_to_state_dict(
+                merged_check, ptm_check, remove_keys=remove_keys
+            )
+
+            # Load the merged state dict into the pretrained model
+            pretrained_model.load_state_dict(merged_state_dict)
+
+        self.print_profile_summary()
         return pretrained_model

fusion_bench/method/we_moe/we_moe.py

@@ -5,7 +5,6 @@ from typing import cast  # noqa: F401
 import lightning as L
 import lightning.fabric.wrappers
 import torch
-from lightning.pytorch.profilers import SimpleProfiler
 from omegaconf import DictConfig
 from torch import Tensor
 from torch.utils.data import DataLoader
@@ -13,6 +12,7 @@ from tqdm.autonotebook import tqdm
 
 from fusion_bench.compat.method.base_algorithm import ModelFusionAlgorithm
 from fusion_bench.compat.modelpool import ModelPool
+from fusion_bench.mixins import SimpleProfilerMixin
 from fusion_bench.models.we_moe import WeightEnsemblingMoE
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.parameters import print_parameters
@@ -34,7 +34,10 @@ def entropy_loss(logits: Tensor) -> Tensor:
     return -torch.sum(probs * torch.log(probs + 1e-8), dim=-1).mean()
 
 
-class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
+class WeightEnsemblingMoEAlgorithm(
+    ModelFusionAlgorithm,
+    SimpleProfilerMixin,
+):
     """
     Algorithm for fusing models using Weight Ensembling Mixture of Experts (MoE).
 
@@ -44,7 +47,6 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
     Attributes:
         _fabric (L.Fabric): The fabric for distributed training.
         modelpool (ModelPool): The pool of models to be fused.
-        profiler (SimpleProfiler): The profiler for measuring performance.
     """
 
     _fabric: L.Fabric = None
@@ -66,9 +68,6 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
             self._fabric.launch()
         else:
             assert "No CUDA device available."
-        self.profiler = SimpleProfiler(
-            self.config.get("cache_dir", "outputs"), "we_moe_profiler.txt"
-        )
 
     @abstractmethod
     def load_checkpoint(self, model, checkpoint):
@@ -177,9 +176,9 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
         for step_idx in pbar:
             if self.config.use_grad_accumulate:
                 for task in self.modelpool.model_names:
-                    with self.profiler.profile("data time"):
+                    with self.profile("data time"):
                         batch = next(self.get_shuffled_test_loader_iter(task))
-                    with self.profiler.profile("forward pass"):
+                    with self.profile("forward pass"):
                         logits = self.compute_logits(module, batch, task)
                         assert (
                             logits.dim() == 2
@@ -187,23 +186,23 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
                         loss = entropy_loss(logits)
                     # .backward() accumulates when .zero_grad() wasn't called
                     # this can save memory
-                    with self.profiler.profile("backward pass"):
+                    with self.profile("backward pass"):
                         self._fabric.backward(loss, retain_graph=True)
             else:
                 loss = 0
                 for task in self.modelpool.model_names:
-                    with self.profiler.profile("data time"):
+                    with self.profile("data time"):
                         batch = next(self.get_shuffled_test_loader_iter(task))
-                    with self.profiler.profile("forward pass"):
+                    with self.profile("forward pass"):
                         logits = self.compute_logits(module, batch, task)
                         assert (
                             logits.dim() == 2
                         ), f"Expected logits to be 2D, got {logits.dim()}"
                         loss = loss + entropy_loss(logits)
-                with self.profiler.profile("backward pass"):
+                with self.profile("backward pass"):
                     self._fabric.backward(loss, retain_graph=True)
 
-            with self.profiler.profile("optimizer step"):
+            with self.profile("optimizer step"):
                 optimizer.step()
                 optimizer.zero_grad()
 
@@ -232,7 +231,7 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
             )
             self.load_checkpoint(moe_model, self.config.checkpoint)
         else:
-            with self.profiler.profile("test-time adaptation"):
+            with self.profile("test-time adaptation"):
                 moe_model = self.test_time_adaptation(moe_model)
             if self.config.get("save_checkpoint", False):
                 log.info(f"save checkpoint to {self.config.save_checkpoint}")
@@ -243,5 +242,5 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
 
         # enable sample-wise adaptation
         moe_model.batch_reduce = False
-        print(self.profiler.summary())
+        self.print_profile_summary()
         return moe_model

fusion_bench/mixins/__init__.py

@@ -6,20 +6,23 @@ from typing_extensions import TYPE_CHECKING
 from fusion_bench.utils.lazy_imports import LazyImporter
 
 _import_structure = {
+    "clip_classification": ["CLIPClassificationMixin"],
+    "fabric_training": ["FabricTrainingMixin"],
+    "hydra_config": ["HydraConfigMixin"],
     "lightning_fabric": ["LightningFabricMixin"],
+    "openclip_classification": ["OpenCLIPClassificationMixin"],
     "serialization": ["YAMLSerializationMixin", "BaseYAMLSerializableModel"],
     "simple_profiler": ["SimpleProfilerMixin"],
-    "clip_classification": ["CLIPClassificationMixin"],
-    "fabric_training": ["FabricTrainingMixin"],
 }
 
 if TYPE_CHECKING:
     from .clip_classification import CLIPClassificationMixin
     from .fabric_training import FabricTrainingMixin
+    from .hydra_config import HydraConfigMixin
     from .lightning_fabric import LightningFabricMixin
+    from .openclip_classification import OpenCLIPClassificationMixin
     from .serialization import BaseYAMLSerializableModel, YAMLSerializationMixin
     from .simple_profiler import SimpleProfilerMixin
-
 else:
     sys.modules[__name__] = LazyImporter(
         __name__,