reafactor: 重构项目

pipeline/processors/__init__.py

@@ -0,0 +1,18 @@
+"""Data processors with factory pattern.
+
+Processor classes are registered at definition time via decorators and
+can be created through :class:`ProcessorFactory`.
+"""
+from pipeline.processors.base import BaseProcessor
+from pipeline.processors.factory import ProcessorFactory
+from pipeline.processors.pretrain import PreTrainProcessor
+from pipeline.processors.sft import SFTProcessor
+from pipeline.processors.dpo import DPOProcessor
+
+__all__ = [
+    "BaseProcessor",
+    "ProcessorFactory",
+    "PreTrainProcessor",
+    "SFTProcessor",
+    "DPOProcessor",
+]

pipeline/processors/base.py

@@ -0,0 +1,31 @@
+"""Processor base class and shared utilities."""
+from abc import ABC, abstractmethod
+from typing import Dict, List, Any, Tuple
+
+import torch
+from torch import Tensor
+
+
+def _encode_with_mask(
+    prompt_tokens: List[int],
+    response_tokens: List[int],
+) -> Tuple[Tensor, Tensor]:
+    """Concatenate token lists and build loss mask (prompt=False, response=True)."""
+    q_len = len(prompt_tokens)
+    combined = torch.tensor(prompt_tokens + response_tokens, dtype=torch.int32)
+    mask = torch.zeros(q_len + len(response_tokens), dtype=torch.bool)
+    mask[q_len:] = True
+    return combined, mask
+
+
+class BaseProcessor(ABC):
+    """Abstract base class for processors."""
+
+    @abstractmethod
+    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        pass
+
+    @property
+    @abstractmethod
+    def output_keys(self) -> List[str]:
+        pass

pipeline/processors/dpo.py

@@ -0,0 +1,51 @@
+"""DPO preference learning data processor."""
+from typing import Dict, List, Any, Optional
+
+import torch
+from torch import Tensor
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.strategies import PromptStrategy, ChatMLStrategy
+from pipeline.processors.base import BaseProcessor, _encode_with_mask
+from pipeline.processors.factory import ProcessorFactory
+
+
+@ProcessorFactory.register("dpo")
+class DPOProcessor(BaseProcessor):
+    """DPO preference learning data processor.
+
+    Supports custom prompt strategy via constructor parameter.
+    """
+
+    def __init__(
+        self,
+        tokenizer: BpeTokenizer,
+        strategy: Optional[PromptStrategy] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.strategy = strategy or ChatMLStrategy(tokenizer)
+
+    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        query_tokens = self.tokenizer.encode(input_dict["query"])
+        chosen_tokens = self.tokenizer.encode(input_dict["chosen"])
+        rejected_tokens = self.tokenizer.encode(input_dict["rejected"])
+
+        prompt = self.strategy.assemble_prompt(query_tokens)
+
+        chosen_t, chosen_m = _encode_with_mask(
+            prompt, self.strategy.assemble_response(chosen_tokens)
+        )
+        rejected_t, rejected_m = _encode_with_mask(
+            prompt, self.strategy.assemble_response(rejected_tokens)
+        )
+
+        return {
+            "chosen": chosen_t,
+            "chosen_mask": chosen_m,
+            "rejected": rejected_t,
+            "rejected_mask": rejected_m,
+        }
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["chosen", "chosen_mask", "rejected", "rejected_mask"]

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())

pipeline/processors/pretrain.py

@@ -0,0 +1,26 @@
+"""Pre-training data processor."""
+from typing import Dict, List, Any
+
+import torch
+from torch import Tensor
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.processors.base import BaseProcessor
+from pipeline.processors.factory import ProcessorFactory
+
+
+@ProcessorFactory.register("pt")
+class PreTrainProcessor(BaseProcessor):
+    """Pre-training data processor."""
+
+    def __init__(self, tokenizer: BpeTokenizer):
+        self.tokenizer = tokenizer
+
+    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        segment = input_dict["text"]
+        tokens = self.tokenizer.encode(f"{segment}<eos>")
+        return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["sequence"]

pipeline/processors/sft.py

@@ -0,0 +1,40 @@
+"""Supervised fine-tuning data processor."""
+from typing import Dict, List, Any, Optional
+
+import torch
+from torch import Tensor
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.strategies import PromptStrategy, ChatMLStrategy
+from pipeline.processors.base import BaseProcessor, _encode_with_mask
+from pipeline.processors.factory import ProcessorFactory
+
+
+@ProcessorFactory.register("sft")
+class SFTProcessor(BaseProcessor):
+    """Supervised fine-tuning data processor.
+
+    Supports custom prompt strategy via constructor parameter.
+    """
+
+    def __init__(
+        self,
+        tokenizer: BpeTokenizer,
+        strategy: Optional[PromptStrategy] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.strategy = strategy or ChatMLStrategy(tokenizer)
+
+    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        query_tokens = self.tokenizer.encode(input_dict["query"])
+        response_tokens = self.tokenizer.encode(input_dict["response"])
+
+        prompt = self.strategy.assemble_prompt(query_tokens)
+        response = self.strategy.assemble_response(response_tokens)
+
+        tokens, loss_mask = _encode_with_mask(prompt, response)
+        return {"sequence": tokens, "loss_mask": loss_mask}
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["sequence", "loss_mask"]