reafactor: 重构项目

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