Address siGit review on loss-mask PR
Two changes from the PR #6 review: - render_chatml: both branches pushed <|im_end|>, so push it once and keep only the assistant-span recording conditional. - Add a regression test that runs the real Qwen3 tokenizer to confirm the reviewers' concern doesn't happen: with add_special_tokens=false the ChatML markers still tokenize as single special tokens with real byte offsets, so the closing <|im_end|> lands inside the loss mask rather than escaping it via a synthesized (0,0) offset. Skips when the model isn't cached locally. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
keypair34 committed
Jul 11, 2026 at 08:13 UTC
ee4b36cff6b098be14ae65da6b81e1f4e26bb4c9
1 file changed
+101
-3
src/finetune.rs
+101
-3
@@ -553,9 +553,9 @@ fn render_chatml(messages: &[DatasetMessage]) -> (String, Vec<(usize, usize)>) {
553
554
let content_start = output.len();
555
output.push_str(&msg.content);
556
+ output.push_str("<|im_end|>");
557
558
if msg.role == "assistant" {
558
- output.push_str("<|im_end|>");
559
// Span end is *after* `<|im_end|>` so the closing marker is masked
560
// in for loss: the model has to learn to emit it, otherwise
561
// generations never terminate. Whether the marker tokenizes as one
@@ -563,8 +563,6 @@ fn render_chatml(messages: &[DatasetMessage]) -> (String, Vec<(usize, usize)>) {
563
// resulting token's byte offset falls inside this span.
564
let span_end = output.len();
565
assistant_spans.push((content_start, span_end));
566
- } else {
567
- output.push_str("<|im_end|>");
566
}
567
568
output.push('\n');
@@ -1409,4 +1407,104 @@ mod tests {
1407
);
1408
assert_eq!(split, vec![true, true, true, true]);
1409
}
1410
+
1411
+ /// Resolve a locally cached Qwen3 `tokenizer.json`, if present, so the
1412
+ /// real-tokenizer test below can skip gracefully on machines without the
1413
+ /// model. Honours `ONDE_TEST_MODEL_DIR` (a dir containing `tokenizer.json`)
1414
+ /// and otherwise falls back to the Onde App Group container cache.
1415
+ fn qwen3_tokenizer_path() -> Option<PathBuf> {
1416
+ if let Ok(dir) = std::env::var("ONDE_TEST_MODEL_DIR") {
1417
+ let p = PathBuf::from(dir).join("tokenizer.json");
1418
+ if p.exists() {
1419
+ return Some(p);
1420
+ }
1421
+ }
1422
+ let snapshots = dirs::home_dir()?
1423
+ .join("Library/Group Containers/group.com.ondeinference.apps")
1424
+ .join("models/hub/models--Qwen--Qwen3-0.6B/snapshots");
1425
+ std::fs::read_dir(&snapshots)
1426
+ .ok()?
1427
+ .flatten()
1428
+ .map(|e| e.path().join("tokenizer.json"))
1429
+ .find(|p| p.exists())
1430
+ }
1431
+
1432
+ /// Regression guard for the offset-based loss mask against the *real* Qwen3
1433
+ /// tokenizer (siGit review, PR #6): confirms that with
1434
+ /// `add_special_tokens=false` the ChatML markers still tokenize as single
1435
+ /// special tokens carrying real byte offsets — not synthesized `(0, 0)` /
1436
+ /// zero-length spans that would let the closing `<|im_end|>` escape the
1437
+ /// mask. Skips when the model is not cached locally.
1438
+ #[test]
1439
+ fn test_real_tokenizer_masks_assistant_and_im_end() {
1440
+ let Some(tok_path) = qwen3_tokenizer_path() else {
1441
+ eprintln!("Qwen3 tokenizer not found locally; skipping");
1442
+ return;
1443
+ };
1444
+ let tokenizer = Tokenizer::from_file(&tok_path).expect("load tokenizer");
1445
+
1446
+ let messages = vec![
1447
+ DatasetMessage {
1448
+ role: "user".to_string(),
1449
+ content: "hi".to_string(),
1450
+ },
1451
+ DatasetMessage {
1452
+ role: "assistant".to_string(),
1453
+ content: "Hello there!".to_string(),
1454
+ },
1455
+ ];
1456
+ let (rendered, spans) = render_chatml(&messages);
1457
+
1458
+ // add_special_tokens=false mirrors the training tokenization path.
1459
+ let encoding = tokenizer.encode(rendered.as_str(), false).unwrap();
1460
+ let offsets = encoding.get_offsets();
1461
+ let tokens = encoding.get_tokens();
1462
+
1463
+ // `<|im_end|>` must survive as a single token with a real, non-empty
1464
+ // offset — not a split into literal pieces and not a `(0, 0)` offset.
1465
+ let im_end_positions: Vec<usize> = tokens
1466
+ .iter()
1467
+ .enumerate()
1468
+ .filter(|(_, t)| t.as_str() == "<|im_end|>")
1469
+ .map(|(i, _)| i)
1470
+ .collect();
1471
+ assert_eq!(
1472
+ im_end_positions.len(),
1473
+ 2,
1474
+ "expected one <|im_end|> per turn as a single token"
1475
+ );
1476
+ for &i in &im_end_positions {
1477
+ let (s, e) = offsets[i];
1478
+ assert!(
1479
+ e > s,
1480
+ "<|im_end|> at token {i} has a zero-length offset {:?}",
1481
+ (s, e)
1482
+ );
1483
+ }
1484
+
1485
+ let mask = compute_loss_mask(offsets, &spans);
1486
+
1487
+ // The assistant's closing `<|im_end|>` (the second one) must be masked
1488
+ // in; if its offset were synthesized it would fall outside the span and
1489
+ // this would fail — exactly the escape the review flagged.
1490
+ let assistant_im_end = *im_end_positions.last().unwrap();
1491
+ assert!(
1492
+ mask[assistant_im_end],
1493
+ "assistant closing <|im_end|> escaped the loss mask"
1494
+ );
1495
+
1496
+ // The assistant content tokens are masked; the user turn and the
1497
+ // trailing newline after the assistant turn are not.
1498
+ let hello = tokens.iter().position(|t| t.contains("Hello")).unwrap();
1499
+ assert!(mask[hello], "assistant content should be masked in");
1500
+ let user_hi = tokens.iter().position(|t| t.as_str() == "hi").unwrap();
1501
+ assert!(!mask[user_hi], "user content must not be masked in");
1502
+ // Newline token immediately after the assistant <|im_end|>.
1503
+ if let Some(nl) = offsets.get(assistant_im_end + 1) {
1504
+ assert!(
1505
+ !compute_loss_mask(&[*nl], &spans)[0],
1506
+ "trailing newline after assistant turn must not be masked"
1507
+ );
1508
+ }
1509
+ }
1510
}