From 6d44c353646fcc17e3b73bc96a092488009c8718 Mon Sep 17 00:00:00 2001 From: Anshul1336 <140470898+Anshul1336@users.noreply.github.com> Date: Sun, 20 Sep 2026 21:54:31 +0530 Subject: [PATCH 1/2] fix: warn when dataset samples get truncated past max_seq_length dataset.py was silently cutting off any sample longer than max_seq_length, no way to know it happened unless you went digging manually. now it measures the real token length before truncating, and logs a warning with the percentage of samples affected and the average tokens lost once tokenization finishes. added tests for the stats helper and for load_and_tokenize actually logging it with a fake tokenizer. fixes #181 --- tests/test_dataset.py | 63 ++++++++++++++++++++++++++++++++++++++++++- trainer/dataset.py | 28 +++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 5a37417..dbfe93c 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -8,7 +8,7 @@ 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 +54,64 @@ 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) diff --git a/trainer/dataset.py b/trainer/dataset.py index d2e5c13..dde9ce0 100644 --- a/trainer/dataset.py +++ b/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, @@ -299,7 +319,12 @@ def load_and_tokenize( } ) + _raw_lengths: list[int] = [] + def _tokenize_and_mask(examples): + _raw_lengths.extend( + len(ids) for ids in tokenizer(examples["text"], truncation=False)["input_ids"] + ) full_enc = tokenizer( examples["text"], truncation=True, @@ -339,4 +364,7 @@ def _tokenize_and_mask(examples): batched=True, remove_columns=raw.column_names, ) + warning = _summarize_truncation(_raw_lengths, max_seq_length) + if warning: + _logger.warning(warning) return tokenized From d0d02ad7de8132cbd50dfdc64bc252771e5f293f Mon Sep 17 00:00:00 2001 From: Sahil Kumar Singh <60318530+SahilKumar75@users.noreply.github.com> Date: Sun, 20 Sep 2026 23:33:26 +0530 Subject: [PATCH 2/2] fix: preserve truncation warning across dataset cache --- tests/test_dataset.py | 22 ++++++++++++++++++++++ trainer/dataset.py | 13 +++++++------ 2 files changed, 29 insertions(+), 6 deletions(-) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index dbfe93c..f327840 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -5,6 +5,7 @@ """ import csv +import json import pytest @@ -115,3 +116,24 @@ def test_load_and_tokenize_warns_when_a_sample_is_truncated(caplog): ) 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) diff --git a/trainer/dataset.py b/trainer/dataset.py index dde9ce0..e4a12a7 100644 --- a/trainer/dataset.py +++ b/trainer/dataset.py @@ -319,12 +319,8 @@ def load_and_tokenize( } ) - _raw_lengths: list[int] = [] - def _tokenize_and_mask(examples): - _raw_lengths.extend( - len(ids) for ids in tokenizer(examples["text"], truncation=False)["input_ids"] - ) + untruncated = tokenizer(examples["text"], truncation=False) full_enc = tokenizer( examples["text"], truncation=True, @@ -357,6 +353,10 @@ 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( @@ -364,7 +364,8 @@ def _tokenize_and_mask(examples): batched=True, remove_columns=raw.column_names, ) - warning = _summarize_truncation(_raw_lengths, max_seq_length) + warning = _summarize_truncation(tokenized["_untruncated_length"], max_seq_length) + tokenized = tokenized.remove_columns("_untruncated_length") if warning: _logger.warning(warning) return tokenized