Audit: every validation example scored once; shift blocks cut by packing; inference and data fixes; tests

.github/workflows/docker-build.yaml

@@ -9,6 +9,7 @@ on:
   pull_request:
     paths:
       - 'Dockerfile'
+      - 'pyproject.toml'
 
 jobs:
   docker-build:

.github/workflows/lint.yaml

@@ -16,4 +16,6 @@ jobs:
       - uses: actions/checkout@v5
 
       - name: Lint code
-        uses: astral-sh/ruff-action@v3
+        uses: astral-sh/ruff-action@v3
+        with:
+          version: "0.15.8"  # Pinned: a new ruff release adds rules

.github/workflows/publish-docker.yaml

@@ -24,7 +24,7 @@ jobs:
 
     steps:
       - name: Checkout repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
 
       # Reclaim ~50–80 GB on the GitHub runner
       - name: Free disk space

benchmarks/flops.py

@@ -14,7 +14,7 @@
 from utf8_tokenizer.tokenizer import PAD_TOKEN_ID
 
 from welt.processor import collate_fn, next_word_labels
-from welt_training.data_utils import pack_words, train_texts
+from welt_training.data_utils import train_texts
 from welt_training.extendable_yaml import load_yaml
 from welt_training.train import build_dataset_provider
 
@@ -38,14 +38,11 @@ def transformer_flops(config, lengths) -> float:
 def step_flops(config: dict, batches: int = 20) -> dict[str, float]:
     model, data = config["model"], config["data"]
     provider = build_dataset_provider(model, data)
-    processor = provider.processor()
     # The training examples (of a single data parallel rank), as the dataloader makes them
-    texts = train_texts(provider, rank=0, world_size=1)
-    examples = (processor.process_single_example(*packed)
-                for packed in pack_words(map(processor.pretokenize, texts), provider.seq_length))
+    examples = provider.make_examples(train_texts(provider, rank=0, world_size=1))
     train_config = config.get("train") or {}
-    micro_batch = train_config.get("micro_batch_size", 32)
-    global_batch = train_config.get("global_batch_size", micro_batch)
+    micro_batch = train_config.get("micro_batch_size", 32)  # build_config's defaults
+    global_batch = train_config.get("global_batch_size", 32)
     configs = {name: AutoConfig.from_pretrained(model[name], trust_remote_code=model.get("trust_remote_code", False))
                for name in ["bytes_encoder", "image_encoder", "latent_transformer", "bytes_decoder"] if model.get(name)}
 

tests/conftest.py

@@ -75,10 +75,12 @@ def tiny_config(tmp_path_factory):
 def oracle_labels(words: list[str], seq_lengths: list[int], continues: list[bool] | None = None) -> list[str | None]:
     """Each word's label, from the words alone: the next word of its sequence, "" after its last word (None if the
     sequence continues in another example), and None for PAD words and for the words of a shift block before its
-    ShiftIn (they see each other)."""
+    ShiftIn (they see each other). A sequence cut from a longer document can begin within a block (a ShiftIn before
+    any ShiftOut)."""
     labels, start = [], 0
     for length, continued in zip(seq_lengths, continues or [False] * len(seq_lengths), strict=True):
-        sequence, inside = words[start:start + length], False
+        sequence = words[start:start + length]
+        inside = SI in sequence and SO not in sequence[:sequence.index(SI)]
         for i, word in enumerate(sequence):
             inside = (inside or word == SO) and word != SI
             last_of_continued = continued and i == length - 1
@@ -89,17 +91,21 @@ def oracle_labels(words: list[str], seq_lengths: list[int], continues: list[bool
 
 def oracle_mask(words: list[str], seq_lengths: list[int]) -> torch.Tensor:
     """(L, L) the words each word sees in the latent transformer, from the words alone: itself and the earlier words
-    of its sequence, and every word of its shift block (ShiftOut to ShiftIn)."""
+    of its sequence, and every word of its shift block (ShiftOut to ShiftIn; a sequence cut from a longer document
+    can begin or end within one)."""
     allowed = torch.zeros(len(words), len(words), dtype=torch.bool)
     start = 0
     for length in seq_lengths:
-        shift_out = None
-        for i in range(start, start + length):
+        shift_out, end = None, start + length
+        for i in range(start, end):
             allowed[i, start:i + 1] = True
-            if words[i] == SO:
+            if words[i] == SO and shift_out is None:
                 shift_out = i
-            elif words[i] == SI and shift_out is not None:
-                allowed[shift_out:i + 1, shift_out:i + 1] = True
+            elif words[i] == SI:
+                block_start = start if shift_out is None else shift_out
+                allowed[block_start:i + 1, block_start:i + 1] = True
                 shift_out = None
+        if shift_out is not None:
+            allowed[shift_out:end, shift_out:end] = True
         start += length
     return allowed

tests/test_config.py

@@ -1,9 +1,12 @@
 import glob
+from pathlib import Path
 
 import pytest
 
 from welt_training.extendable_yaml import load_yaml
 
+ROOT = Path(__file__).resolve().parents[1]
+
 
 def test_nested_extends_and_overrides(tmp_path):
     (tmp_path / "base.yaml").write_text("model:\n  a: 1\n  b: 2\ndata:\n  seq_length: 128\n")
@@ -15,23 +18,34 @@ def test_nested_extends_and_overrides(tmp_path):
                       "output_dir": "out", "train": {"train_iters": 5}, "optimizer": {"lr": 1e-4}}
 
 
-CONFIGS = sorted(glob.glob("welt_training/experiments/*/*.yaml") + ["benchmarks/welt-bench.yaml"])
+CONFIGS = sorted([*glob.glob(str(ROOT / "welt_training/experiments/*/*.yaml")),
+                  str(ROOT / "benchmarks/welt-bench.yaml")])
 
 
 @pytest.mark.parametrize("path", CONFIGS)
-def test_experiment_configs_build(path):
+def test_experiment_configs_reference_existing_models(path):
+    """Model configs given as JSON files exist (paths are relative to the repository)."""
+    for value in load_yaml(path)["model"].values():
+        if isinstance(value, str) and value.endswith(".json"):
+            assert (ROOT / value).exists(), value
+
+
+@pytest.mark.parametrize("path", CONFIGS)
+def test_experiment_configs_build(path, monkeypatch):
     """Every shipped config builds a Megatron-Bridge config (unknown options raise)."""
     pytest.importorskip("megatron.bridge", reason="Requires the NeMo container")
     from welt_training import baseline, train
 
     config = load_yaml(path)
+    monkeypatch.chdir(ROOT)  # Model JSON paths are relative to the repository
     (baseline if "transformer" in config["model"] else train).build(config)  # baseline: a causal LM
 
 
-def test_null_sections_and_model_overrides():
+def test_null_sections_and_model_overrides(monkeypatch):
     pytest.importorskip("megatron.bridge", reason="Requires the NeMo container")
     from welt_training import train
 
+    monkeypatch.chdir(ROOT)
     config = load_yaml("welt_training/experiments/easy-tasks/string-repetition.yaml",
                        ["optimizer=null", "model.hidden_dropout=0.2"])
     cfg, _ = train.build(config)
@@ -40,3 +54,15 @@ def test_null_sections_and_model_overrides():
         train.build(load_yaml("welt_training/experiments/easy-tasks/string-repetition.yaml", ["train.typo=1"]))
     with pytest.raises(ValueError, match="Unknown config sections"):
         train.build(load_yaml("welt_training/experiments/easy-tasks/string-repetition.yaml", ["optimiser.lr=0.1"]))
+
+
+def test_baseline_model_overrides(monkeypatch):
+    pytest.importorskip("megatron.bridge", reason="Requires the NeMo container")
+    from welt_training import baseline
+
+    monkeypatch.chdir(ROOT)
+    path = "welt_training/experiments/language-modeling/climbmix-baseline.yaml"
+    cfg, _ = baseline.build(load_yaml(path, ["model.hidden_dropout=0.2"]))
+    assert cfg.model.hidden_dropout == 0.2
+    with pytest.raises(ValueError, match="Unknown model option"):
+        baseline.build(load_yaml(path, ["model.typo=1"]))

tests/test_data.py

@@ -1,9 +1,11 @@
 from itertools import islice
 
 import pytest
+import torch
 from datasets import Dataset
-from utf8_tokenizer.tokenizer import PAD_TOKEN
+from utf8_tokenizer.tokenizer import EOS_TOKEN_ID, PAD_TOKEN, PAD_TOKEN_ID
 
+from welt_training import data_utils
 from welt_training.data_utils import TextDataConfig, build_iterators, pack_words, train_texts
 
 
@@ -68,11 +70,82 @@ def test_build_iterators_batches_examples(tmp_path):
     def make_examples(texts):
         yield from ({"text": text} for text in texts)
 
-    train, validation = build_iterators(config, make_examples, list, rank=0, world_size=1)
+    def collate(examples):
+        return {"texts": [example["text"] for example in examples]}
+
+    train, validation = build_iterators(config, make_examples, collate, rank=0, world_size=1)
     batches = list(islice(train, 4))
-    assert all(len(batch) == 2 for batch in batches)
-    assert {e["text"] for batch in batches for e in batch} <= {f"t{i}" for i in range(3, 10)}
-    # Each evaluation sees the validation examples (repeated to fill its micro batches), the same ones every time
+    assert all(len(batch["texts"]) == 2 for batch in batches)
+    assert {text for batch in batches for text in batch["texts"]} <= {f"t{i}" for i in range(3, 10)}
+    # Each evaluation sees every validation example once, filled up with copies that count for nothing (score 0),
+    # the same ones every time
     evaluations = [[next(validation) for _ in range(3)] for _ in range(2)]
-    assert evaluations[0] == evaluations[1]
-    assert evaluations[0][:2] == [[{"text": "t0"}, {"text": "t1"}], [{"text": "t2"}]]
+    assert [[(batch["texts"], batch["score"].item()) for batch in evaluation] for evaluation in evaluations] == [
+        [(["t0", "t1"], 1), (["t2"], 1), (["t0", "t1"], 0)]] * 2
+
+
+def test_build_iterators_refuse_to_evaluate_a_subset(tmp_path):
+    config = json_config(tmp_path, {"text": [f"t{i}" for i in range(10)]}, max_eval_samples=5)
+    config.num_workers, config.pin_memory, config.persistent_workers = 0, False, False
+    config.micro_batch_size, config.eval_micro_batches = 2, 2
+    with pytest.raises(ValueError, match="increase validation.eval_iters"):
+        build_iterators(config, lambda texts: ({"text": text} for text in texts),
+                        lambda examples: {"n": len(examples)}, rank=0, world_size=1)
+
+
+def test_texts_drop_nul_bytes_and_empty_texts(tmp_path):
+    config = json_config(tmp_path, {"text": ["a\x00b", "", "\x00", "c"]}, max_eval_samples=4)
+    assert [e["text"] for e in config.texts("validation")] == ["ab", "c"]
+    config = json_config(tmp_path, {"text": ["a\x00b", "", "\x00", "c"]}, max_eval_samples=0)
+    assert sorted(islice(train_texts(config, rank=0, world_size=1), 2)) == ["ab", "c"]
+
+
+def test_train_texts_are_split_across_dataloader_workers(tmp_path, monkeypatch):
+    """Each dataloader worker keeps every num_workers-th text of its rank: together, one epoch."""
+    config = json_config(tmp_path, {"text": [f"t{i}" for i in range(10)]}, max_eval_samples=0)
+    workers = []
+    for worker_id in range(2):
+        info = type("WorkerInfo", (), {"id": worker_id, "num_workers": 2})
+        monkeypatch.setattr(torch.utils.data, "get_worker_info", lambda info=info: info)
+        workers.append(list(islice(train_texts(config, rank=0, world_size=1), 5)))
+    assert not set(workers[0]) & set(workers[1])
+    assert sorted(workers[0] + workers[1]) == sorted(f"t{i}" for i in range(10))
+
+
+def test_baseline_counts_the_bytes_of_its_texts():
+    """Each token counts the text since the previous one: every byte once, also with characters split across tokens."""
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m")
+    texts = ["Hello wörld 👋 שלום 日本語 🧑‍🤝‍🧑", "a"]
+    (ids, _, nbytes), = data_utils.token_chunks(texts, tokenizer, length=64)
+    assert sum(nbytes) == sum(len(text.encode()) for text in texts)
+    assert [n for i, n in zip(ids, nbytes, strict=True) if i == tokenizer.eos_token_id] == [0] * ids.count(
+        tokenizer.eos_token_id)
+
+
+@pytest.mark.parametrize("seq_length", [4, 5, 6, 8])
+def test_baseline_scores_the_bytes_welt_scores(seq_length):
+    """Every token is a label once (chunks overlap by one token); shift blocks are not scored, as in WeLT."""
+    from transformers import AutoTokenizer
+
+    from welt.processor import TextImageProcessor, next_word_labels
+
+    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m")
+    texts = ["<en>\x0eHello world\x0f<he> שלום עולם", "no shift block here", "<a>\x0eb\x0f c"]
+    examples = list(data_utils.token_examples(texts, tokenizer, length=4))
+    labels = [token for example in examples for token in example["labels"][example["loss_mask"].bool()].tolist()]
+    # The labels: each document (after the first, after its EOS) without the text of its shift block
+    assert tokenizer.decode(labels) == tokenizer.eos_token.join(data_utils.SHIFT_BLOCK.sub("\x0e", t) for t in texts)
+    baseline_bytes = sum(int(example["label_bytes"].sum()) for example in examples)
+
+    processor = TextImageProcessor.create(max_word_length=32, render_images=False)
+    welt_bytes = 0
+    # Packed in short examples: documents are split across examples, also within shift blocks, as long documents are
+    # in training
+    for packed in pack_words(map(processor.pretokenize, texts), seq_length=seq_length):
+        example = processor.process_single_example(*packed)
+        inputs = (example[k][None] for k in ("input_ids", "sequence_ids", "label_mask"))
+        labels = next_word_labels(*inputs)
+        welt_bytes += int(((labels != PAD_TOKEN_ID) & (labels != EOS_TOKEN_ID)).sum())
+    assert baseline_bytes == welt_bytes == sum(len(data_utils.SHIFT_BLOCK.sub("\x0e", t).encode()) for t in texts)

tests/test_inference.py

@@ -0,0 +1,44 @@
+import pytest
+import torch
+
+pytest.importorskip("vllm")
+
+from welt import ContextTooLongError  # noqa: E402
+from welt.inference import WeLTGenerator  # noqa: E402
+from welt.processor import TextImageProcessor  # noqa: E402
+
+
+def stub_generator(next_words, max_model_len=100):
+    """A WeLTGenerator without engines: _next_words(prefixes) returns each active text's next word (bytes)."""
+    generator = object.__new__(WeLTGenerator)
+    generator.processor = TextImageProcessor.create(max_word_length=32, render_images=False)
+    model_config = type("ModelConfig", (), {"max_model_len": max_model_len})
+    generator.latent = type("LLM", (), {"llm_engine": type("Engine", (), {"model_config": model_config})})
+    generator._latents = lambda sequences: torch.zeros(len(sequences), 1)
+    calls = []
+
+    def _next_words(latents, prefixes, sampling):
+        calls.append((list(prefixes), [params.seed for params in sampling]))
+        return [next_words(prefix) for prefix in prefixes]
+
+    generator._next_words = _next_words
+    return generator, calls
+
+
+def test_generation_continues_a_partial_word_and_stops_on_an_empty_word():
+    words = iter([b"world ", b""])
+    generator, calls = stub_generator(lambda prefix: prefix + b"ld " if prefix else next(words))
+    assert generator.generate_words(["hello wor"]) == [["ld ", "world "]]
+    assert [prefixes for prefixes, _ in calls] == [[b"wor"], [b""], [b""]]
+
+
+def test_generation_stops_at_max_generated_words_with_a_seed_per_text():
+    generator, calls = stub_generator(lambda prefix: b"x ")
+    assert generator.generate(["a ", "b "], max_generated_words=2, seed=10) == ["x x ", "x x "]
+    assert [seeds for _, seeds in calls] == [[10, 11], [12, 13]]  # Distinct across texts and steps
+
+
+def test_generation_refuses_prompts_beyond_the_context():
+    generator, _ = stub_generator(lambda prefix: b"x ", max_model_len=4)
+    with pytest.raises(ContextTooLongError):
+        generator.generate(["a b c"], max_generated_words=2)

tests/test_metrics.py

@@ -25,46 +25,6 @@ def test_welt_loss_func_metrics():
     assert report["word accuracy"].tolist() == [2, 3]
 
 
-def test_baseline_counts_the_bytes_of_its_texts():
-    """Each token counts the text since the previous one: every byte once, also with characters split across tokens."""
-    from transformers import AutoTokenizer
-
-    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m")
-    texts = ["Hello wörld 👋 שלום 日本語 🧑‍🤝‍🧑", "a"]
-    (ids, _, nbytes), = baseline.token_chunks(texts, tokenizer, length=64)
-    assert sum(nbytes) == sum(len(text.encode()) for text in texts)
-    assert [n for i, n in zip(ids, nbytes, strict=True) if i == tokenizer.eos_token_id] == [0] * ids.count(
-        tokenizer.eos_token_id)
-
-
-def test_baseline_scores_the_bytes_welt_scores():
-    """Every token is a label once (chunks overlap by one token); shift blocks are not scored, as in WeLT."""
-    from transformers import AutoTokenizer
-
-    from welt.processor import TextImageProcessor, next_word_labels
-
-    tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m")
-    texts = ["<en>\x0eHello world\x0f<he> שלום עולם", "no shift block here", "<a>\x0eb\x0f c"]
-    examples = list(baseline.token_examples(texts, tokenizer, length=4))
-    labels = [token for example in examples for token in example["labels"][example["loss_mask"].bool()].tolist()]
-    # The labels: each document (after the first, after its EOS) without the text of its shift block
-    assert tokenizer.decode(labels) == tokenizer.eos_token.join(baseline.SHIFT_BLOCK.sub("\x0e", t) for t in texts)
-    baseline_bytes = sum(int(example["label_bytes"].sum()) for example in examples)
-
-    from welt_training.data_utils import pack_words
-
-    processor = TextImageProcessor.create(max_word_length=32, render_images=False)
-    welt_bytes = 0
-    # Packed in examples of 6 words: documents are split across examples (not within a shift block), as long
-    # documents are in training
-    for packed in pack_words(map(processor.pretokenize, texts), seq_length=6):
-        example = processor.process_single_example(*packed)
-        inputs = (example[k][None] for k in ("input_ids", "sequence_ids", "label_mask"))
-        labels = next_word_labels(*inputs)
-        welt_bytes += int(((labels != PAD_TOKEN_ID) & (labels != EOS_TOKEN_ID)).sum())
-    assert baseline_bytes == welt_bytes == sum(len(baseline.SHIFT_BLOCK.sub("\x0e", t).encode()) for t in texts)
-
-
 def test_baseline_loss_func_bits_per_byte():
     losses = torch.tensor([[2.0, 3.0, 1.0, 1.5, 9.0]])
     loss_mask = torch.tensor([[1, 1, 1, 1, 0]])  # The last label is padding