fix: warn when dataset samples get truncated past max_seq_length

tests/test_dataset.py

@@ -5,10 +5,11 @@
 """
 
 import csv
+import json
 
 import pytest
 
-from trainer.dataset import format_prompt, load_and_tokenize
+from trainer.dataset import _summarize_truncation, format_prompt, load_and_tokenize
 
 
 def test_format_prompt_uses_instruction_and_output():
@@ -54,3 +55,85 @@ def _boom(*_a, **_k):
 
     with pytest.raises(ValueError, match="Output column"):
         load_and_tokenize(str(csv_path), _boom, output_col="output")
+
+
+# ── truncation warning (#181) ───────────────────────────────────────
+
+
+def test_summarize_truncation_returns_none_when_nothing_truncated():
+    assert _summarize_truncation([10, 20, 30], max_seq_length=50) is None
+
+
+def test_summarize_truncation_reports_percentage_and_average_overflow():
+    # two samples fit, two exceed max_seq_length by 10 and 20 tokens
+    msg = _summarize_truncation([50, 60, 70, 50], max_seq_length=50)
+    assert msg is not None
+    assert "50.0%" in msg
+    assert "2/4" in msg
+    assert "max_seq_length=50" in msg
+    assert "avg 15.0 tokens" in msg
+
+
+class _FakeTokenizer:
+    """Minimal stand-in for a HF tokenizer: token count = word count (+BOS)."""
+
+    bos_token_id = 0
+
+    def __call__(
+        self, texts, truncation=True, max_length=None, padding=None, add_special_tokens=True
+    ):
+        if isinstance(texts, str):
+            texts = [texts]
+        out = []
+        for t in texts:
+            ids = ([self.bos_token_id] if add_special_tokens else []) + list(
+                range(1, len(t.split()) + 1)
+            )
+            if truncation and max_length is not None:
+                ids = ids[:max_length]
+            if padding == "max_length" and max_length is not None:
+                ids = ids + [0] * (max_length - len(ids))
+            out.append(ids)
+        return {"input_ids": out}
+
+
+def test_load_and_tokenize_warns_when_a_sample_is_truncated(caplog):
+    from datasets import Dataset
+
+    preloaded = Dataset.from_dict(
+        {
+            "instruction": ["short one", "this instruction has quite a lot more words in it"],
+            "output": ["ok", "also a fairly long response with several words"],
+        }
+    )
+
+    with caplog.at_level("WARNING", logger="trainer.dataset"):
+        load_and_tokenize(
+            "unused.json",
+            _FakeTokenizer(),
+            max_seq_length=6,
+            preloaded=preloaded,
+        )
+
+    assert any("exceeded max_seq_length=6" in r.message for r in caplog.records)
+
+
+def test_load_and_tokenize_warns_when_tokenization_is_cached(tmp_path, caplog):
+    jsonl_path = tmp_path / "data.jsonl"
+    jsonl_path.write_text(
+        json.dumps(
+            {
+                "instruction": "this instruction contains enough words to overflow",
+                "output": "and so does this response",
+            }
+        )
+        + "\n"
+    )
+    tokenizer = _FakeTokenizer()
+
+    load_and_tokenize(str(jsonl_path), tokenizer, max_seq_length=6)
+    caplog.clear()
+    with caplog.at_level("WARNING", logger="trainer.dataset"):
+        load_and_tokenize(str(jsonl_path), tokenizer, max_seq_length=6)
+
+    assert any("exceeded max_seq_length=6" in r.message for r in caplog.records)

trainer/dataset.py

@@ -1,3 +1,5 @@
+import logging
+
 import pandas as pd
 from datasets import Dataset, load_dataset
 from transformers import PreTrainedTokenizer
@@ -6,6 +8,8 @@
 # (single source of truth — see trainer/prompt_templates.py).
 from trainer.prompt_templates import PROMPT_TEMPLATE, PROMPT_TEMPLATES  # noqa: F401
 
+_logger = logging.getLogger(__name__)
+
 
 def _format_prompt_prefix(row: dict, template: str = "alpaca") -> str:
     """Return just the instruction part of the formatted text (no output).
@@ -264,6 +268,22 @@ def _process(batch):
     return raw.map(_process, batched=True, remove_columns=raw.column_names)
 
 
+def _summarize_truncation(lengths: list[int], max_seq_length: int) -> str | None:
+    """Return a warning message describing how many samples exceeded
+    ``max_seq_length`` and were silently truncated, or None if none were."""
+    total = len(lengths)
+    overflow = [length - max_seq_length for length in lengths if length > max_seq_length]
+    if not overflow:
+        return None
+    pct = 100 * len(overflow) / total
+    avg_overflow = sum(overflow) / len(overflow)
+    return (
+        f"{pct:.1f}% of samples ({len(overflow)}/{total}) exceeded "
+        f"max_seq_length={max_seq_length} and were truncated "
+        f"(avg {avg_overflow:.1f} tokens lost per truncated sample)."
+    )
+
+
 def load_and_tokenize(
     file_path: str,
     tokenizer: PreTrainedTokenizer,
@@ -300,6 +320,7 @@ def load_and_tokenize(
     )
 
     def _tokenize_and_mask(examples):
+        untruncated = tokenizer(examples["text"], truncation=False)
         full_enc = tokenizer(
             examples["text"],
             truncation=True,
@@ -332,11 +353,19 @@ def _tokenize_and_mask(examples):
             lbl[bos_offset : bos_offset + prompt_len] = [-100] * prompt_len
             labels.append(lbl)
         full_enc["labels"] = labels
+        # Store lengths in the map result instead of mutating an outer list.
+        # Hugging Face may serve this map from cache without calling this
+        # function, so side effects would make the warning disappear on reload.
+        full_enc["_untruncated_length"] = [len(ids) for ids in untruncated["input_ids"]]
         return full_enc
 
     tokenized = raw.map(
         _tokenize_and_mask,
         batched=True,
         remove_columns=raw.column_names,
     )
+    warning = _summarize_truncation(tokenized["_untruncated_length"], max_seq_length)
+    tokenized = tokenized.remove_columns("_untruncated_length")
+    if warning:
+        _logger.warning(warning)
     return tokenized