Fix zero-length span test assertion in loss mask test
The zero-length-span test asserted the wrong expected mask. Per the documented edge case, a token spanning the position of a zero-length assistant span (empty content) should mask true, since the model needs to learn to emit <|im_end|> right after the role header even when the assistant said nothing. The compute_loss_mask overlap logic already did this correctly; only the test's assertion was inverted. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01KfXzV7X5JZEenoH6MSdNtg
Claude committed
Jul 10, 2026 at 17:17 UTC
1651c748c8119560a17bd398b24a7af11bd853c0
1 file changed
+9
-3
src/finetune.rs
+9
-3
@@ -1331,13 +1331,19 @@ mod tests {
1331
1332
#[test]
1333
fn test_compute_loss_mask_zero_length_span() {
1334
- // A zero-length span (empty assistant content) should not match tokens
1334
+ // A zero-length span occurs when assistant content is empty — the
1335
+ // span collapses to a single point where content_start == end_tag_start.
1336
+ // Any token whose range contains that point (e.g. the immediately
1337
+ // following `<|im_end|>` token) must still be masked true, since the
1338
+ // model needs to learn to emit `<|im_end|>` right after the role
1339
+ // header even for an empty response.
1340
let offsets = vec![(10, 20), (20, 30)];
1341
let spans = vec![(15, 15)]; // zero-length at byte 15
1342
let mask = compute_loss_mask(&offsets, &spans);
1343
1339
- // Zero-length spans match nothing (start == end, so condition fails)
1340
- assert_eq!(mask, vec![false, false]);
1344
+ // Token [10, 20) contains the zero-length span's point (15) → true.
1345
+ // Token [20, 30) does not → false.
1346
+ assert_eq!(mask, vec![true, false]);
1347
}
1348
1349
#[test]