Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / text / chat_template.py: 91%
295 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Chat-template implementations used by LLM conversation transforms.
17Canonical merge (05 §11.3) of the former
18``components/data/chat_template.py`` and
19``components/datasets/llm/chat_template.py`` duplicates: one
20``CHAT_TEMPLATE_REGISTRY``, the ``DatasetLogger`` logging path, and
21``IGNORE_INDEX`` shared from the constants module.
22"""
24import os
25from abc import ABC, abstractmethod
26from typing import (
27 TYPE_CHECKING,
28 Any,
29 Callable,
30 Dict,
31 Iterator,
32 List,
33 Mapping,
34 MutableMapping,
35 Optional,
36 Sequence,
37 Type,
38 Union,
39)
41from hyper_parallel.data.constants import IGNORE_INDEX
42from hyper_parallel.data.dataset_logging import get_dataset_logger
44if TYPE_CHECKING:
45 from transformers import PreTrainedTokenizer
47logger = get_dataset_logger(__name__)
50class Registry(MutableMapping):
51 """Domain registry for chat-template formatters.
53 Inlined from the former ``auto_models/components/utils/registry.py`` in
54 stage 7 (05 §10.3 row 565): the generic Registry was merged into the
55 domain registries that consume it instead of living in a utils package.
56 """
58 # Class instance object, so that a call to `register` can be reflected into all other files correctly, even if
59 # a new instance is created (in order to locally override a given function)
60 registry = []
62 def __init__(self, name: str) -> None:
63 """Initialize the registry.
65 Args:
66 name: Human-readable name of the registered category, used in
67 error messages.
68 """
69 self._name = name
70 self.registry.append(name)
71 self._local_mapping = {}
72 self._global_mapping = {}
74 def __getitem__(self, key: str) -> Union[Type, Callable]:
75 """Look up a registered class or function by key.
77 Args:
78 key: Registration key.
80 Returns:
81 The locally overridden value when present, otherwise the globally
82 registered value.
84 Raises:
85 ValueError: If ``key`` is not registered.
86 """
87 # First check if instance has a local override
88 if key not in self.valid_keys():
89 raise ValueError(f"Unknown {self._name} name: {key}. No {self._name} registered for this source.")
90 if key in self._local_mapping:
91 return self._local_mapping[key]
92 return self._global_mapping[key]
94 def __setitem__(self, key: str, value: Union[Type, Callable]) -> None:
95 """Set a local override for ``key`` without affecting other instances."""
96 # Allow local update of the default functions without impacting other instances
97 self._local_mapping.update({key: value})
99 def __delitem__(self, key: str) -> None:
100 """Delete the local override for ``key``."""
101 del self._local_mapping[key]
103 def __iter__(self) -> Iterator[str]:
104 """Iterate over all valid keys, local overrides taking precedence."""
105 # Ensure we use all keys, with the overwritten ones on top
106 return iter({**self._global_mapping, **self._local_mapping})
108 def __len__(self) -> int:
109 """Return the number of distinct registered keys."""
110 return len(self._global_mapping.keys() | self._local_mapping.keys())
112 def register(
113 self,
114 key: str,
115 cls_or_func: Optional[Union[Type, Callable]] = None,
116 ) -> Union[Type, Callable]:
117 """Register a class or function under ``key``.
119 Can be used either as a plain call or as a decorator::
121 registry.register("name", MyClass)
123 @registry.register("name")
124 class MyClass: ...
126 Args:
127 key: Registration key.
128 cls_or_func: The class or function to register. When ``None``, a
129 decorator is returned instead.
131 Returns:
132 The registered class/function, or a decorator when ``cls_or_func``
133 is ``None``.
135 Raises:
136 ValueError: If ``key`` is already registered (decorator form).
137 """
138 if cls_or_func is not None:
139 self._global_mapping[key] = cls_or_func
140 return cls_or_func
142 def decorator(cls_or_func: Union[Type, Callable]) -> Union[Type, Callable]:
143 """Register the decorated class or function under ``key``."""
144 if key in self._global_mapping:
145 raise ValueError(
146 f"{self._name} for '{key}' is already registered. Cannot register duplicate {self._name}."
147 )
148 self._global_mapping.update({key: cls_or_func})
149 return cls_or_func
151 return decorator
153 def valid_keys(self) -> List[str]:
154 """Return the list of all registered keys."""
155 return list(self.keys())
158CHAT_TEMPLATE_REGISTRY = Registry("ChatTemplate")
161def _format_janus_message(
162 message: Dict[str, str], index: int, message_count: int, task_type: str, assistant_count: int
163) -> tuple[str, int]:
164 """Format one Janus message and update the generation assistant count."""
165 role = message["role"]
166 content = message["content"]
167 separators = ["\n\n", "<｜end▁of▁sentence｜>"]
168 if content == "":
169 return role + ":", assistant_count
170 if "assistant" in role and (
171 "wikihow_generation" in task_type or "interleave_generation" in task_type
172 ):
173 prefix = "Assistant: " if assistant_count == 0 else ""
174 suffix = separators[1] if index + 1 == message_count else separators[0]
175 return prefix + content.strip() + suffix, assistant_count + 1
176 if "assistant" in role:
177 return "Assistant: " + content.strip() + separators[1], assistant_count
178 if "user" in role:
179 return "User: " + content.strip() + separators[0], assistant_count
180 if "system" in role and "wikihow_generation" in task_type:
181 instruction = "Please generate a step-by-step tutorial with images for the following question."
182 return content.strip() + separators[0] + instruction + separators[0], assistant_count
183 if "system" in role:
184 return content.strip() + separators[0], assistant_count
185 raise ValueError(f"Unknown role {role}, should be one of {{system, user, assistant}}.")
188def _janus_labels(content_ids: List[int], image_token_id: int, loss_mask: int, task_type: str) -> List[int]:
189 """Build labels for one encoded Janus message."""
190 if loss_mask != 1:
191 return [IGNORE_INDEX] * len(content_ids)
192 if (
193 image_token_id in content_ids
194 and "wikihow_generation" not in task_type
195 and "interleave_generation" not in task_type
196 ):
197 return [image_token_id if token == image_token_id else IGNORE_INDEX for token in content_ids]
198 return content_ids
201def build_chat_template(
202 template_name: str,
203 tokenizer: "PreTrainedTokenizer",
204 *,
205 chat_template_kwargs: Optional[Dict[str, Any]] = None,
206 log_first_rendered_template: bool = False,
207) -> "ChatTemplate":
208 """Build the registered chat template for a tokenizer.
210 Args:
211 template_name: Registry key of the chat template implementation.
212 tokenizer: Hugging Face tokenizer used for encoding and special tokens.
213 chat_template_kwargs: Optional keyword arguments forwarded to the tokenizer's native chat template.
214 log_first_rendered_template: Whether the native tokenizer template logs its first rendered conversation.
216 Returns:
217 The configured chat template instance.
219 Raises:
220 ValueError: If native-template options are requested for a non-tokenizer template.
221 """
222 if template_name != "tokenizer":
223 if chat_template_kwargs or log_first_rendered_template:
224 raise ValueError("Native chat-template options require chat_template='tokenizer'")
225 return CHAT_TEMPLATE_REGISTRY[template_name](tokenizer)
226 return CHAT_TEMPLATE_REGISTRY[template_name](
227 tokenizer,
228 chat_template_kwargs=chat_template_kwargs,
229 log_first_rendered_template=log_first_rendered_template,
230 )
233class ChatTemplate(ABC):
234 """
235 Abstract class for chat template.
236 """
238 def __init__(self, tokenizer: Any) -> None:
239 """Store the tokenizer used for encoding and special tokens."""
240 self.tokenizer = tokenizer
242 def save_pretrained(self, output_dir: str) -> None:
243 """Attach the jinja template to the tokenizer and save it.
245 Args:
246 output_dir: Directory the tokenizer assets are written to.
247 """
248 self.tokenizer.chat_template = self.get_jinja_template()
249 try:
250 self.tokenizer.save_pretrained(output_dir)
251 except OSError:
252 logger.warning("Failed to save tokenizer.")
254 @abstractmethod
255 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:
256 """
257 Encodes messages to a dictionary of input_ids, attention_mask, and labels.
258 """
260 @abstractmethod
261 def get_jinja_template(self) -> str:
262 """
263 Gets the jinja template for the chat template.
264 """
267@CHAT_TEMPLATE_REGISTRY.register("default")
268class DefaultTemplate(ChatTemplate):
269 """Render each message as ``Role: content`` followed by the EOS token."""
271 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:
272 """Encode messages into input_ids, attention_mask, and loss-masked labels.
274 Args:
275 messages: Conversation records with ``role``, ``content``, and ``loss_mask`` fields.
276 max_seq_len: Maximum sequence length kept from the end of the encoding.
278 Returns:
279 Model inputs with input_ids, attention_mask, and labels.
280 """
281 input_ids, attention_mask, labels = [], [], []
282 for message in messages:
283 content_str = message["role"].title() + ": " + message["content"].strip() + self.tokenizer.eos_token + "\n"
284 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)
285 input_ids += content_ids
286 attention_mask += [1] * len(content_ids)
287 if message["loss_mask"] == 1:
288 labels += content_ids
289 else:
290 labels += [IGNORE_INDEX] * len(content_ids)
292 model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
293 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}
294 return model_inputs
296 def get_jinja_template(self) -> str:
297 """Return the jinja template matching the default ``Role: content`` rendering."""
298 return (
299 "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"
300 "{% for message in messages %}"
301 "{{ message['role'].title() + ': ' + message['content'] | trim + eos_token + '\n' }}"
302 "{% endfor %}"
303 "{% if add_generation_prompt %}{{ 'Assistant: ' }}{% endif %}"
304 )
307@CHAT_TEMPLATE_REGISTRY.register("tokenizer")
308class TokenizerTemplate(ChatTemplate):
309 """Use the tokenizer's native chat template and train only selected messages."""
311 _RESERVED_TEMPLATE_KWARGS = frozenset({"messages", "tools", "tokenize", "add_generation_prompt", "return_dict"})
312 _REASONING_EFFORTS = frozenset({"low", "medium", "xhigh"})
314 def __init__(
315 self,
316 tokenizer: Any,
317 chat_template_kwargs: Optional[Dict[str, Any]] = None,
318 log_first_rendered_template: bool = False,
319 ) -> None:
320 """Initialize the native tokenizer template wrapper.
322 Args:
323 tokenizer: Tokenizer providing the native Jinja chat template.
324 chat_template_kwargs: Keyword arguments forwarded to every native chat-template rendering.
325 log_first_rendered_template: Whether to log the first input messages and their fully rendered template.
327 Raises:
328 ValueError: If a reserved argument is configured or a known reasoning option is invalid.
329 """
330 super().__init__(tokenizer)
331 self.chat_template_kwargs = dict(chat_template_kwargs or {})
332 self._validate_chat_template_kwargs()
333 self.log_first_rendered_template = log_first_rendered_template
334 self._has_logged_first_rendered_template = False
336 def _validate_chat_template_kwargs(self) -> None:
337 """Validate Trainer-owned arguments and known Qwen reasoning options."""
338 reserved_kwargs = self._RESERVED_TEMPLATE_KWARGS.intersection(self.chat_template_kwargs)
339 if reserved_kwargs:
340 raise ValueError(
341 "chat_template_kwargs cannot override Trainer-owned arguments: "
342 f"{sorted(reserved_kwargs)!r}"
343 )
345 enable_thinking = self.chat_template_kwargs.get("enable_thinking")
346 if enable_thinking is not None and not isinstance(enable_thinking, bool):
347 raise ValueError("chat_template_kwargs.enable_thinking must be a bool")
349 preserve_thinking = self.chat_template_kwargs.get("preserve_thinking")
350 if preserve_thinking is not None and not isinstance(preserve_thinking, bool):
351 raise ValueError("chat_template_kwargs.preserve_thinking must be a bool")
353 reasoning_effort = self.chat_template_kwargs.get("reasoning_effort")
354 if reasoning_effort is not None and reasoning_effort not in self._REASONING_EFFORTS:
355 raise ValueError(
356 "chat_template_kwargs.reasoning_effort must be one of "
357 f"{sorted(self._REASONING_EFFORTS)!r}"
358 )
360 def _apply_chat_template(
361 self,
362 messages: Sequence[Mapping[str, Any]],
363 *,
364 tokenize: bool,
365 return_dict: bool,
366 add_generation_prompt: bool = False,
367 tools: Optional[Sequence[Mapping[str, Any]]] = None,
368 ) -> Any:
369 """Render messages with one consistent set of native-template arguments."""
370 template_kwargs = dict(self.chat_template_kwargs)
371 if tools:
372 template_kwargs["tools"] = tools
373 return self.tokenizer.apply_chat_template(
374 messages,
375 tokenize=tokenize,
376 add_generation_prompt=add_generation_prompt,
377 return_dict=return_dict,
378 **template_kwargs,
379 )
381 def _log_first_template(
382 self,
383 messages: Sequence[Mapping[str, Any]],
384 tools: Optional[Sequence[Mapping[str, Any]]] = None,
385 ) -> None:
386 """Log the first raw conversation and native-template rendering once."""
387 if not self.log_first_rendered_template or self._has_logged_first_rendered_template:
388 return
389 if os.environ.get("RANK", "0") != "0":
390 self._has_logged_first_rendered_template = True
391 return
393 rendered_template = self._apply_chat_template(
394 messages,
395 tokenize=False,
396 return_dict=False,
397 tools=tools,
398 )
399 logger.info("First chat-template input messages: %r", messages)
400 template_kwargs = dict(self.chat_template_kwargs)
401 if tools:
402 template_kwargs["tools"] = tools
403 logger.info("First chat-template kwargs: %r", template_kwargs)
404 logger.info("First rendered chat template:\n%s", rendered_template)
405 self._has_logged_first_rendered_template = True
407 def _update_prefix_labels(
408 self,
409 previous_ids: List[int],
410 current_ids: List[int],
411 labels: List[int],
412 ) -> None:
413 """Ensure appending a message does not rewrite the existing token prefix."""
414 # ``labels`` is only adjusted by overrides that permit a terminal rewrite.
415 del labels
416 previous_length = len(previous_ids)
417 if current_ids[:previous_length] != previous_ids:
418 raise ValueError(
419 "The tokenizer chat template structurally rewrote an earlier conversation prefix; "
420 "the generic tokenizer template requires prefix-stable rendering."
421 )
423 def _encode_text_continuation(
424 self,
425 text: str,
426 previous_text: str,
427 previous_ids: List[int],
428 ) -> List[int]:
429 """Keep earlier source/target token boundaries when the rendered text extends history."""
430 if text.startswith(previous_text):
431 # Re-encoding history can merge tokens across a masked source/target boundary
432 # (e.g. the two newlines in empty reasoning), invalidating earlier labels.
433 return previous_ids + self.tokenizer.encode(text[len(previous_text):], add_special_tokens=False)
434 # Actual template rewrites still go through the strict prefix validator, including
435 # the explicitly supported GPT-OSS terminal-token substitution.
436 return self.tokenizer.encode(text, add_special_tokens=False)
438 def _encode_assistant_source_target(
439 self,
440 messages: Sequence[Mapping[str, Any]],
441 previous_text: str,
442 previous_ids: List[int],
443 tools: Optional[Sequence[Mapping[str, Any]]] = None,
444 ) -> tuple[List[int], List[int], str]:
445 """Encode history plus the prefilled header separately from the current assistant target.
447 ``messages`` ends at the assistant being supervised. The source includes
448 earlier assistant targets, whose labels must be preserved by the caller;
449 only newly appended source tokens should receive ``IGNORE_INDEX``.
450 """
451 completed_text = self._apply_chat_template(
452 messages, tokenize=False, return_dict=False, tools=tools,
453 )
454 source_text = self._apply_chat_template(
455 messages[:-1], tokenize=False, return_dict=False, add_generation_prompt=True, tools=tools,
456 )
457 if not completed_text.startswith(source_text):
458 raise ValueError(
459 "The tokenizer generation prompt is not a text prefix of the completed assistant message; "
460 "training and inference assistant headers would be inconsistent."
461 )
462 # Separate encoding prevents BPE from merging a prefilled header token with the first target token.
463 source_ids = self._encode_text_continuation(source_text, previous_text, previous_ids)
464 target_ids = self.tokenizer.encode(completed_text[len(source_text):], add_special_tokens=False)
465 return source_ids, target_ids, completed_text
467 def encode_messages(
468 self,
469 messages: Sequence[Mapping[str, Any]],
470 max_seq_len: int = 8192,
471 tools: Optional[Sequence[Mapping[str, Any]]] = None,
472 ) -> Dict[str, List[int]]:
473 """Encode messages with the tokenizer chat template and mask non-assistant loss.
475 Single-turn and multi-turn conversations share the same traversal. Each
476 supervised assistant uses its generation prompt as the source/target
477 boundary, preserving tokens and labels from previous turns. Only newly
478 appended text is encoded, so BPE cannot merge across earlier loss boundaries.
479 Rendering must remain prefix-stable. Leading system messages are
480 first rendered together with the following non-system message because
481 some native templates reject a system-only conversation prefix.
483 Args:
484 messages: Conversation records with ``role``, ``content``, and optional ``loss_mask``.
485 max_seq_len: Maximum sequence length kept from the end of the encoding.
486 tools: Optional tool definitions forwarded to the native tokenizer template.
488 Returns:
489 Model inputs with input_ids, attention_mask, and labels.
491 Raises:
492 ValueError: If the template rewrites or shortens an earlier conversation prefix, or if leading
493 system messages cannot be grouped without changing the configured loss mask.
494 """
495 self._log_first_template(messages, tools)
496 input_ids: List[int] = []
497 labels: List[int] = []
498 previous_length = 0
499 previous_text = ""
501 has_pending_system_messages = False
502 for end, message in enumerate(messages, start=1):
503 loss_mask = message.get("loss_mask", 1 if message["role"] == "assistant" else 0)
504 if previous_length == 0 and message["role"] == "system":
505 if loss_mask == 1:
506 raise ValueError(
507 "A leading system message with loss enabled cannot be grouped with the first user message."
508 )
509 has_pending_system_messages = True
510 continue
511 if has_pending_system_messages and previous_length == 0 and loss_mask == 1:
512 raise ValueError(
513 "The first message after leading system messages must have loss disabled so their token "
514 "boundaries can be grouped safely."
515 )
517 supervise_assistant = message["role"] == "assistant" and loss_mask == 1
518 if supervise_assistant:
519 source_ids, target_ids, rendered_text = self._encode_assistant_source_target(
520 messages[:end], previous_text, input_ids, tools,
521 )
522 rendered_conversation_ids = source_ids + target_ids
523 else:
524 rendered_text = self._apply_chat_template(
525 messages[:end],
526 tokenize=False,
527 return_dict=False,
528 tools=tools,
529 )
530 rendered_conversation_ids = self._encode_text_continuation(rendered_text, previous_text, input_ids)
531 current_length = len(rendered_conversation_ids)
532 if current_length < previous_length:
533 raise ValueError(
534 "The tokenizer chat template shortened the conversation after adding a message; "
535 "assistant-only loss masking requires monotonic message boundaries."
536 )
538 self._update_prefix_labels(input_ids, rendered_conversation_ids, labels)
539 if supervise_assistant:
540 if (
541 len(source_ids) < previous_length
542 or source_ids[:previous_length] != input_ids
543 ):
544 raise ValueError(
545 "The tokenizer generation prompt structurally rewrote the conversation prefix; "
546 "assistant-header loss masking requires the inference prompt to extend the training history."
547 )
548 # History contains supervised targets from earlier turns: mask only the new source suffix.
549 labels.extend([IGNORE_INDEX] * (len(source_ids) - previous_length))
550 labels.extend(target_ids)
551 else:
552 new_ids = rendered_conversation_ids[previous_length:]
553 labels.extend(new_ids if loss_mask == 1 else [IGNORE_INDEX] * len(new_ids))
554 input_ids = rendered_conversation_ids
555 previous_length = current_length
556 previous_text = rendered_text
558 if has_pending_system_messages and previous_length == 0:
559 raise ValueError("A conversation containing only system messages cannot be rendered for training.")
561 input_ids = input_ids[-max_seq_len:]
562 labels = labels[-max_seq_len:]
563 return {
564 "input_ids": input_ids,
565 "attention_mask": [1] * len(input_ids),
566 "labels": labels,
567 }
569 def get_jinja_template(self) -> str:
570 """Return the tokenizer's native jinja chat template.
572 Raises:
573 ValueError: If the tokenizer does not define a chat template.
574 """
575 if not self.tokenizer.chat_template:
576 raise ValueError("The tokenizer does not define a native chat template.")
577 return self.tokenizer.chat_template
580@CHAT_TEMPLATE_REGISTRY.register("gpt_oss")
581class GptOssTokenizerTemplate(TokenizerTemplate):
582 """Support GPT-OSS terminal rewrites from <|return|> to <|end|>."""
584 def __init__(self, tokenizer: "PreTrainedTokenizer") -> None:
585 """Resolve the GPT-OSS terminal token IDs required for prefix rewrites.
587 Args:
588 tokenizer: Hugging Face tokenizer that must define <|return|> and <|end|>.
590 Raises:
591 ValueError: If either terminal token is missing from the vocabulary.
592 """
593 super().__init__(tokenizer)
594 self.return_token_id = tokenizer.convert_tokens_to_ids("<|return|>")
595 self.end_token_id = tokenizer.convert_tokens_to_ids("<|end|>")
596 if tokenizer.unk_token_id in (self.return_token_id, self.end_token_id):
597 raise ValueError("The GPT-OSS chat template requires <|return|> and <|end|> tokenizer tokens.")
599 def _update_prefix_labels(
600 self,
601 previous_ids: List[int],
602 current_ids: List[int],
603 labels: List[int],
604 ) -> None:
605 """Accept only the terminal <|return|>-to-<|end|> rewrite and fix its label.
607 Args:
608 previous_ids: Token IDs rendered before the latest message was appended.
609 current_ids: Token IDs rendered after the latest message was appended.
610 labels: Loss labels aligned with ``previous_ids``, updated in place.
612 Raises:
613 ValueError: If the rewrite is not the supported terminal substitution.
614 """
615 previous_length = len(previous_ids)
616 rewritten_positions = [
617 index
618 for index in range(previous_length)
619 if previous_ids[index] != current_ids[index]
620 ]
621 if not rewritten_positions:
622 return
624 is_terminal_rewrite = (
625 rewritten_positions == [previous_length - 1]
626 and len(current_ids) > previous_length
627 and previous_ids[-1] == self.return_token_id
628 and current_ids[previous_length - 1] == self.end_token_id
629 and self.return_token_id not in current_ids[previous_length:]
630 )
631 if not is_terminal_rewrite:
632 raise ValueError(
633 "The GPT-OSS tokenizer chat template structurally rewrote an earlier conversation prefix; "
634 "only the terminal <|return|>-to-<|end|> substitution is supported."
635 )
637 if labels[-1] != IGNORE_INDEX:
638 labels[-1] = self.end_token_id
641@CHAT_TEMPLATE_REGISTRY.register("llama2")
642class Llama2Template(ChatTemplate):
643 """Render messages with the Llama-2 [INST]/<<SYS>> prompt format."""
645 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:
646 """Encode messages into input_ids, attention_mask, and loss-masked labels.
648 Args:
649 messages: Conversation records with ``role``, ``content``, and ``loss_mask`` fields.
650 max_seq_len: Maximum sequence length kept from the end of the encoding.
652 Returns:
653 Model inputs with input_ids, attention_mask, and labels.
655 Raises:
656 ValueError: If a message role is not one of system, user, assistant, or tool.
657 """
658 input_ids, attention_mask, labels = [], [], []
659 for message in messages:
660 if message["role"] == "system":
661 content_str = "<<SYS>>\n" + message["content"].strip() + "\n<</SYS>>\n\n"
662 elif message["role"] == "user":
663 content_str = self.tokenizer.bos_token + "[INST] " + message["content"].strip() + " [/INST]"
664 elif message["role"] == "assistant":
665 content_str = " " + message["content"].strip() + " " + self.tokenizer.eos_token
666 elif message["role"] == "tool":
667 content_str = self.tokenizer.bos_token + "[TOOL] " + message["content"].strip() + " [/TOOL]"
668 else:
669 raise ValueError(
670 f"Unknown role {message['role']}, should be one of {{system, user, assistant, tool}}."
671 )
673 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)
674 input_ids += content_ids
675 attention_mask += [1] * len(content_ids)
676 if message["loss_mask"] == 1:
677 labels += content_ids
678 else:
679 labels += [IGNORE_INDEX] * len(content_ids)
681 model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
682 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}
683 return model_inputs
685 def get_jinja_template(self) -> str:
686 """Return the jinja template matching the Llama-2 prompt format."""
687 return (
688 "{% if messages[0]['role'] == 'system' %}"
689 "{{ '<<SYS>>\n' + messages[0]['content'] | trim + '\n<</SYS>>\n\n' }}"
690 "{% set loop_messages = messages[1:] %}"
691 "{% else %}"
692 "{% set loop_messages = messages %}"
693 "{% endif %}"
694 "{% for message in loop_messages %}"
695 "{% set content = message['content'] %}"
696 "{% if message['role'] == 'user' %}"
697 "{{ bos_token + '[INST] ' + content | trim + ' [/INST]' }}"
698 "{% elif message['role'] == 'tool' %}"
699 "{{ bos_token + '[TOOL] ' + content | trim + ' [/TOOL]' }}"
700 "{% elif message['role'] == 'assistant' %}"
701 "{{ ' ' + content | trim + ' ' + eos_token }}"
702 "{% endif %}"
703 "{% endfor %}"
704 )
707@CHAT_TEMPLATE_REGISTRY.register("Janus")
708class JanusTemplate(ChatTemplate):
709 """Render Janus multimodal conversations with image placeholder masks."""
711 def encode_messages(
712 self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192, task_type: str = ""
713 ) -> Dict[str, List[int]]:
714 """Encode messages into token, label, and image mask fields.
716 Args:
717 messages: Conversation records with ``role``, ``content``, and ``loss_mask`` fields.
718 max_seq_len: Maximum sequence length kept from the end of the encoding.
719 task_type: Optional task selector enabling generation-specific separators.
721 Returns:
722 Model inputs with input_ids, attention_mask, labels, images_seq_mask,
723 and images_emb_mask.
725 Raises:
726 ValueError: If a message role or image placeholder count is invalid.
727 """
728 # pylint: disable=too-many-locals
729 input_ids, attention_mask, labels = [], [], []
730 images_seq_mask, images_emb_mask = [], []
731 assistant_count = 0
732 for idx, message in enumerate(messages):
733 content_str, assistant_count = _format_janus_message(
734 message, idx, len(messages), task_type, assistant_count
735 )
736 if "system" in message["role"]:
737 content_ids = self.tokenizer.encode(content_str)
738 else:
739 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)
740 input_ids += content_ids
741 attention_mask += [1] * len(content_ids)
742 image_token_id = self.tokenizer.vocab.get("<image_placeholder>")
743 content_image_mask = [token == image_token_id for token in content_ids]
744 images_seq_mask += content_image_mask
745 num_image_tokens = sum(content_image_mask)
746 if num_image_tokens % 576:
747 raise ValueError("Each Janus image must have 576 placeholder tokens")
748 images_emb_mask.extend([[True] * 576 for _ in range(num_image_tokens // 576)])
750 labels += _janus_labels(content_ids, image_token_id, message["loss_mask"], task_type)
752 model_inputs = {
753 "input_ids": input_ids,
754 "attention_mask": attention_mask,
755 "labels": labels,
756 "images_seq_mask": images_seq_mask,
757 "images_emb_mask": images_emb_mask,
758 }
759 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}
760 return model_inputs
762 def get_jinja_template(self) -> str:
763 """Return the jinja template matching the Janus ChatML-style rendering."""
764 return (
765 "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"
766 "{% for message in messages %}"
767 "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] | trim + '<|im_end|>\n' }}"
768 "{% endfor %}"
769 "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
770 )
773@CHAT_TEMPLATE_REGISTRY.register("chatml")
774class ChatmlTemplate(ChatTemplate):
775 """Render messages with the ChatML <|im_start|>/<|im_end|> format."""
777 def encode_messages(self, messages: Sequence[Dict[str, str]], max_seq_len: int = 8192) -> Dict[str, List[int]]:
778 """Encode messages into input_ids, attention_mask, and loss-masked labels.
780 Args:
781 messages: Conversation records with ``role``, ``content``, and optional ``loss_mask``.
782 max_seq_len: Maximum sequence length kept from the end of the encoding.
784 Returns:
785 Model inputs with input_ids, attention_mask, and labels.
786 """
787 input_ids, attention_mask, labels = [], [], []
788 for message in messages:
789 content_str = "<|im_start|>" + message["role"] + "\n" + message["content"].strip() + "<|im_end|>\n"
790 content_ids = self.tokenizer.encode(content_str, add_special_tokens=False)
791 input_ids += content_ids
792 attention_mask += [1] * len(content_ids)
794 if "loss_mask" in message:
795 loss_mask = message["loss_mask"]
796 else:
797 loss_mask = 1 if message["role"] == "assistant" else 0
798 if loss_mask == 1:
799 labels += content_ids
800 else:
801 labels += [IGNORE_INDEX] * len(content_ids)
803 model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
804 model_inputs = {k: v[-max_seq_len:] for k, v in model_inputs.items()}
805 return model_inputs
807 def get_jinja_template(self) -> str:
808 """Return the jinja template matching the ChatML rendering."""
809 return (
810 "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}"
811 "{% for message in messages %}"
812 "{{ '<|im_start|>' + message['role'] + '\n' + message['content'] | trim + '<|im_end|>\n' }}"
813 "{% endfor %}"
814 "{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
815 )