MTL Config Union -> dataclass

adapter-hub · Feb 15, 2025 · cbda47a · cbda47a
1 parent 9efc666
commit cbda47a
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/src/adapters/configuration/adapter_config.py b/src/adapters/configuration/adapter_config.py
@@ -551,14 +551,11 @@ class MTLLoRAConfig(LoRAConfig, MultiTaskConfig):
     weights_sharpness: float = 0.05
 
 
+@dataclass
 class MultiTaskConfigUnion(AdapterConfig):
     architecture: Optional[str] = "mtl_union"
-    base_config: MultiTaskConfig
-    task_names: List[str]
-
-    def __init__(self, base_config: MultiTaskConfig, task_names: List[str]):
-        self.base_config = base_config
-        self.task_names = task_names
+    base_config: Optional[MultiTaskConfig] = None
+    task_names: Optional[List[str]] = None
 
     def to_dict(self):
         return {