diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 5a37417..f327840 100644 --- a/tests/test_dataset.py +++ b/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) diff --git a/trainer/dataset.py b/trainer/dataset.py index d2e5c13..e4a12a7 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, @@ -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,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( @@ -339,4 +364,8 @@ def _tokenize_and_mask(examples): 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