[data] feat: support assistant turn loss mask - #89
Conversation
|
@claude review this PR |
|
Claude finished @nullnonenilNULL's task in 2m 8s —— View job Verdict: COMMENTSummaryThis PR adds per-assistant-turn 🟠 Major
🟡 Minor
🟢 Nits
TestsNo automated tests are in the diff. The PR description mentions manual CPU helper checks and a local tokenizer validation, but these aren't committed. Given the partial-label edge case above, at least one test covering mixed explicit/implicit assistant masks would be valuable. Checklist
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| cls._message_loss_mask_to_bool(message["loss_mask"]) | ||
| ) | ||
| else: | ||
| assistant_loss_flags.append(True) |
There was a problem hiding this comment.
🟠 Major: Mixed-explicit semantics create a silent training gap.
When only some assistant messages carry loss_mask (and others don't), the unlabelled turns default to True (append(True) at line 616). This means has_explicit_assistant_loss_mask is set to True, triggering the length-check and the filter — but the unlabelled turns silently inherit "train on this turn", which may not be what the author intended when providing a partial mask.
The PR description says "Preserve existing behavior when no assistant message has explicit loss_mask", but it doesn't address the partial case (some turns labelled, some not). The current behavior is all-or-nothing is not enforced: a dataset with one masked turn and one un-annotated turn will pass validation and silently train on the unannotated turn.
Suggested fix: either (a) require that if any assistant message carries loss_mask, all assistant messages must carry it (raise if the counts differ before the length check), or (b) document this partial-label behavior explicitly in the docstring so callers understand the "missing = True" default applies even in mixed data.
| """ | ||
| if isinstance(value, bool): | ||
| return value | ||
| if isinstance(value, int) and value in (0, 1): |
There was a problem hiding this comment.
🟡 Minor: _message_loss_mask_to_bool accepts any int value of 0 or 1, but silently rejects other integers (e.g. 2) by falling through to the ValueError. That is correct, but the isinstance(value, int) branch also matches bool values since bool is a subclass of int in Python — however this is harmless here because the isinstance(value, bool) check comes first.
More importantly, a float like 1.0 or 0.0 will fall through to the ValueError. Whether this is intentional should be documented or validated more explicitly. If floats should be rejected, a note in the docstring would prevent future surprises.
| ) | ||
| ) | ||
| ori_total_len = len(input_ids) | ||
| if has_explicit_loss_mask and train_on_prompt: |
There was a problem hiding this comment.
🟡 Minor: The conflict check fires after _tokenize_with_generation_indices has already done the tokenization work (rendering, aligning masks, etc.). For train_on_prompt / history_mask_loss, that work is wasted. Moving the conflict guard before calling _tokenize_with_generation_indices would save time and make it obvious that these options are mutually exclusive at the call site — not a correctness issue, just a needless cost on every conflicting call.
| tools: Optional[Sequence[Dict[str, Any]]] = None, | ||
| ) -> Tuple[List[int], List[int]]: | ||
| ) -> Tuple[List[int], List[int], bool]: | ||
| """Render OpenAI chat messages and build assistant-token masks.""" |
There was a problem hiding this comment.
🟢 Nit: The returned bool third element of the tuple now has a different meaning from the first two elements (it's a metadata flag, not token data). The return type annotation Tuple[List[int], List[int], bool] at line 625 is correct but the docstring still says only "Render OpenAI chat messages and build assistant-token masks" — it doesn't mention the new third return value. Worth a one-line update for future readers.
2 participants
Summary
loss_masksupport inHFChatTemplate.encode_openai()loss_mask=1keeps that assistant turn trainable,loss_mask=0masks it outloss_maskloss_maskon non-assistant messages for uniform input schemas; these fields are ignored because only assistant messages map to HF generation rangestrain_on_prompt/history_mask_lossValidation
compile(..., "loongforge/data/chat_template.py", "exec")assistant_flags=[False, False, True])