refactor: 重构流水线架构，添加Pipeline抽象并拆分IOHandler

pipeline/processors/__init__.py

@@ -4,15 +4,18 @@ 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.base import BaseProcessor, ProcessorSchema, encode_with_mask
+from pipeline.processors.factory import ProcessorFactory, ProcessorConfig
 from pipeline.processors.pretrain import PreTrainProcessor
 from pipeline.processors.sft import SFTProcessor
 from pipeline.processors.dpo import DPOProcessor
 
 __all__ = [
     "BaseProcessor",
+    "ProcessorSchema",
+    "ProcessorConfig",
     "ProcessorFactory",
+    "encode_with_mask",
     "PreTrainProcessor",
     "SFTProcessor",
     "DPOProcessor",

pipeline/processors/base.py

@@ -1,32 +1,126 @@
 """Processor base class and shared utilities."""
 
 from abc import ABC, abstractmethod
-from typing import Dict, List, Any, Tuple
+from dataclasses import dataclass
+from typing import Any, Dict, List, 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
+@dataclass(frozen=True)
+class ProcessorSchema:
+    """Schema definition for processor input/output contracts.
+
+    Attributes:
+        input_fields: Required input field names and their types.
+        output_fields: Output field names and their tensor dtypes.
+    """
+
+    input_fields: Dict[str, type]
+    output_fields: Dict[str, torch.dtype]
 
 
 class BaseProcessor(ABC):
-    """Abstract base class for processors."""
+    """Abstract base class for data processors.
+
+    Processors transform raw data (e.g., text, JSON) into tokenized tensors
+    suitable for model training.
+
+    Subclasses must implement:
+    - process(): Transform a single input sample
+    - output_keys: Declare output tensor names
+    - schema: Define input/output contracts (optional but recommended)
+
+    Example::
+
+        class MyProcessor(BaseProcessor):
+            @property
+            def schema(self) -> ProcessorSchema:
+                return ProcessorSchema(
+                    input_fields={"text": str},
+                    output_fields={"tokens": torch.int32}
+                )
+
+            def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+                tokens = self.tokenizer.encode(input_dict["text"])
+                return {"tokens": torch.tensor(tokens, dtype=torch.int32)}
+
+            @property
+            def output_keys(self) -> List[str]:
+                return ["tokens"]
+    """
+
+    @property
+    def schema(self) -> ProcessorSchema:
+        """Return the input/output schema for this processor.
+
+        Override this property to define explicit contracts.
+        Default returns None, meaning schema is not defined.
+        """
+        return None
 
     @abstractmethod
     def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        """Process a single input sample.
+
+        Args:
+            input_dict: Dictionary containing input fields as defined by schema.
+
+        Returns:
+            Dictionary mapping output key names to tensors.
+
+        Raises:
+            KeyError: If required input fields are missing.
+            ValueError: If input data is invalid.
+        """
         pass
 
     @property
     @abstractmethod
     def output_keys(self) -> List[str]:
+        """Return list of output tensor key names."""
         pass
+
+    def validate_input(self, input_dict: Dict[str, Any]) -> None:
+        """Validate input against schema before processing.
+
+        Args:
+            input_dict: Input dictionary to validate.
+
+        Raises:
+            KeyError: If required fields are missing.
+            TypeError: If field types don't match schema.
+        """
+        schema = self.schema
+        if schema is None:
+            return
+
+        for field_name, expected_type in schema.input_fields.items():
+            if field_name not in input_dict:
+                raise KeyError(f"Missing required input field: '{field_name}'")
+            if not isinstance(input_dict[field_name], expected_type):
+                raise TypeError(
+                    f"Field '{field_name}' expected type {expected_type.__name__}, "
+                    f"got {type(input_dict[field_name]).__name__}"
+                )
+
+
+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).
+
+    Args:
+        prompt_tokens: Token IDs for the prompt/question.
+        response_tokens: Token IDs for the response/answer.
+
+    Returns:
+        Tuple of (combined_tokens, loss_mask) where loss_mask is True for response tokens.
+    """
+    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

pipeline/processors/dpo.py

@@ -1,21 +1,32 @@
 """DPO preference learning data processor."""
 
-from typing import Dict, List, Any, Optional
+from typing import Any, Dict, List, Optional
 
 import torch
 from torch import Tensor
 
 from pipeline.tokenize import AutoTokenizer
 from pipeline.strategies import PromptStrategy, ChatMLStrategy
-from pipeline.processors.base import BaseProcessor, _encode_with_mask
+from pipeline.processors.base import BaseProcessor, ProcessorSchema, encode_with_mask
 from pipeline.processors.factory import ProcessorFactory
 
 
 @ProcessorFactory.register("dpo")
 class DPOProcessor(BaseProcessor):
-    """DPO preference learning data processor.
+    """DPO (Direct Preference Optimization) data processor.
 
-    Supports custom prompt strategy via constructor parameter.
+    Processes query, chosen, and rejected responses for preference learning.
+
+    Input schema:
+        - query: str - User query/prompt
+        - chosen: str - Preferred assistant response
+        - rejected: str - Dispreferred assistant response
+
+    Output schema:
+        - chosen: int32 tensor - Token IDs for preferred response
+        - chosen_mask: bool tensor - True for response tokens
+        - rejected: int32 tensor - Token IDs for dispreferred response
+        - rejected_mask: bool tensor - True for response tokens
     """
 
     def __init__(
@@ -26,6 +37,22 @@ class DPOProcessor(BaseProcessor):
         self.tokenizer = tokenizer
         self.strategy = strategy or ChatMLStrategy(tokenizer)
 
+    @property
+    def schema(self) -> ProcessorSchema:
+        return ProcessorSchema(
+            input_fields={
+                "query": str,
+                "chosen": str,
+                "rejected": str,
+            },
+            output_fields={
+                "chosen": torch.int32,
+                "chosen_mask": torch.bool,
+                "rejected": torch.int32,
+                "rejected_mask": torch.bool,
+            },
+        )
+
     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"])
@@ -33,10 +60,10 @@ class DPOProcessor(BaseProcessor):
 
         prompt = self.strategy.assemble_prompt(query_tokens)
 
-        chosen_t, chosen_m = _encode_with_mask(
+        chosen_t, chosen_m = encode_with_mask(
             prompt, self.strategy.assemble_response(chosen_tokens)
         )
-        rejected_t, rejected_m = _encode_with_mask(
+        rejected_t, rejected_m = encode_with_mask(
             prompt, self.strategy.assemble_response(rejected_tokens)
         )
 

pipeline/processors/factory.py

@@ -1,12 +1,32 @@
-"""Factory for creating and registering processors."""
+"""Factory for creating and registering processors with unified interface."""
 
-from typing import Dict, List, Any, Optional, Type
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Type, Union
 
 from pipeline.processors.base import BaseProcessor
 from pipeline.tokenize import AutoTokenizer
 from pipeline.strategies import PromptStrategy, StrategyFactory
 
 
+@dataclass
+class ProcessorConfig:
+    """Configuration for creating a processor.
+
+    Attributes:
+        processor_type: Type name for the processor ("pt", "sft", "dpo").
+        tokenizer: Tokenizer instance (required).
+        strategy_name: Name of the strategy to use (optional).
+        strategy: Pre-created strategy instance (optional).
+        strategy_kwargs: Additional arguments for strategy creation.
+    """
+
+    processor_type: str
+    tokenizer: AutoTokenizer
+    strategy_name: Optional[str] = None
+    strategy: Optional[PromptStrategy] = None
+    strategy_kwargs: Optional[Dict] = None
+
+
 class ProcessorFactory:
     """Registry and factory for BaseProcessor implementations.
 
@@ -18,7 +38,16 @@ class ProcessorFactory:
         class CustomProcessor(BaseProcessor):
             ...
 
-        processor = ProcessorFactory.create(optimizer, "custom", **kwargs)
+        # Using config object (recommended)
+        config = ProcessorConfig(
+            processor_type="sft",
+            tokenizer=tokenizer,
+            strategy_name="alpaca"
+        )
+        processor = ProcessorFactory.create_from_config(config)
+
+        # Using direct arguments
+        processor = ProcessorFactory.create("pt", tokenizer)
     """
 
     PROCESSOR_MAP: Dict[str, Type[BaseProcessor]] = {}
@@ -46,10 +75,10 @@ class ProcessorFactory:
 
     @classmethod
     def create(cls, processor_type: str, tokenizer: AutoTokenizer) -> BaseProcessor:
-        """Create a processor by type name (uses default ChatMLStrategy for SFT/DPO).
+        """Create a processor by type name.
 
         Args:
-            processor_type: Registered processor name (e.g. ``"pt"``, ``"sft"``, ``"dpo"``).
+            processor_type: Registered processor name (e.g. "pt", "sft", "dpo").
             tokenizer: Tokenizer instance.
 
         Returns:
@@ -72,14 +101,12 @@ class ProcessorFactory:
         tokenizer: AutoTokenizer,
         strategy: PromptStrategy,
     ) -> BaseProcessor:
-        """Create a processor with a custom strategy.
-
-        Only SFT and DPO processors accept a strategy; PreTrain ignores it.
+        """Create a processor with a pre-configured strategy.
 
         Args:
             processor_type: Registered processor name.
             tokenizer: Tokenizer instance.
-            strategy: Prompt strategy instance.
+            strategy: Pre-created strategy instance.
 
         Returns:
             Processor instance configured with strategy.
@@ -108,8 +135,8 @@ class ProcessorFactory:
         Args:
             processor_type: Registered processor name.
             tokenizer: Tokenizer instance.
-            strategy_name: Registered strategy name (``"chatml"``, ``"alpaca"``, etc.).
-            **strategy_kwargs: Forwarded to the strategy constructor.
+            strategy_name: Registered strategy name ("chatml", "alpaca", etc.).
+            **strategy_kwargs: Additional arguments forwarded to strategy constructor.
 
         Returns:
             Processor instance.
@@ -117,6 +144,44 @@ class ProcessorFactory:
         strategy = StrategyFactory.create(strategy_name, tokenizer, **strategy_kwargs)
         return cls.create_with_strategy(processor_type, tokenizer, strategy)
 
+    @classmethod
+    def create_from_config(cls, config: ProcessorConfig) -> BaseProcessor:
+        """Create a processor from a configuration object (unified interface).
+
+        Args:
+            config: ProcessorConfig with all creation parameters.
+
+        Returns:
+            Processor instance.
+
+        Raises:
+            ValueError: If processor_type is not registered or strategy is invalid.
+        """
+        if config.processor_type not in cls.PROCESSOR_MAP:
+            raise ValueError(
+                f"Unknown processor type: '{config.processor_type}'. "
+                f"Supported types: {sorted(cls.PROCESSOR_MAP.keys())}"
+            )
+
+        tokenizer = config.tokenizer
+        strategy_kwargs = config.strategy_kwargs or {}
+
+        # Determine strategy to use
+        strategy: Optional[PromptStrategy] = None
+        if config.strategy is not None:
+            strategy = config.strategy
+        elif config.strategy_name is not None:
+            strategy = StrategyFactory.create(
+                config.strategy_name, tokenizer, **strategy_kwargs
+            )
+
+        # Create processor
+        if strategy is not None:
+            return cls.create_with_strategy(
+                config.processor_type, tokenizer, strategy
+            )
+        return cls.create(config.processor_type, tokenizer)
+
     @classmethod
     def available_types(cls) -> List[str]:
         """Return list of registered processor type names."""

pipeline/processors/pretrain.py

@@ -1,25 +1,46 @@
 """Pre-training data processor."""
 
-from typing import Dict, List, Any
+from typing import Any, Dict, List
 
 import torch
 from torch import Tensor
 
 from pipeline.tokenize import AutoTokenizer
-from pipeline.processors.base import BaseProcessor
+from pipeline.processors.base import BaseProcessor, ProcessorSchema
 from pipeline.processors.factory import ProcessorFactory
 
 
 @ProcessorFactory.register("pt")
 class PreTrainProcessor(BaseProcessor):
-    """Pre-training data processor."""
+    """Pre-training data processor.
 
-    def __init__(self, tokenizer: AutoTokenizer):
+    Processes raw text into tokenized sequences with EOS tokens.
+
+    Input schema:
+        - text: str - Raw text string to tokenize
+
+    Output schema:
+        - sequence: int32 tensor - Token IDs with EOS appended
+    """
+
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        eos_token: str = "<｜end▁of▁sentence｜>",
+    ):
         self.tokenizer = tokenizer
+        self._eos_token = eos_token
+
+    @property
+    def schema(self) -> ProcessorSchema:
+        return ProcessorSchema(
+            input_fields={"text": str},
+            output_fields={"sequence": torch.int32},
+        )
 
     def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
         segment = input_dict["text"]
-        tokens = self.tokenizer.encode(f"{segment}<｜end▁of▁sentence｜>")
+        tokens = self.tokenizer.encode(f"{segment}{self._eos_token}")
         return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
 
     @property

pipeline/processors/sft.py

@@ -1,13 +1,13 @@
 """Supervised fine-tuning data processor."""
 
-from typing import Dict, List, Any, Optional
+from typing import Any, Dict, List, Optional
 
 import torch
 from torch import Tensor
 
 from pipeline.tokenize import AutoTokenizer
 from pipeline.strategies import PromptStrategy, ChatMLStrategy
-from pipeline.processors.base import BaseProcessor, _encode_with_mask
+from pipeline.processors.base import BaseProcessor, ProcessorSchema, encode_with_mask
 from pipeline.processors.factory import ProcessorFactory
 
 
@@ -15,7 +15,15 @@ from pipeline.processors.factory import ProcessorFactory
 class SFTProcessor(BaseProcessor):
     """Supervised fine-tuning data processor.
 
-    Supports custom prompt strategy via constructor parameter.
+    Processes query-response pairs into tokenized sequences with loss masks.
+
+    Input schema:
+        - query: str - User query/prompt
+        - response: str - Assistant response
+
+    Output schema:
+        - sequence: int32 tensor - Combined token IDs (query + response)
+        - loss_mask: bool tensor - True for response tokens (compute loss)
     """
 
     def __init__(
@@ -26,6 +34,16 @@ class SFTProcessor(BaseProcessor):
         self.tokenizer = tokenizer
         self.strategy = strategy or ChatMLStrategy(tokenizer)
 
+    @property
+    def schema(self) -> ProcessorSchema:
+        return ProcessorSchema(
+            input_fields={"query": str, "response": str},
+            output_fields={
+                "sequence": torch.int32,
+                "loss_mask": torch.bool,
+            },
+        )
+
     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"])
@@ -33,7 +51,7 @@ class SFTProcessor(BaseProcessor):
         prompt = self.strategy.assemble_prompt(query_tokens)
         response = self.strategy.assemble_response(response_tokens)
 
-        tokens, loss_mask = _encode_with_mask(prompt, response)
+        tokens, loss_mask = encode_with_mask(prompt, response)
         return {"sequence": tokens, "loss_mask": loss_mask}
 
     @property