42 lines
1.4 KiB
Python
42 lines
1.4 KiB
Python
"""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
|