ViperEkura
598e1ce4ae
- 将 pipeline/packing.py 拆分为 packing/ 子包 (base/stream/binpack) - 新增 BfdPacker(默认)/FfDPacker/GreedyPacker，移除 StreamingPacker - 超长序列直接截断至 pack_size - group_size 语义改为"每 N 个 chunk 合并为一块"，默认 1000 - 新增 AutoTokenizer.token_to_id()，修复 ChatML 中 hacky 的 nl_id 获取 - pad_value 默认改为 2(pad_token_id)，position_ids pad=0, loss_mask pad=False - 新增 position_ids 打包后归零一致性测试 - scripts/cache_h5.py 新增 --pack-algo 参数
44 lines
1.4 KiB
Python
44 lines
1.4 KiB
Python
"""ChatML format strategy."""
|
||
|
||
from typing import List
|
||
|
||
from pipeline.tokenize import AutoTokenizer
|
||
from pipeline.strategies.base import PromptStrategy
|
||
from pipeline.strategies.factory import StrategyFactory
|
||
|
||
|
||
@StrategyFactory.register("chatml")
|
||
class ChatMLStrategy(PromptStrategy):
|
||
"""ChatML format strategy."""
|
||
|
||
def __init__(
|
||
self,
|
||
tokenizer: AutoTokenizer,
|
||
user_start: str = "<｜im▁start｜>user",
|
||
user_end: str = "<｜im▁end｜>",
|
||
assistant_start: str = "<｜im▁start｜>assistant",
|
||
assistant_end: str = "<｜im▁end｜>",
|
||
):
|
||
super().__init__(tokenizer)
|
||
nl_id = tokenizer.token_to_id("\n")
|
||
|
||
self._user_start_ids = self._encode_format(user_start) + [nl_id]
|
||
self._user_end_ids = self._encode_format(user_end) + [nl_id]
|
||
self._assistant_start_ids = self._encode_format(assistant_start) + [nl_id]
|
||
self._assistant_end_ids = self._encode_format(assistant_end) + [nl_id]
|
||
|
||
@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
|