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

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