reafactor: 重构项目

pipeline/strategies/__init__.py

@@ -0,0 +1,14 @@
+"""Strategy pattern for prompt/response format abstraction."""
+from pipeline.strategies.base import PromptStrategy
+from pipeline.strategies.factory import StrategyFactory
+
+# Import strategy implementations to trigger decorator registration
+from pipeline.strategies.chatml import ChatMLStrategy  # noqa: F401
+from pipeline.strategies.alpaca import AlpacaStrategy  # noqa: F401
+
+__all__ = [
+    "PromptStrategy",
+    "StrategyFactory",
+    "ChatMLStrategy",
+    "AlpacaStrategy",
+]

pipeline/strategies/alpaca.py

@@ -0,0 +1,43 @@
+"""Alpaca format strategy."""
+from typing import List
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.strategies.base import PromptStrategy
+from pipeline.strategies.factory import StrategyFactory
+
+
+@StrategyFactory.register("alpaca")
+class AlpacaStrategy(PromptStrategy):
+    """Alpaca format: ``### Instruction: ... \\n\\n### Response: ... <eos>``"""
+
+    def __init__(
+        self,
+        tokenizer: BpeTokenizer,
+        instruction_start: str = "### Instruction:\n",
+        response_start: str = "### Response:\n",
+        response_suffix: str = "\n<eos>",
+    ):
+        super().__init__(tokenizer)
+        self.instruction_start = instruction_start
+        self.response_start = response_start
+        self.response_suffix = response_suffix
+
+        self._instruction_start_ids = self._encode_format(instruction_start)
+        self._separator_ids = self._encode_format("\n\n")
+        self._response_start_ids = self._encode_format(response_start)
+        self._response_suffix_ids = self._encode_format(response_suffix)
+
+    @property
+    def name(self) -> str:
+        return "alpaca"
+
+    def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
+        return (
+            self._instruction_start_ids
+            + query_tokens
+            + self._separator_ids
+            + self._response_start_ids
+        )
+
+    def assemble_response(self, response_tokens: List[int]) -> List[int]:
+        return response_tokens + self._response_suffix_ids

pipeline/strategies/base.py

@@ -0,0 +1,41 @@
+"""Abstract base class for prompt construction strategies."""
+from abc import ABC, abstractmethod
+from typing import List
+
+from pipeline.tokenizer import BpeTokenizer
+
+
+class PromptStrategy(ABC):
+    """Abstract base for prompt/response format strategies.
+
+    Strategies operate at the token level: the Processor tokenizes raw
+    text (query, response, …) and passes token lists to the Strategy,
+    which assembles them with pre-encoded format tokens.
+    """
+
+    def __init__(self, tokenizer: BpeTokenizer):
+        self.tokenizer = tokenizer
+
+    def _encode_format(self, text: str) -> List[int]:
+        """Encode a format string that may contain special tokens."""
+        return self.tokenizer.encode(text, add_special_tokens=False)
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+    @abstractmethod
+    def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
+        """Assemble query tokens into a complete prompt with format tokens.
+
+        The prompt includes all tokens up to (and including) the response
+        start marker, e.g. ``<|im_start|>assistant\n``.
+        """
+
+    @abstractmethod
+    def assemble_response(self, response_tokens: List[int]) -> List[int]:
+        """Wrap response tokens with format tokens (suffix, eos, etc)."""
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(name='{self.name}')"

pipeline/strategies/chatml.py

@@ -0,0 +1,41 @@
+"""ChatML format strategy."""
+from typing import List
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.strategies.base import PromptStrategy
+from pipeline.strategies.factory import StrategyFactory
+
+
+@StrategyFactory.register("chatml")
+class ChatMLStrategy(PromptStrategy):
+    """ChatML format: ``<|im_start|>user ... <|im_end|> <|im_start|>assistant ... <|im_end|> <eos>``"""
+
+    def __init__(
+        self,
+        tokenizer: BpeTokenizer,
+        user_start: str = "<|im_start|>user\n",
+        user_end: str = "<|im_end|>\n",
+        assistant_start: str = "<|im_start|>assistant\n",
+        assistant_end: str = "<|im_end|>\n<eos>",
+    ):
+        super().__init__(tokenizer)
+
+        self._user_start_ids = self._encode_format(user_start)
+        self._user_end_ids = self._encode_format(user_end)
+        self._assistant_start_ids = self._encode_format(assistant_start)
+        self._assistant_end_ids = self._encode_format(assistant_end)
+
+    @property
+    def name(self) -> str:
+        return "chatml"
+
+    def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
+        return (
+            self._user_start_ids
+            + query_tokens
+            + self._user_end_ids
+            + self._assistant_start_ids
+        )
+
+    def assemble_response(self, response_tokens: List[int]) -> List[int]:
+        return response_tokens + self._assistant_end_ids

pipeline/strategies/factory.py

@@ -0,0 +1,70 @@
+"""Factory for creating and registering prompt strategies."""
+from typing import Dict, List, Type
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.strategies.base import PromptStrategy
+
+
+class StrategyFactory:
+    """Registry and factory for PromptStrategy implementations.
+
+    Supports decorator-based registration for extensible strategy types.
+
+    Example usage::
+
+        @StrategyFactory.register("custom")
+        class CustomStrategy(PromptStrategy):
+            ...
+
+        strategy = StrategyFactory.create("custom", tokenizer, **kwargs)
+    """
+
+    STRATEGY_MAP: Dict[str, Type[PromptStrategy]] = {}
+
+    @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[PromptStrategy]) -> Type[PromptStrategy]:
+            if not issubclass(strategy_cls, PromptStrategy):
+                raise TypeError(
+                    f"{strategy_cls.__name__} must inherit from PromptStrategy"
+                )
+            cls.STRATEGY_MAP[name] = strategy_cls
+            return strategy_cls
+
+        return decorator
+
+    @classmethod
+    def create(cls, name: str, tokenizer: BpeTokenizer, **kwargs) -> PromptStrategy:
+        """Create a strategy by name.
+
+        Args:
+            name: Registered strategy name (e.g. ``"chatml"``, ``"alpaca"``).
+            tokenizer: Tokenizer instance (required by all strategies).
+            **kwargs: Forwarded to the strategy constructor.
+
+        Returns:
+            Strategy instance.
+
+        Raises:
+            ValueError: If name is not registered.
+        """
+        if name not in cls.STRATEGY_MAP:
+            raise ValueError(
+                f"Unknown strategy: '{name}'. "
+                f"Supported types: {sorted(cls.STRATEGY_MAP.keys())}"
+            )
+        return cls.STRATEGY_MAP[name](tokenizer=tokenizer, **kwargs)
+
+    @classmethod
+    def available_types(cls) -> List[str]:
+        """Return list of registered strategy type names."""
+        return list(cls.STRATEGY_MAP.keys())