reafactor: 重构项目

pipeline/processors/factory.py

@@ -0,0 +1,122 @@
+"""Factory for creating and registering processors."""
+from typing import Dict, List, Any, Optional, Type
+
+from pipeline.processors.base import BaseProcessor
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.strategies import PromptStrategy, StrategyFactory
+
+
+class ProcessorFactory:
+    """Registry and factory for BaseProcessor implementations.
+
+    Supports decorator-based registration for extensible processor types.
+
+    Example usage::
+
+        @ProcessorFactory.register("custom")
+        class CustomProcessor(BaseProcessor):
+            ...
+
+        processor = ProcessorFactory.create(optimizer, "custom", **kwargs)
+    """
+
+    PROCESSOR_MAP: Dict[str, Type[BaseProcessor]] = {}
+
+    @classmethod
+    def register(cls, name: str):
+        """Decorator to register a new processor class.
+
+        Args:
+            name: Registration name for the processor.
+
+        Returns:
+            Decorator function that registers the processor class.
+        """
+
+        def decorator(processor_cls: Type[BaseProcessor]) -> Type[BaseProcessor]:
+            if not issubclass(processor_cls, BaseProcessor):
+                raise TypeError(
+                    f"{processor_cls.__name__} must inherit from BaseProcessor"
+                )
+            cls.PROCESSOR_MAP[name] = processor_cls
+            return processor_cls
+
+        return decorator
+
+    @classmethod
+    def create(cls, processor_type: str, tokenizer: BpeTokenizer) -> BaseProcessor:
+        """Create a processor by type name (uses default ChatMLStrategy for SFT/DPO).
+
+        Args:
+            processor_type: Registered processor name (e.g. ``"pt"``, ``"sft"``, ``"dpo"``).
+            tokenizer: Tokenizer instance.
+
+        Returns:
+            Processor instance.
+
+        Raises:
+            ValueError: If processor_type is not registered.
+        """
+        if processor_type not in cls.PROCESSOR_MAP:
+            raise ValueError(
+                f"Unknown processor type: '{processor_type}'. "
+                f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
+            )
+        return cls.PROCESSOR_MAP[processor_type](tokenizer)
+
+    @classmethod
+    def create_with_strategy(
+        cls,
+        processor_type: str,
+        tokenizer: BpeTokenizer,
+        strategy: PromptStrategy,
+    ) -> BaseProcessor:
+        """Create a processor with a custom strategy.
+
+        Only SFT and DPO processors accept a strategy; PreTrain ignores it.
+
+        Args:
+            processor_type: Registered processor name.
+            tokenizer: Tokenizer instance.
+            strategy: Prompt strategy instance.
+
+        Returns:
+            Processor instance configured with strategy.
+        """
+        if processor_type not in cls.PROCESSOR_MAP:
+            raise ValueError(
+                f"Unknown processor type: '{processor_type}'. "
+                f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
+            )
+
+        processor_cls = cls.PROCESSOR_MAP[processor_type]
+        if processor_type == "pt":
+            return processor_cls(tokenizer)
+        return processor_cls(tokenizer, strategy=strategy)
+
+    @classmethod
+    def create_with_strategy_name(
+        cls,
+        processor_type: str,
+        tokenizer: BpeTokenizer,
+        strategy_name: str,
+        **strategy_kwargs,
+    ) -> BaseProcessor:
+        """Create a processor with a strategy selected by name.
+
+        Args:
+            processor_type: Registered processor name.
+            tokenizer: Tokenizer instance.
+            strategy_name: Registered strategy name (``"chatml"``, ``"alpaca"``, etc.).
+            **strategy_kwargs: Forwarded to the strategy constructor.
+
+        Returns:
+            Processor instance.
+        """
+        strategy = StrategyFactory.create(strategy_name, tokenizer, **strategy_kwargs)
+        return cls.create_with_strategy(processor_type, tokenizer, strategy)
+
+    @classmethod
+    def available_types(cls) -> List[str]:
+        """Return list of registered processor type names."""
+        return list(cls.PROCESSOR_MAP.keys())