refactor: 优化工厂模式结构

astrai/trainer/schedule.py

@@ -6,6 +6,8 @@ from typing import Any, Dict, List, Type
 
 from torch.optim.lr_scheduler import LRScheduler
 
+from astrai.core.factory import BaseFactory
+
 
 class BaseScheduler(LRScheduler, ABC):
     """Base scheduler class for all other schedulers."""
@@ -25,7 +27,7 @@ class BaseScheduler(LRScheduler, ABC):
         super().load_state_dict(state_dict)
 
 
-class SchedulerFactory:
+class SchedulerFactory(BaseFactory["BaseScheduler"]):
     """Factory class for creating learning rate schedulers.
 
     Supports decorator-based registration for extensible scheduler types.
@@ -36,34 +38,21 @@ class SchedulerFactory:
         class CustomScheduler(BaseScheduler):
             ...
 
-        scheduler = SchedulerFactory.create(optimizer, "custom", **kwargs)
+        scheduler = SchedulerFactory.create("custom", optimizer, **kwargs)
     """
 
-    SCHEDULER_MAP: Dict[str, Type[BaseScheduler]] = {}
+    _registry: Dict[str, Type[BaseScheduler]] = {}
 
     @classmethod
-    def register(cls, name: str):
-        """Decorator to register a new scheduler class.
-
-        Args:
-            name: Registration name for the scheduler
-
-        Returns:
-            Decorator function that registers the scheduler class
-        """
-
-        def decorator(scheduler_cls: Type[BaseScheduler]) -> Type[BaseScheduler]:
-            if not issubclass(scheduler_cls, BaseScheduler):
-                raise TypeError(
-                    f"{scheduler_cls.__name__} must inherit from BaseScheduler"
-                )
-            cls.SCHEDULER_MAP[name] = scheduler_cls
-            return scheduler_cls
-
-        return decorator
+    def _validate_component(cls, scheduler_cls: Type[BaseScheduler]) -> None:
+        """Validate that the scheduler class inherits from BaseScheduler."""
+        if not issubclass(scheduler_cls, BaseScheduler):
+            raise TypeError(f"{scheduler_cls.__name__} must inherit from BaseScheduler")
 
     @classmethod
-    def create(cls, optimizer, schedule_type: str = "none", **kwargs) -> BaseScheduler:
+    def create(
+        cls, optimizer, schedule_type: str = "none", **kwargs
+    ) -> "BaseScheduler":
         """Create a scheduler instance by type name.
 
         Args:
@@ -73,23 +62,13 @@ class SchedulerFactory:
 
         Returns:
             Scheduler instance
-
-        Raises:
-            ValueError: If schedule_type is not supported
         """
-        if schedule_type not in cls.SCHEDULER_MAP:
-            raise ValueError(
-                f"Unknown schedule type: '{schedule_type}'. "
-                f"Supported types: {sorted(cls.SCHEDULER_MAP.keys())}"
-            )
-
-        scheduler_cls = cls.SCHEDULER_MAP[schedule_type]
-        return scheduler_cls(optimizer, **kwargs)
+        return super().create(schedule_type, optimizer, **kwargs)
 
     @classmethod
     def available_types(cls) -> list:
         """Return list of registered scheduler type names."""
-        return list(cls.SCHEDULER_MAP.keys())
+        return cls.list_registered()
 
 
 # ----------- Scheduler implementations -----------

astrai/trainer/strategy.py

@@ -10,6 +10,8 @@ import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.parallel import DistributedDataParallel as DDP
 
+from astrai.core.factory import BaseFactory
+
 
 def unwrap_model(model: nn.Module) -> nn.Module:
     """Unwrap DDP wrapper if present to get the original model."""
@@ -106,7 +108,7 @@ class BaseStrategy(ABC):
         return self.compute_loss(batch)
 
 
-class StrategyFactory:
+class StrategyFactory(BaseFactory["BaseStrategy"]):
     """Factory class for creating training strategy instances.
 
     Supports decorator-based registration for extensible strategy types.
@@ -117,68 +119,36 @@ class StrategyFactory:
         class CustomStrategy(BaseStrategy):
             ...
 
-        strategy = StrategyFactory.create(model, "custom", device)
+        strategy = StrategyFactory.create("custom", model, device)
     """
 
-    SUPPORTED_STRATEGIES = frozenset({"seq", "sft", "dpo", "grpo"})
-    STRATEGY_MAP: Dict[str, type] = {}
+    _registry: Dict[str, type] = {}
 
     @classmethod
-    def register(cls, name: str):
-        """Decorator to register a new strategy class.
-
-        Args:
-            name: Registration name for the strategy
-
-        Returns:
-            Decorator function that registers the strategy class
-        """
-
-        def decorator(strategy_cls: type) -> type:
-            if not issubclass(strategy_cls, BaseStrategy):
-                raise TypeError(
-                    f"{strategy_cls.__name__} must inherit from BaseStrategy"
-                )
-            cls.STRATEGY_MAP[name] = strategy_cls
-            return strategy_cls
-
-        return decorator
+    def _validate_component(cls, strategy_cls: type) -> None:
+        """Validate that the strategy class inherits from BaseStrategy."""
+        if not issubclass(strategy_cls, BaseStrategy):
+            raise TypeError(f"{strategy_cls.__name__} must inherit from BaseStrategy")
 
     @classmethod
-    def create(cls, model, train_type: str, device: str, **kwargs) -> BaseStrategy:
+    def create(cls, train_type: str, model, device: str, **kwargs) -> "BaseStrategy":
         """Create a strategy instance based on training type.
 
         Args:
-            model: Model instance for the strategy
             train_type: Type of training ("seq", "sft", "dpo", "grpo")
+            model: Model instance for the strategy
             device: Device to run the strategy on
             **kwargs: Additional arguments passed to strategy constructor
 
         Returns:
             Strategy instance
-
-        Raises:
-            ValueError: If train_type is not supported
-            NotImplementedError: If train_type is in supported list but not implemented
         """
-        if train_type not in cls.SUPPORTED_STRATEGIES:
-            raise ValueError(
-                f"Unknown training strategy: '{train_type}'. "
-                f"Supported strategies: {sorted(cls.SUPPORTED_STRATEGIES)}"
-            )
-
-        if train_type not in cls.STRATEGY_MAP:
-            raise NotImplementedError(
-                f"Strategy '{train_type}' is supported but not yet implemented."
-            )
-
-        strategy_cls = cls.STRATEGY_MAP[train_type]
-        return strategy_cls(model, device, **kwargs)
+        return super().create(train_type, model, device, **kwargs)
 
     @classmethod
     def available_strategies(cls) -> list:
         """Return list of registered strategy names."""
-        return list(cls.STRATEGY_MAP.keys())
+        return cls.list_registered()
 
 
 # ============== Strategy Classes ==============