@@ -77,6 +77,9 @@ class AssistantMaskBoundaryConfig:
7777 ``role_end_token_variants`` provides ordered fallback end sequences for
7878 formats whose final turn may omit a trailing delimiter such as a
7979 newline. The primary ``role_end_tokens`` sequence is always preferred.
80+ ``rendered_turn_end_roles`` covers formats such as GLM-5.2 whose
81+ assistant turn has no explicit end token. For those roles, the
82+ conversation-prefix render determines the exclusive content boundary.
8083
8184 Example::
8285
@@ -105,6 +108,7 @@ class AssistantMaskBoundaryConfig:
105108 trim_leading_token_ids : Sequence [int ] = ()
106109 trim_leading_token_sequences : Sequence [Sequence [int ]] = ()
107110 role_end_token_variants : Mapping [str , Sequence [Sequence [int ]]] = field (default_factory = dict )
111+ rendered_turn_end_roles : Sequence [str ] = ()
108112
109113
110114@dataclass (frozen = True )
@@ -491,7 +495,9 @@ def resolve_chat_template_kwargs(template_owner: Any, template_kwargs: Mapping[s
491495
492496 MBridge exposes ``truncate_history_thinking`` consistently. Templates that
493497 implement the inverse ``preserve_thinking`` control receive the equivalent
494- negated value only at the rendering boundary.
498+ negated value only at the rendering boundary, while GLM-style
499+ ``clear_thinking`` receives the same value. Custom tokenizers such as
500+ Kimi-K3 use ``thinking`` instead of ``enable_thinking``.
495501
496502 Args:
497503 template_owner: Processor or tokenizer that owns the chat template.
@@ -501,13 +507,29 @@ def resolve_chat_template_kwargs(template_owner: Any, template_kwargs: Mapping[s
501507 A copied mapping adapted to the selected template.
502508 """
503509 kwargs = dict (template_kwargs )
504- if "truncate_history_thinking" not in kwargs :
505- return kwargs
506510 tokenizer = get_processor_tokenizer (template_owner )
507511 template = _get_chat_template (template_owner , tokenizer )
508- if template is not None and "preserve_thinking" in template and "truncate_history_thinking" not in template :
509- truncate_history = kwargs .pop ("truncate_history_thinking" )
510- kwargs ["preserve_thinking" ] = not truncate_history
512+
513+ if "truncate_history_thinking" in kwargs and (template is None or "truncate_history_thinking" not in template ):
514+ truncate_history = kwargs ["truncate_history_thinking" ]
515+ if template is not None and "preserve_thinking" in template :
516+ kwargs .pop ("truncate_history_thinking" )
517+ kwargs ["preserve_thinking" ] = not truncate_history
518+ elif template is not None and "clear_thinking" in template :
519+ kwargs .pop ("truncate_history_thinking" )
520+ kwargs ["clear_thinking" ] = truncate_history
521+
522+ if "enable_thinking" in kwargs and (template is None or "enable_thinking" not in template ):
523+ parameters : set [str ] = set ()
524+ for owner in (template_owner , tokenizer ):
525+ apply_chat_template = _get_explicit_attribute (owner , "apply_chat_template" )
526+ try :
527+ parameters .update (inspect .signature (apply_chat_template ).parameters if apply_chat_template else {})
528+ except (TypeError , ValueError ):
529+ continue
530+ if "thinking" in parameters and "enable_thinking" not in parameters :
531+ enable_thinking = kwargs .pop ("enable_thinking" )
532+ kwargs .setdefault ("thinking" , enable_thinking )
511533 return kwargs
512534
513535
@@ -556,7 +578,8 @@ def normalize_chat_conversation(
556578 """Normalize supported text-chat schemas to OpenAI-style messages.
557579
558580 Rows may use ``messages``, singular ``conversation``, or legacy plural
559- ``conversations`` with ``from``/``value`` keys.
581+ ``conversations`` with ``from``/``value`` keys. Null OpenAI message content
582+ is converted to an empty string before chat-template rendering.
560583 """
561584 if isinstance (example_or_conversation , Mapping ):
562585 source = example_or_conversation
@@ -581,6 +604,8 @@ def normalize_chat_conversation(
581604 raise ValueError ("Chat conversation turns must be dictionaries." )
582605 turn = dict (raw_turn )
583606 if "role" in turn and "content" in turn :
607+ if turn ["content" ] is None :
608+ turn ["content" ] = ""
584609 _normalize_assistant_tool_call_arguments (turn )
585610 normalized .append (turn )
586611 continue
@@ -853,7 +878,12 @@ def tokenize_chat_example(
853878 )
854879
855880 chat_example : dict [str , Any ] = {"conversation" : conversation }
856- chat_example .update (template_kwargs )
881+ mask_template_kwargs = dict (template_kwargs )
882+ mask_tools = mask_template_kwargs .pop ("tools" , None )
883+ if mask_template_kwargs :
884+ chat_example ["chat_template_kwargs" ] = mask_template_kwargs
885+ if mask_tools is not None :
886+ chat_example ["tools" ] = mask_tools
857887 effective_boundary_config = boundary_config or infer_assistant_mask_boundary_config (processor )
858888 if returned_mask is not None and effective_boundary_config is None :
859889 assistant_mask = returned_mask
@@ -1038,9 +1068,42 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar
10381068 """Infer assistant role boundaries from known HF chat-template formats."""
10391069 tokenizer = get_processor_tokenizer (processor )
10401070 template = _get_chat_template (tokenizer , processor )
1071+
10411072 if template is None :
1073+ # Kimi-K3 owns rendering in a custom tokenizer rather than a Jinja
1074+ # ``chat_template``. Detect its XTML structural vocabulary instead of a
1075+ # Python class name so local/remote-code wrappers retain the same behavior.
1076+ apply_chat_template = _get_explicit_attribute (tokenizer , "apply_chat_template" )
1077+ try :
1078+ template_parameters = inspect .signature (apply_chat_template ).parameters if apply_chat_template else {}
1079+ except (TypeError , ValueError ):
1080+ template_parameters = {}
1081+ if "thinking" in template_parameters :
1082+ get_vocab = _get_explicit_attribute (tokenizer , "get_vocab" )
1083+ try :
1084+ vocab = get_vocab () if callable (get_vocab ) else {}
1085+ except (AttributeError , TypeError , ValueError ):
1086+ vocab = {}
1087+ kimi_markers = ("<|open|>" , "<|close|>" , "<|sep|>" , "<|end_of_msg|>" )
1088+ if isinstance (vocab , Mapping ) and all (marker in vocab for marker in kimi_markers ):
1089+ return assistant_mask_boundary_config_from_markers (
1090+ processor ,
1091+ assistant_start = '<|open|>message role="assistant"<|sep|>' ,
1092+ assistant_end = "<|close|>message<|sep|><|end_of_msg|>" ,
1093+ )
10421094 return None
10431095
1096+ # GLM-5.2 does not terminate completed assistant turns with a structural
1097+ # token. Prefix rendering is therefore the only loss boundary that remains
1098+ # exact for both terminal and historical turns.
1099+ if all (marker in template for marker in ("<|assistant|>" , "clear_thinking" , "<think>" )):
1100+ return assistant_mask_boundary_config_from_markers (
1101+ processor ,
1102+ assistant_start = "<|assistant|>" ,
1103+ assistant_end = None ,
1104+ rendered_turn_end_roles = (CHATML_ASSISTANT_ROLE ,),
1105+ )
1106+
10441107 candidates : tuple [tuple [tuple [str , ...], str , str , tuple [str , ...], Mapping [str , str ]], ...] = (
10451108 (
10461109 ("<|im_start|>assistant" , "<|im_end|>" ),
@@ -1087,6 +1150,18 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar
10871150 ("<|end|>" , "<|call|>" ),
10881151 {role : f"<|start|>{ role } " for role in ("system" , "developer" , "user" , "tool" )},
10891152 ),
1153+ (
1154+ ("]~!b[" , "]~b]" , "[e~[" , "<mm:think>" ),
1155+ "]~b]ai\n " ,
1156+ "[e~[\n " ,
1157+ (),
1158+ {
1159+ "system" : "]~!b[]~b]system\n " ,
1160+ "developer" : "]~b]developer\n " ,
1161+ "user" : "]~b]user\n " ,
1162+ "tool" : "]~b]tool" ,
1163+ },
1164+ ),
10901165 )
10911166 for required_markers , assistant_start , assistant_end , assistant_end_fallbacks , other_role_starts in candidates :
10921167 if not all (marker in template for marker in required_markers ):
@@ -1126,25 +1201,26 @@ def assistant_mask_boundary_config_from_markers(
11261201 processor : Any ,
11271202 * ,
11281203 assistant_start : str ,
1129- assistant_end : str ,
1204+ assistant_end : str | None ,
11301205 assistant_end_fallbacks : Sequence [str ] = (),
11311206 role_start_markers : Mapping [str , str ] | None = None ,
11321207 loss_roles : Sequence [str ] = (CHATML_ASSISTANT_ROLE ,),
11331208 include_start_tokens_for_roles : Sequence [str ] = (),
11341209 include_end_tokens_for_roles : Sequence [str ] = (CHATML_ASSISTANT_ROLE ,),
11351210 trim_leading_token_ids : Sequence [int ] = (),
11361211 trim_leading_token_sequences : Sequence [Sequence [int ]] = (),
1212+ rendered_turn_end_roles : Sequence [str ] = (),
11371213) -> AssistantMaskBoundaryConfig :
11381214 """Build explicit assistant boundary config from model-declared marker strings."""
11391215 tokenizer = get_processor_tokenizer (processor )
11401216 start_tokens = tokenize_text_without_special_tokens (tokenizer , assistant_start )
1141- end_tokens = tokenize_text_without_special_tokens (tokenizer , assistant_end )
1217+ end_tokens = tokenize_text_without_special_tokens (tokenizer , assistant_end ) if assistant_end is not None else []
11421218 end_token_variants = [
11431219 token_ids
11441220 for fallback in assistant_end_fallbacks
11451221 if (token_ids := tokenize_text_without_special_tokens (tokenizer , fallback ))
11461222 ]
1147- if not start_tokens or not end_tokens :
1223+ if not start_tokens or ( not end_tokens and CHATML_ASSISTANT_ROLE not in rendered_turn_end_roles ) :
11481224 raise ValueError (
11491225 "Unable to tokenize assistant loss-mask boundary markers. "
11501226 "Check the model collator's assistant_start and assistant_end markers."
@@ -1157,13 +1233,14 @@ def assistant_mask_boundary_config_from_markers(
11571233 role_start_tokens [role ] = marker_tokens
11581234 return AssistantMaskBoundaryConfig (
11591235 role_start_tokens = role_start_tokens ,
1160- role_end_tokens = {role : end_tokens for role in role_start_tokens },
1236+ role_end_tokens = {role : end_tokens for role in role_start_tokens if role not in rendered_turn_end_roles },
11611237 loss_roles = loss_roles ,
11621238 include_start_tokens_for_roles = include_start_tokens_for_roles ,
11631239 include_end_tokens_for_roles = include_end_tokens_for_roles ,
11641240 trim_leading_token_ids = trim_leading_token_ids ,
11651241 trim_leading_token_sequences = trim_leading_token_sequences ,
11661242 role_end_token_variants = {role : end_token_variants for role in role_start_tokens },
1243+ rendered_turn_end_roles = rendered_turn_end_roles ,
11671244 )
11681245
11691246
@@ -1316,13 +1393,14 @@ def _assistant_mask_from_conversation_turns(
13161393 }
13171394 include_start_roles = set (boundary_config .include_start_tokens_for_roles )
13181395 include_end_roles = set (boundary_config .include_end_tokens_for_roles )
1396+ rendered_turn_end_roles = set (boundary_config .rendered_turn_end_roles )
13191397 rendered_mask = torch .zeros (len (rendered_ids ), dtype = torch .float32 )
13201398
13211399 for turn_index , turn in loss_turns :
13221400 role = turn ["role" ]
13231401 start_tokens = role_start_tokens .get (role )
13241402 end_patterns = [role_end_tokens .get (role , []), * role_end_token_variants .get (role , [])]
1325- if not start_tokens or not any (end_patterns ):
1403+ if not start_tokens or ( not any (end_patterns ) and role not in rendered_turn_end_roles ):
13261404 return None
13271405
13281406 try :
@@ -1360,20 +1438,23 @@ def _assistant_mask_from_conversation_turns(
13601438 if content_start > turn_limit :
13611439 return None
13621440
1363- candidate_ends : list [tuple [int , int , int , list [int ]]] = []
1364- for priority , pattern in enumerate (end_patterns ):
1365- if not pattern :
1366- continue
1367- search_start = content_start
1368- while search_start < turn_limit :
1369- end_start , after_end = find_token_span (rendered_ids , pattern , search_start )
1370- if end_start < 0 or after_end > turn_limit :
1371- break
1372- candidate_ends .append ((end_start , - priority , after_end , pattern ))
1373- search_start = end_start + 1
1374- if not candidate_ends :
1375- return None
1376- content_end , _ , after_end , _ = max (candidate_ends )
1441+ if role in rendered_turn_end_roles :
1442+ content_end = after_end = turn_limit
1443+ else :
1444+ candidate_ends : list [tuple [int , int , int , list [int ]]] = []
1445+ for priority , pattern in enumerate (end_patterns ):
1446+ if not pattern :
1447+ continue
1448+ search_start = content_start
1449+ while search_start < turn_limit :
1450+ end_start , after_end = find_token_span (rendered_ids , pattern , search_start )
1451+ if end_start < 0 or after_end > turn_limit :
1452+ break
1453+ candidate_ends .append ((end_start , - priority , after_end , pattern ))
1454+ search_start = end_start + 1
1455+ if not candidate_ends :
1456+ return None
1457+ content_end , _ , after_end , _ = max (candidate_ends )
13771458
13781459 loss_start = _trim_leading_token_sequences (
13791460 rendered_ids ,
@@ -1391,7 +1472,7 @@ def _assistant_mask_from_conversation_turns(
13911472 rendered_mask [loss_start :content_end ] = 1.0
13921473 if role in include_start_roles :
13931474 rendered_mask [role_start : role_start + matched_start_length ] = 1.0
1394- if role in include_end_roles :
1475+ if role in include_end_roles and role not in rendered_turn_end_roles :
13951476 rendered_mask [content_end :after_end ] = 1.0
13961477
13971478 aligned_mask = _align_assistant_mask_to_padded_ids (rendered_ids , rendered_mask .tolist (), ids , tokenizer )
@@ -1551,6 +1632,12 @@ def _assistant_mask_from_boundary_config(
15511632 for role , raw_variants in boundary_config .role_end_token_variants .items ()
15521633 }
15531634 loss_roles = set (boundary_config .loss_roles )
1635+ rendered_turn_end_roles = set (boundary_config .rendered_turn_end_roles )
1636+ if rendered_turn_end_roles :
1637+ # A raw token scan cannot distinguish the end of an unterminated turn
1638+ # from following non-loss roles. Only the conversation-prefix renderer
1639+ # can establish this boundary safely.
1640+ return None
15541641 if not role_start_tokens or not role_end_tokens or not loss_roles :
15551642 raise ValueError (
15561643 "AssistantMaskBoundaryConfig must provide non-empty role_start_tokens, role_end_tokens, and loss_roles."
@@ -1657,7 +1744,7 @@ def build_assistant_loss_mask(
16571744
16581745 mask = _assistant_mask_from_hf_chat_template (example_or_conversation , ids , processor , tokenizer )
16591746 if boundary_config is not None :
1660- boundary_scan_is_safe = (
1747+ boundary_scan_is_safe = not boundary_config . rendered_turn_end_roles and (
16611748 not conversation
16621749 or _conversation_contains_boundary_tokens (example_or_conversation , tokenizer , boundary_config ) is False
16631750 )
0 commit comments