diff --git a/angelslim/data/text_dataset.py b/angelslim/data/text_dataset.py index 8c0b2ead..54838a3a 100644 --- a/angelslim/data/text_dataset.py +++ b/angelslim/data/text_dataset.py @@ -174,15 +174,40 @@ def _load_jsonl_data(self, data_path: str, num_samples: int): prompt_messages = messages[:last_assistant_idx] assistant_msg = messages[last_assistant_idx] + # Collect optional chat-template kwargs. Two shapes are + # supported so we can consume both legacy and new records: + # (a) top-level fields on ``data`` (legacy): + # ``tools`` / ``reasoning_effort`` / ``is_ai_search`` + # (b) grouped under ``apply_chat_template_kwargs`` (new + # agent SFT / rollout format), which may also carry + # fields like ``interleaved_thinking`` etc. + tpl_kwargs = dict(data.get("apply_chat_template_kwargs") or {}) + for legacy_key in ("tools", "reasoning_effort", "is_ai_search"): + if legacy_key in data and legacy_key not in tpl_kwargs: + tpl_kwargs[legacy_key] = data[legacy_key] + # ``tools`` in the template is invoked as a positional arg + # in the legacy path; keep it separate so we can pass it in + # the same way when present. + tools = tpl_kwargs.pop("tools", None) + # Tokenize the prompt (up to the generation marker) and the # full conversation separately so we know exactly where the - # assistant reply starts. - prompt_text = self.processor.apply_chat_template( - prompt_messages, tokenize=False, add_generation_prompt=True + # assistant reply starts. The two calls MUST use the same + # kwargs, otherwise the tokenized prefix of ``full_text`` + # will not align with ``prompt_text`` and the label mask + # below would be off by a few tokens. + prompt_text = self._apply_chat_template_safe( + prompt_messages, + tools=tools, + add_generation_prompt=True, + **tpl_kwargs, ) full_messages = prompt_messages + [assistant_msg] - full_text = self.processor.apply_chat_template( - full_messages, tokenize=False, add_generation_prompt=False + full_text = self._apply_chat_template_safe( + full_messages, + tools=tools, + add_generation_prompt=False, + **tpl_kwargs, ) # Legacy branch: thinking-style data without a chat template. @@ -249,6 +274,62 @@ def _load_jsonl_data(self, data_path: str, num_samples: int): line_count += 1 + def _apply_chat_template_safe( + self, + messages: List[Dict], + tools=None, + add_generation_prompt: bool = False, + **extra_kwargs, + ) -> str: + """Apply chat template with graceful fallback. + + Newer Hy3-style tokenizers accept extra kwargs such as + ``reasoning_effort`` / ``is_ai_search`` / ``interleaved_thinking`` + and a positional ``tools`` argument, while stock HF tokenizers + may not. Try the richest signature first and progressively drop + unsupported arguments so this loader stays compatible with both. + """ + + # Filter out None values so we don't push unnecessary kwargs + # through — some templates treat ``None`` differently from + # missing. + kwargs = {k: v for k, v in extra_kwargs.items() if v is not None} + base = dict( + tokenize=False, + add_generation_prompt=add_generation_prompt, + ) + # Preserve legacy behavior: signal that we are producing training + # text so templates that vary between train/eval render the + # correct branch. + base["is_training"] = True + + # Attempt 1: full signature (tools + all extras). + try: + if tools is not None: + return self.processor.apply_chat_template(messages, tools, **base, **kwargs) + return self.processor.apply_chat_template(messages, **base, **kwargs) + except TypeError: + pass + + # Attempt 2: drop is_training, keep the rest. + base.pop("is_training", None) + try: + if tools is not None: + return self.processor.apply_chat_template(messages, tools, **base, **kwargs) + return self.processor.apply_chat_template(messages, **base, **kwargs) + except TypeError: + pass + + # Attempt 3: drop all optional extras (reasoning_effort etc.). + try: + if tools is not None: + return self.processor.apply_chat_template(messages, tools, **base) + return self.processor.apply_chat_template(messages, **base) + except TypeError: + # Attempt 4: also drop ``tools`` — final fallback for stock + # HF tokenizers that don't understand tool-call templates. + return self.processor.apply_chat_template(messages, **base) + def _prepare_messages(self, data: Dict) -> List[Dict]: """Prepare chat messages from data entry""" if "messages" in data: