From 81bf0e628cf8ffc0fde15a4236cc235287e11e51 Mon Sep 17 00:00:00 2001 From: w00958190 Date: Wed, 12 Aug 2026 16:46:02 +0800 Subject: [PATCH 1/3] feat(scripts): build a draft vocabulary mapping without regenerating features Choosing a draft vocabulary currently requires a full hidden-state capture first: process_token_dict_to_mappings runs inside prepare_hidden_states, so trying a second draft_vocab_size, or building a mapping for a corpus whose features were captured before the mapping was needed, means paying for the capture again. On a 600k-sample corpus that is hours of GPU time to answer a question about token frequencies. scripts/build_vocab_mapping.py derives the mapping from either source, with no GPU involved: - --data-path counts loss-bearing tokens from the source JSONL by applying the same chat template and truncation the trainer would; - --hidden-states-path counts them from already prepared features. Both stream: the tally is accumulated per shard instead of materializing a token column for the whole corpus, and the per-corpus counts are cached under a key that includes the corpus identity, so sweeping several draft_vocab_size values re-reads nothing. Two supporting changes: - specforge/training/vocab_mapping.py moves to specforge/data/. Scripts may not import specforge.training (tests/test_runtime/ test_package_architecture.py enforces it), and counting tokens in prepared features is a data concern that the trainer merely happens to be the first caller of. The module is unchanged. - process_token_dict_to_mappings builds t2d by scattering into a zeroed mask instead of testing membership in a Python list per target id. The old form is quadratic in the vocabulary: tens of seconds at V=248320 and K=64000, inside a phase that prints nothing, for a result that takes milliseconds. Co-Authored-By: Claude Opus 5 --- scripts/build_vocab_mapping.py | 448 ++++++++++++++++++ specforge/data/preprocessing.py | 17 +- specforge/{training => data}/vocab_mapping.py | 7 +- specforge/training/assembly.py | 2 +- .../test_vocab_mapping_construction.py | 81 ++++ .../test_offline_vocab_mapping.py | 2 +- .../test_scripts/test_build_vocab_mapping.py | 304 ++++++++++++ 7 files changed, 851 insertions(+), 10 deletions(-) create mode 100644 scripts/build_vocab_mapping.py rename specforge/{training => data}/vocab_mapping.py (90%) create mode 100644 tests/test_data/test_vocab_mapping_construction.py create mode 100644 tests/test_scripts/test_build_vocab_mapping.py diff --git a/scripts/build_vocab_mapping.py b/scripts/build_vocab_mapping.py new file mode 100644 index 000000000..7216e3489 --- /dev/null +++ b/scripts/build_vocab_mapping.py @@ -0,0 +1,448 @@ +"""Build a draft vocabulary mapping from prepared offline features. + +Training can derive this map on its own, but only for a colocated offline run, +and only by reading every feature file first -- which for gzipped features means +decompressing the whole dataset before the first step. This script does that +pass once, caches the token counts, and then answers any number of +``draft_vocab_size`` questions from the cache in milliseconds. + +That separation is the point: counting is the expensive half and depends only on +the dataset, while choosing the top-K is cheap and is the half you actually want +to iterate on. Changing K therefore never requires regenerating hidden states, +and never requires a second pass over them. + +The map is emitted at the *draft config's* ``vocab_size``, which is the length +the model's ``t2d`` buffer is registered with. Sizing it from the target config +instead would produce a file that silently fails to load whenever the target +declares ``padded_vocab_size``. + +Two sources, same numbers. ``--hidden-states-path`` reads the prepared +features, which is exact but serial -- for a large gzipped dataset it is not +merely slow, it is impractical, since every file is decompressed in full to +recover two small tensors. ``--data-path`` re-tokenizes the source JSONL with +the same stack the capture used, in parallel, without touching the features at +all; pass it the same tokenizer, template, max length, and filters. + +Survey several sizes before committing to one (writes nothing): + + python scripts/build_vocab_mapping.py \ + --data-path ./cache/dataset/train.jsonl \ + --tokenizer-path Qwen/Qwen3-8B --chat-template qwen --max-length 4096 \ + --draft-model-config configs/qwen3.6-27b-dspark.json \ + --draft-vocab-size 16000,32000,48000,64000 + +Then write the chosen one, reusing the cached counts: + + python scripts/build_vocab_mapping.py \ + --data-path ./cache/dataset/train.jsonl \ + --tokenizer-path Qwen/Qwen3-8B --chat-template qwen --max-length 4096 \ + --draft-model-config configs/qwen3.6-27b-dspark-draftvocab64k.json \ + --output-path ./cache/vocab_mapping/qwen3.6-27b-k64000.pt +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +from collections import Counter +from pathlib import Path +from typing import Optional + +import torch + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Derive t2d/d2t from prepared offline features, without " + "regenerating them and without a second pass per vocabulary size." + ) + ) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument( + "--hidden-states-path", + type=Path, + help=( + "Directory of prepared offline features (.ckpt / .ckpt.gz). Exact, " + "but reads every file serially -- impractical for a large gzipped " + "dataset; prefer --data-path there." + ), + ) + source.add_argument( + "--data-path", + type=Path, + help=( + "Raw conversation JSONL. Re-tokenizes with the same stack " + "prepare_hidden_states.py uses, in parallel and without touching " + "the features at all. Pass the same tokenizer/template/max-length " + "and --minimum-valid-tokens the capture ran with, or the counts " + "will describe a different corpus than training sees." + ), + ) + parser.add_argument( + "--tokenizer-path", + default=None, + help="Target model/tokenizer path. Required with --data-path.", + ) + parser.add_argument( + "--chat-template", + default=None, + help="Chat template used at capture time. Required with --data-path.", + ) + parser.add_argument( + "--is-preformatted", + action="store_true", + help="Source rows already have the chat template applied.", + ) + parser.add_argument( + "--minimum-valid-tokens", + type=int, + default=None, + help=( + "Mirror prepare_hidden_states.py's filter so dropped samples do " + "not contribute frequencies." + ), + ) + parser.add_argument( + "--num-samples", + type=int, + default=None, + help="Mirror prepare_hidden_states.py's --num-samples.", + ) + parser.add_argument( + "--build-dataset-num-proc", + type=int, + default=8, + help="Tokenization worker processes for --data-path.", + ) + parser.add_argument( + "--dataset-cache-dir", + type=Path, + default=Path("./cache"), + help=( + "Root for the conversation and tokenized-dataset caches " + "(/hf_dataset and /processed_dataset). Point it at a " + "partition with room; matching data.cache_dir lets a later " + "training run reuse the tokenization." + ), + ) + parser.add_argument( + "--draft-model-config", + type=Path, + required=True, + help=( + "Draft config JSON. Supplies vocab_size (the map's length, matching " + "the model's t2d buffer) and the default draft_vocab_size." + ), + ) + parser.add_argument( + "--draft-vocab-size", + default=None, + help=( + "One size, or a comma-separated list to survey. Defaults to the " + "draft config's draft_vocab_size. A list writes nothing." + ), + ) + parser.add_argument( + "--output-path", + type=Path, + default=None, + help=("Where to write the {t2d, d2t} file. Omit to only report coverage."), + ) + parser.add_argument( + "--max-length", + type=int, + default=None, + help="Truncate each sample's ids/mask, matching data.max_length.", + ) + parser.add_argument( + "--counts-cache", + type=Path, + default=None, + help=( + "Token-count cache. Defaults to /.token_counts.pt " + "for --hidden-states-path, and " + "/vocab_mapping/.token_counts.pt for --data-path. " + "Reused only when the corpus fingerprint is unchanged." + ), + ) + parser.add_argument( + "--recount", + action="store_true", + help="Ignore any cached counts and rescan the features.", + ) + return parser + + +def _feature_identity(hidden_states_path: str, max_length: Optional[int]) -> str: + """Fingerprint the feature set so a stale cache is never silently reused.""" + from specforge.runtime.data_plane.offline_reader import list_feature_files + + entries = [] + for path in list_feature_files(hidden_states_path): + stat = os.stat(path) + entries.append((os.path.abspath(path), stat.st_size, stat.st_mtime_ns)) + payload = json.dumps( + {"kind": "offline-features-v1", "files": entries, "max_length": max_length}, + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest() + + +def _dataset_identity(args, vocab_size: int) -> str: + """Fingerprint the tokenization inputs, so a cache answers for its own corpus.""" + stat = os.stat(args.data_path) + payload = json.dumps( + { + "kind": "conversations-jsonl-v1", + "path": os.path.abspath(args.data_path), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + "tokenizer": args.tokenizer_path, + "chat_template": args.chat_template, + "max_length": args.max_length, + "is_preformatted": bool(args.is_preformatted), + "minimum_valid_tokens": args.minimum_valid_tokens, + "num_samples": args.num_samples, + "vocab_size": vocab_size, + }, + sort_keys=True, + ) + return hashlib.sha256(payload.encode()).hexdigest() + + +def count_dataset_tokens(args, *, vocab_size: int) -> Counter: + """Count loss-bearing tokens by re-tokenizing the source conversations. + + Runs the same tokenizer, chat template, truncation, and trainable-token + filter that ``prepare_hidden_states.py`` applied, so the frequencies match + the captured features without reading them. Any of those knobs differing + from the capture yields a mapping for a different corpus, which is why they + are all explicit rather than defaulted. + """ + from datasets import Dataset + + from specforge.data.preprocessing import build_eagle3_dataset + from specforge.utils import load_tokenizer, safe_conversations_generator + + tokenizer = load_tokenizer(args.tokenizer_path) + dataset = Dataset.from_generator( + generator=safe_conversations_generator, + gen_kwargs={"file_path": str(args.data_path)}, + # Pinned under --dataset-cache-dir like prepare_hidden_states.py does. + # Left to its default this lands in ~/.cache/huggingface, which is + # rarely the partition with room for a 600k-conversation corpus. + cache_dir=str(args.dataset_cache_dir / "hf_dataset"), + num_proc=min(args.build_dataset_num_proc, 32), + ) + if args.num_samples is not None: + dataset = dataset.select(range(args.num_samples)) + processed = build_eagle3_dataset( + dataset=dataset, + tokenizer=tokenizer, + chat_template=args.chat_template, + max_length=args.max_length, + cache_dir=str(args.dataset_cache_dir / "processed_dataset"), + cache_key=_dataset_identity(args, vocab_size), + is_preformatted=args.is_preformatted, + num_proc=args.build_dataset_num_proc, + minimum_valid_tokens=args.minimum_valid_tokens, + ) + print(f"Tokenized {len(processed)} samples") + return tally_loss_tokens(processed, vocab_size=vocab_size) + + +def tally_loss_tokens(dataset, *, vocab_size: int, batch_size: int = 512) -> Counter: + """Sum loss-bearing token frequencies over a tokenized dataset. + + Streams batches rather than touching ``dataset["input_ids"]``: column access + materializes every sequence as Python ints at once, which for a corpus this + feature targets is tens of gigabytes and looks like a hang. Frequencies then + accumulate into one dense bincount per batch instead of per token, keeping + the work in tensors rather than in a billion-iteration Python loop. + """ + from tqdm import tqdm + + totals = torch.zeros(vocab_size, dtype=torch.int64) + batches = (len(dataset) + batch_size - 1) // batch_size + for batch in tqdm( + dataset.iter(batch_size=batch_size), + total=batches, + desc="Counting tokens for vocab mapping", + ): + selected = [] + for input_ids, loss_mask in zip(batch["input_ids"], batch["loss_mask"]): + ids = torch.as_tensor(input_ids).reshape(-1) + mask = torch.as_tensor(loss_mask).reshape(-1) + kept = ids[mask.to(dtype=torch.bool)] + if kept.numel(): + selected.append(kept) + if not selected: + continue + flat = torch.cat(selected).long() + if int(flat.min()) < 0 or int(flat.max()) >= vocab_size: + raise ValueError( + f"token id {int(flat.max())} is outside the draft config's " + f"vocab_size {vocab_size}" + ) + totals += torch.bincount(flat, minlength=vocab_size) + + present = torch.nonzero(totals, as_tuple=False).flatten() + return Counter({int(token): int(totals[token]) for token in present.tolist()}) + + +def load_or_count_tokens(args, *, vocab_size: int, counts_cache: Path) -> Counter: + """Return loss-bearing token frequencies, counting at most once per corpus.""" + from_features = args.hidden_states_path is not None + identity = ( + _feature_identity(str(args.hidden_states_path), args.max_length) + if from_features + else _dataset_identity(args, vocab_size) + ) + if not args.recount and counts_cache.exists(): + cached = torch.load(counts_cache, map_location="cpu", weights_only=False) + if cached.get("identity") == identity: + print(f"Reusing token counts from {counts_cache}") + return Counter(cached["counts"]) + print(f"{counts_cache} describes a different corpus; recounting.") + + if from_features: + from specforge.data.vocab_mapping import count_effective_feature_tokens + + print(f"Counting loss-bearing tokens under {args.hidden_states_path} ...") + counts = count_effective_feature_tokens( + str(args.hidden_states_path), + max_length=args.max_length, + target_vocab_size=vocab_size, + ) + else: + print(f"Tokenizing {args.data_path} to count loss-bearing tokens ...") + counts = count_dataset_tokens(args, vocab_size=vocab_size) + + counts_cache.parent.mkdir(parents=True, exist_ok=True) + temporary = counts_cache.with_suffix(f".{os.getpid()}.tmp") + torch.save({"identity": identity, "counts": dict(counts)}, temporary) + os.replace(temporary, counts_cache) + print(f"Cached token counts at {counts_cache}") + return counts + + +def coverage_ratio(counts: Counter, draft_vocab_size: int) -> float: + """Share of loss-bearing token occurrences the top-K tokens account for. + + This is the ceiling on acceptance: a target token outside the draft + vocabulary can never be proposed, so that position is always rejected. + """ + total = sum(counts.values()) + if total == 0: + return 0.0 + kept = sum(frequency for _, frequency in counts.most_common(draft_vocab_size)) + return kept / total + + +def write_mapping( + counts: Counter, + *, + draft_vocab_size: int, + vocab_size: int, + output_path: Path, +) -> None: + from specforge.core.compact_teacher import validate_vocab_mapping_consistency + from specforge.data.preprocessing import process_token_dict_to_mappings + + d2t, t2d = process_token_dict_to_mappings( + Counter(counts), draft_vocab_size, vocab_size + ) + # The same invariant the model checks on install; catching it here keeps a + # broken file from ever reaching a training run. + validate_vocab_mapping_consistency(t2d, d2t) + output_path.parent.mkdir(parents=True, exist_ok=True) + temporary = output_path.with_suffix(f".{os.getpid()}.tmp") + torch.save({"d2t": d2t, "t2d": t2d}, temporary) + os.replace(temporary, output_path) + + +def main(argv=None) -> int: + args = build_parser().parse_args(argv) + + with open(args.draft_model_config, encoding="utf-8") as handle: + draft_config = json.load(handle) + vocab_size = int(draft_config["vocab_size"]) + + if args.draft_vocab_size is not None: + sizes = [int(item) for item in str(args.draft_vocab_size).split(",") if item] + else: + configured = draft_config.get("draft_vocab_size") + if configured is None: + raise ValueError( + f"{args.draft_model_config} has no draft_vocab_size; pass " + "--draft-vocab-size explicitly" + ) + sizes = [int(configured)] + for size in sizes: + if not 0 < size <= vocab_size: + raise ValueError( + f"draft_vocab_size must be in (0, {vocab_size}], got {size}" + ) + if len(sizes) > 1 and args.output_path is not None: + raise ValueError( + "--output-path writes a single mapping; pass one --draft-vocab-size" + ) + + if args.data_path is not None: + missing = [ + name + for name, value in ( + ("--tokenizer-path", args.tokenizer_path), + ("--chat-template", args.chat_template), + ("--max-length", args.max_length), + ) + if value is None + ] + if missing: + raise ValueError( + f"--data-path re-tokenizes the corpus and must reproduce the " + f"capture exactly; missing {missing}" + ) + + if args.counts_cache is not None: + counts_cache = args.counts_cache + elif args.hidden_states_path is not None: + counts_cache = args.hidden_states_path / ".token_counts.pt" + else: + counts_cache = args.dataset_cache_dir / "vocab_mapping" / ".token_counts.pt" + counts = load_or_count_tokens( + args, vocab_size=vocab_size, counts_cache=counts_cache + ) + distinct = len(counts) + print(f"Distinct loss-bearing tokens: {distinct} of {vocab_size}") + + for size in sizes: + ratio = coverage_ratio(counts, size) + note = "" + if size > distinct: + note = f" (only {distinct} tokens ever appear; the rest are padding)" + print(f" top {size:>7} token frequency ratio: {ratio:7.2%}{note}") + + if args.output_path is None: + print( + "\nNo --output-path given, so nothing was written. The ratio above " + "is the acceptance ceiling for that size." + ) + return 0 + + write_mapping( + counts, + draft_vocab_size=sizes[0], + vocab_size=vocab_size, + output_path=args.output_path, + ) + print(f"\nWrote mapping to {args.output_path}") + print("Point model.vocab_mapping_path at it to skip the training-time scan.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/specforge/data/preprocessing.py b/specforge/data/preprocessing.py index a2f268101..8b1fea8f1 100644 --- a/specforge/data/preprocessing.py +++ b/specforge/data/preprocessing.py @@ -757,12 +757,15 @@ def process_token_dict_to_mappings( top_N_ratio = top_N_frequency_sum / total_frequency print(f"top {draft_vocab_size} token frequency ratio: {top_N_ratio:.2%}") - used_tokens = [key for key, freq in top_N] - used_tokens.sort() - - d2t = [used_tokens[i] - i for i in range(len(used_tokens))] - t2d = [i in used_tokens for i in range(target_vocab_size)] - d2t = torch.tensor(d2t) - t2d = torch.tensor(t2d) + used_tokens = torch.tensor( + sorted(key for key, _frequency in top_N), dtype=torch.int64 + ) + # ``t2d`` used to be built as ``[i in used_tokens for i in range(V)]`` over a + # list, i.e. a linear scan per target id -- tens of seconds at V=248320 and + # K=64000, in a phase that prints nothing. Scattering into a zeroed mask is + # the same result in milliseconds. + d2t = used_tokens - torch.arange(used_tokens.numel(), dtype=torch.int64) + t2d = torch.zeros(target_vocab_size, dtype=torch.bool) + t2d[used_tokens] = True return d2t, t2d diff --git a/specforge/training/vocab_mapping.py b/specforge/data/vocab_mapping.py similarity index 90% rename from specforge/training/vocab_mapping.py rename to specforge/data/vocab_mapping.py index 63dac6063..4e6c21ed7 100644 --- a/specforge/training/vocab_mapping.py +++ b/specforge/data/vocab_mapping.py @@ -6,7 +6,12 @@ # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 -"""Dataset vocabulary mappings shared by online and offline EAGLE training.""" +"""Token frequencies for draft vocabulary mappings, read from offline features. + +Lives beside preprocessing's top-K selection rather than under training: it +reads prepared feature files and is equally useful to scripts, which are not +allowed to reach into specforge.training. +""" from __future__ import annotations diff --git a/specforge/training/assembly.py b/specforge/training/assembly.py index 5d773ccb9..6980117d2 100644 --- a/specforge/training/assembly.py +++ b/specforge/training/assembly.py @@ -479,8 +479,8 @@ def _ensure_offline_vocab_mapping( ): return + from specforge.data.vocab_mapping import count_effective_feature_tokens from specforge.runtime.data_plane.offline_reader import list_feature_files - from specforge.training.vocab_mapping import count_effective_feature_tokens identity_parts = [] for path in list_feature_files(cfg.data.hidden_states_path): diff --git a/tests/test_data/test_vocab_mapping_construction.py b/tests/test_data/test_vocab_mapping_construction.py new file mode 100644 index 000000000..675469279 --- /dev/null +++ b/tests/test_data/test_vocab_mapping_construction.py @@ -0,0 +1,81 @@ +# coding=utf-8 +"""t2d/d2t construction from token frequencies. + +The mapping is built once per run but consumed by every training step and by +serving, so this pins it against the formulation it replaced rather than +against hand-written expectations: the two must agree bit for bit, including +dtypes, or a rebuilt mapping would silently reorder a checkpoint's head rows. +""" + +import random +import unittest +from collections import Counter + +import torch + +from specforge.data.preprocessing import process_token_dict_to_mappings + + +def _reference_mappings(top_n, target_vocab_size): + """The pre-vectorization construction, transcribed and frozen. + + Kept as literal Python so it cannot drift with the implementation under + test; its ``i in used_tokens`` list scan is exactly the cost that made the + real thing worth replacing. + """ + used_tokens = [key for key, _frequency in top_n] + used_tokens.sort() + d2t = [used_tokens[i] - i for i in range(len(used_tokens))] + t2d = [i in used_tokens for i in range(target_vocab_size)] + return torch.tensor(d2t), torch.tensor(t2d) + + +class VocabMappingConstructionTest(unittest.TestCase): + def test_matches_the_reference_construction(self): + random.seed(0) + for vocab_size, draft_vocab_size in ((256, 64), (5000, 1500), (20000, 7000)): + with self.subTest(vocab_size=vocab_size, draft=draft_vocab_size): + distinct = min(vocab_size, int(draft_vocab_size * 1.7)) + counts = Counter( + { + token: random.randint(1, 1000) + for token in random.sample(range(vocab_size), distinct) + } + ) + expected = _reference_mappings( + counts.most_common(draft_vocab_size), vocab_size + ) + actual = process_token_dict_to_mappings( + Counter(counts), draft_vocab_size, vocab_size + ) + + for name, got, want in zip(("d2t", "t2d"), actual, expected): + self.assertEqual(want.dtype, got.dtype, name) + self.assertTrue(torch.equal(want, got), name) + + def test_satisfies_the_invariant_the_model_checks_on_install(self): + """nonzero(t2d) == d2t + arange, which d2t being an offset table needs.""" + from specforge.core.compact_teacher import validate_vocab_mapping_consistency + + counts = Counter({token: token + 1 for token in range(0, 200, 3)}) + d2t, t2d = process_token_dict_to_mappings(counts, 32, 256) + + validate_vocab_mapping_consistency(t2d, d2t) + self.assertEqual(32, int(t2d.sum())) + selected = torch.nonzero(t2d, as_tuple=False).flatten() + self.assertTrue(torch.equal(selected, d2t + torch.arange(32))) + + def test_pads_when_the_corpus_has_too_few_distinct_tokens(self): + """A short corpus must still yield exactly draft_vocab_size entries.""" + counts = Counter({1: 5, 7: 3}) + d2t, t2d = process_token_dict_to_mappings(counts, 8, 64) + + self.assertEqual((8,), tuple(d2t.shape)) + self.assertEqual((64,), tuple(t2d.shape)) + self.assertEqual(8, int(t2d.sum())) + # The observed tokens survive the padding. + self.assertTrue(bool(t2d[1]) and bool(t2d[7])) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runtime/test_offline_vocab_mapping.py b/tests/test_runtime/test_offline_vocab_mapping.py index 3985f45e3..dd5d49ca5 100644 --- a/tests/test_runtime/test_offline_vocab_mapping.py +++ b/tests/test_runtime/test_offline_vocab_mapping.py @@ -11,12 +11,12 @@ from specforge.algorithms.builtin import builtin_algorithm_registry from specforge.config import Config +from specforge.data.vocab_mapping import count_effective_feature_tokens from specforge.training.assembly import ( _ensure_offline_vocab_mapping, _install_dataset_vocab_mapping, _prompt_cache_key, ) -from specforge.training.vocab_mapping import count_effective_feature_tokens ALGORITHM = builtin_algorithm_registry().resolve("eagle3") diff --git a/tests/test_scripts/test_build_vocab_mapping.py b/tests/test_scripts/test_build_vocab_mapping.py new file mode 100644 index 000000000..4e1575679 --- /dev/null +++ b/tests/test_scripts/test_build_vocab_mapping.py @@ -0,0 +1,304 @@ +# coding=utf-8 +"""Standalone vocabulary-mapping builder. + +Exercised against real feature files (both plain and gzipped) and a real draft +model, because the two things most likely to break silently -- the map's length +matching the model's ``t2d`` buffer, and the count cache being reused for the +wrong dataset -- are invisible to argument-level checks. +""" + +import gzip +import json +import tempfile +import unittest +from collections import Counter +from pathlib import Path + +import torch + +from scripts.build_vocab_mapping import coverage_ratio, main + +VOCAB_SIZE = 256 +DRAFT_VOCAB_SIZE = 64 + + +def _write_features(directory: Path, *, compress: bool, seed: int = 0) -> None: + """Write two feature files whose loss-bearing ids are a known skewed set.""" + torch.manual_seed(seed) + for index in range(2): + # Even ids appear often, odd ids once, so top-K is a scattered set and + # d2t comes out as a non-trivial offset table rather than all zeros. + ids = torch.cat( + [ + torch.arange(0, VOCAB_SIZE, 2).repeat(3), + torch.arange(1, VOCAB_SIZE, 2), + ] + ) + record = { + "input_ids": ids, + "loss_mask": torch.ones_like(ids), + "hidden_states": torch.zeros(ids.numel(), 4), + "target_last_hidden_states": torch.zeros(ids.numel(), 4), + } + path = directory / f"sample_{index}.ckpt{'.gz' if compress else ''}" + if compress: + with gzip.open(path, "wb") as handle: + torch.save(record, handle) + else: + torch.save(record, path) + + +def _write_draft_config(path: Path, draft_vocab_size=None) -> None: + payload = {"vocab_size": VOCAB_SIZE} + if draft_vocab_size is not None: + payload["draft_vocab_size"] = draft_vocab_size + path.write_text(json.dumps(payload), encoding="utf-8") + + +class BuildVocabMappingTest(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.root = Path(self._tmp.name) + self.features = self.root / "features" + self.features.mkdir() + self.config = self.root / "draft.json" + + def tearDown(self): + self._tmp.cleanup() + + def _run(self, *extra, compress=False, draft_vocab_size=DRAFT_VOCAB_SIZE): + if not any(self.features.iterdir()): + _write_features(self.features, compress=compress) + _write_draft_config(self.config, draft_vocab_size) + return main( + [ + "--hidden-states-path", + str(self.features), + "--draft-model-config", + str(self.config), + *extra, + ] + ) + + def test_written_mapping_loads_into_a_real_draft_model(self): + """The whole point: the file must fit the model's buffers, not just parse.""" + from transformers import LlamaConfig + + from specforge.modeling.draft.llama3_eagle import LlamaForCausalLMEagle3 + + out = self.root / "mapping.pt" + self.assertEqual(0, self._run("--output-path", str(out))) + + config = LlamaConfig( + vocab_size=VOCAB_SIZE, + draft_vocab_size=DRAFT_VOCAB_SIZE, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=128, + ) + model = LlamaForCausalLMEagle3(config) + model.load_vocab_mapping(str(out)) + + self.assertTrue(model.vocab_mapping_loaded) + self.assertEqual(int(model.t2d.sum()), DRAFT_VOCAB_SIZE) + # The frequent (even) ids must be the ones kept. + kept = torch.nonzero(model.t2d, as_tuple=False).flatten().tolist() + self.assertTrue(all(token % 2 == 0 for token in kept)) + + def test_gzipped_features_are_supported(self): + out = self.root / "mapping.pt" + self.assertEqual(0, self._run("--output-path", str(out), compress=True)) + mapping = torch.load(out, map_location="cpu") + self.assertEqual(tuple(mapping["t2d"].shape), (VOCAB_SIZE,)) + self.assertEqual(tuple(mapping["d2t"].shape), (DRAFT_VOCAB_SIZE,)) + + def test_counts_are_cached_and_reused(self): + """Counting is the expensive half; changing K must not repeat it.""" + cache = self.features / ".token_counts.pt" + self._run("--output-path", str(self.root / "a.pt")) + self.assertTrue(cache.exists()) + + stamp = cache.stat().st_mtime_ns + self._run("--draft-vocab-size", "32", "--output-path", str(self.root / "b.pt")) + self.assertEqual(stamp, cache.stat().st_mtime_ns, "counts were recomputed") + + def test_stale_cache_is_not_reused_for_different_features(self): + """A cache keyed only by path would silently answer for the wrong data.""" + cache = self.features / ".token_counts.pt" + self._run("--output-path", str(self.root / "a.pt")) + before = torch.load(cache, map_location="cpu", weights_only=False)["identity"] + + _write_features(self.features, compress=False, seed=1) + (self.features / "extra.ckpt").write_bytes( + (self.features / "sample_0.ckpt").read_bytes() + ) + self._run("--output-path", str(self.root / "b.pt")) + after = torch.load(cache, map_location="cpu", weights_only=False)["identity"] + self.assertNotEqual(before, after) + + def test_survey_reports_several_sizes_and_writes_nothing(self): + self.assertEqual(0, self._run("--draft-vocab-size", "16,32,64")) + self.assertEqual([], sorted(self.root.glob("*.pt"))) + + def test_survey_refuses_to_write_a_single_output(self): + with self.assertRaisesRegex(ValueError, "single mapping"): + self._run( + "--draft-vocab-size", "16,32", "--output-path", str(self.root / "x.pt") + ) + + def test_config_without_draft_vocab_size_requires_the_flag(self): + with self.assertRaisesRegex(ValueError, "no draft_vocab_size"): + self._run(draft_vocab_size=None) + + def test_size_outside_the_vocabulary_is_rejected(self): + for bad in ("0", "-1", str(VOCAB_SIZE + 1)): + with self.subTest(draft_vocab_size=bad): + with self.assertRaisesRegex(ValueError, "draft_vocab_size must be"): + self._run("--draft-vocab-size", bad) + + def test_data_path_requires_the_capture_settings(self): + """A JSONL source that guesses the tokenization describes another corpus.""" + _write_draft_config(self.config, DRAFT_VOCAB_SIZE) + jsonl = self.root / "train.jsonl" + jsonl.write_text("{}\n", encoding="utf-8") + with self.assertRaisesRegex(ValueError, "must reproduce the capture"): + main( + [ + "--data-path", + str(jsonl), + "--draft-model-config", + str(self.config), + ] + ) + + def test_the_two_sources_are_mutually_exclusive(self): + _write_draft_config(self.config, DRAFT_VOCAB_SIZE) + with self.assertRaises(SystemExit): + main( + [ + "--hidden-states-path", + str(self.features), + "--data-path", + str(self.root / "train.jsonl"), + "--draft-model-config", + str(self.config), + ] + ) + + def test_a_source_is_required(self): + _write_draft_config(self.config, DRAFT_VOCAB_SIZE) + with self.assertRaises(SystemExit): + main(["--draft-model-config", str(self.config)]) + + def test_dataset_counts_are_cached_under_the_dataset_cache_dir(self): + """The JSONL route caches too, so surveying K stays a one-time cost.""" + from types import SimpleNamespace + + from scripts.build_vocab_mapping import load_or_count_tokens + + cache = self.root / "counts.pt" + args = SimpleNamespace( + hidden_states_path=None, + data_path=self.root / "train.jsonl", + tokenizer_path="tok", + chat_template="qwen", + max_length=128, + is_preformatted=False, + minimum_valid_tokens=None, + num_samples=None, + build_dataset_num_proc=1, + dataset_cache_dir=self.root, + recount=False, + ) + args.data_path.write_text("{}\n", encoding="utf-8") + + calls = [] + + def fake_count(_args, *, vocab_size): + calls.append(vocab_size) + return Counter({0: 5, 2: 3}) + + import scripts.build_vocab_mapping as module + + original = module.count_dataset_tokens + module.count_dataset_tokens = fake_count + try: + first = load_or_count_tokens( + args, vocab_size=VOCAB_SIZE, counts_cache=cache + ) + second = load_or_count_tokens( + args, vocab_size=VOCAB_SIZE, counts_cache=cache + ) + finally: + module.count_dataset_tokens = original + + self.assertEqual(1, len(calls), "the corpus was tokenized twice") + self.assertEqual(first, second) + + def test_tally_counts_only_loss_bearing_tokens(self): + from datasets import Dataset + + from scripts.build_vocab_mapping import tally_loss_tokens + + dataset = Dataset.from_dict( + { + "input_ids": [[5, 7, 5, 9], [5, 200], [11, 11]], + "loss_mask": [[1, 1, 1, 0], [0, 1], [0, 0]], + } + ) + counts = tally_loss_tokens(dataset, vocab_size=VOCAB_SIZE, batch_size=2) + + # 9 is masked out, 11 is entirely masked, 5 appears twice in row 0. + self.assertEqual(Counter({5: 2, 7: 1, 200: 1}), counts) + + def test_tally_never_materializes_a_whole_column(self): + """Column access loads every sequence as Python ints; at scale that hangs.""" + from datasets import Dataset + + from scripts.build_vocab_mapping import tally_loss_tokens + + dataset = Dataset.from_dict( + { + "input_ids": [[1, 2], [3, 4]], + "loss_mask": [[1, 1], [1, 1]], + } + ) + original = Dataset.__getitem__ + column_reads = [] + + def tracking_getitem(self, key): + if isinstance(key, str): + column_reads.append(key) + return original(self, key) + + Dataset.__getitem__ = tracking_getitem + try: + counts = tally_loss_tokens(dataset, vocab_size=VOCAB_SIZE, batch_size=1) + finally: + Dataset.__getitem__ = original + + self.assertEqual([], column_reads, f"read whole columns: {column_reads}") + self.assertEqual(Counter({1: 1, 2: 1, 3: 1, 4: 1}), counts) + + def test_tally_rejects_ids_outside_the_vocabulary(self): + from datasets import Dataset + + from scripts.build_vocab_mapping import tally_loss_tokens + + dataset = Dataset.from_dict({"input_ids": [[VOCAB_SIZE]], "loss_mask": [[1]]}) + with self.assertRaisesRegex(ValueError, "outside the draft config"): + tally_loss_tokens(dataset, vocab_size=VOCAB_SIZE) + + def test_coverage_ratio_matches_the_definition(self): + counts = Counter({0: 90, 1: 9, 2: 1}) + self.assertAlmostEqual(0.9, coverage_ratio(counts, 1)) + self.assertAlmostEqual(0.99, coverage_ratio(counts, 2)) + self.assertAlmostEqual(1.0, coverage_ratio(counts, 3)) + self.assertAlmostEqual(0.0, coverage_ratio(Counter(), 5)) + + +if __name__ == "__main__": + unittest.main() From 6eef8987ea0385d157233e4212d6f517a4ec32fe Mon Sep 17 00:00:00 2001 From: w00958190 Date: Wed, 12 Aug 2026 17:06:52 +0800 Subject: [PATCH 2/3] docs: building a draft vocabulary mapping from the corpus Co-Authored-By: Claude Opus 5 --- docs/basic_usage/data_preparation.md | 57 ++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/docs/basic_usage/data_preparation.md b/docs/basic_usage/data_preparation.md index 3537309d1..9d70a4457 100644 --- a/docs/basic_usage/data_preparation.md +++ b/docs/basic_usage/data_preparation.md @@ -341,6 +341,63 @@ See the [Training](training.md) guide for the complete run schema and supported combinations. +## 🔤 Build a draft vocabulary mapping + +A draft that predicts over a subset of the target vocabulary needs a `t2d`/`d2t` +mapping. `prepare_hidden_states.py` writes one as a side effect of capture, so +trying a second `draft_vocab_size` — or building a mapping for a corpus whose +features were captured before one was needed — otherwise means repeating the +capture. `scripts/build_vocab_mapping.py` derives it directly, without a GPU. + +From the source JSONL, re-tokenizing with the same stack the capture used. Pass +the same tokenizer, chat template, `--max-length`, and `--minimum-valid-tokens`, +or the counts describe a different corpus than training sees: + +```bash +python scripts/build_vocab_mapping.py \ + --data-path ./cache/dataset/sharegpt_train.jsonl \ + --tokenizer-path Qwen/Qwen3-8B \ + --chat-template qwen \ + --max-length 2048 \ + --dataset-cache-dir ./cache \ + --draft-model-config configs/qwen3-8b-eagle3.json \ + --output-path ./cache/vocab_mapping/qwen3-8b-32k.pt +``` + +Or from features that already exist, which needs neither the tokenizer nor the +chat template: + +```bash +python scripts/build_vocab_mapping.py \ + --hidden-states-path ./cache/hidden_states/sharegpt_train_Qwen3-8B \ + --max-length 2048 \ + --draft-model-config configs/qwen3-8b-eagle3.json \ + --output-path ./cache/vocab_mapping/qwen3-8b-32k.pt +``` + +The draft config supplies `vocab_size` — the map's length, which must match the +model's `t2d` buffer — and the default `draft_vocab_size`. Both modes cache the +per-corpus token counts, reusing them only while the corpus fingerprint is +unchanged, so surveying sizes costs one pass: + +```bash +python scripts/build_vocab_mapping.py \ + --hidden-states-path ./cache/hidden_states/sharegpt_train_Qwen3-8B \ + --draft-model-config configs/qwen3-8b-eagle3.json \ + --draft-vocab-size 16000,32000,64000 +``` + +A comma-separated list only reports the coverage each size would reach and +writes nothing; omitting `--output-path` reports coverage for a single size. + +Point `model.vocab_mapping_path` at the written file. A disaggregated run +requires this, because producer and consumer cannot derive one shared mapping +independently. + +`--hidden-states-path` is exact but reads every feature file serially, which is +impractical for a large gzipped corpus; prefer `--data-path` there. + + ## ➕ Handling Multiple Datasets If you have multiple datasets, you can just merge them into the one jsonl file. For example, you can do something like this From 80cb14c61346bb2f8a02c1e0290cb7a1ac4bd821 Mon Sep 17 00:00:00 2001 From: canghua Date: Mon, 17 Aug 2026 15:45:12 +0800 Subject: [PATCH 3/3] update description --- docs/basic_usage/data_preparation.md | 32 +++++++------ scripts/build_vocab_mapping.py | 46 +++++++++++++------ .../test_scripts/test_build_vocab_mapping.py | 39 ++++++++++++++-- 3 files changed, 86 insertions(+), 31 deletions(-) diff --git a/docs/basic_usage/data_preparation.md b/docs/basic_usage/data_preparation.md index 9d70a4457..84a922ee5 100644 --- a/docs/basic_usage/data_preparation.md +++ b/docs/basic_usage/data_preparation.md @@ -341,13 +341,15 @@ See the [Training](training.md) guide for the complete run schema and supported combinations. -## 🔤 Build a draft vocabulary mapping +## 🔤 Build an EAGLE-family draft vocabulary mapping -A draft that predicts over a subset of the target vocabulary needs a `t2d`/`d2t` -mapping. `prepare_hidden_states.py` writes one as a side effect of capture, so -trying a second `draft_vocab_size` — or building a mapping for a corpus whose -features were captured before one was needed — otherwise means repeating the -capture. `scripts/build_vocab_mapping.py` derives it directly, without a GPU. +EAGLE3 and P-EAGLE drafts that predict over a subset of the target vocabulary +need a `t2d`/`d2t` mapping. DFlash, DSpark, and Domino currently train over the +full target vocabulary and do not consume this mapping. `prepare_hidden_states.py` +writes one as a side effect of EAGLE-family capture, so trying a second +`draft_vocab_size` — or building a mapping for a corpus whose features were +captured before one was needed — otherwise means repeating the capture. +`scripts/build_vocab_mapping.py` derives it directly, without a GPU. From the source JSONL, re-tokenizing with the same stack the capture used. Pass the same tokenizer, chat template, `--max-length`, and `--minimum-valid-tokens`, @@ -356,6 +358,7 @@ or the counts describe a different corpus than training sees: ```bash python scripts/build_vocab_mapping.py \ --data-path ./cache/dataset/sharegpt_train.jsonl \ + --target-model-path Qwen/Qwen3-8B \ --tokenizer-path Qwen/Qwen3-8B \ --chat-template qwen \ --max-length 2048 \ @@ -370,19 +373,22 @@ chat template: ```bash python scripts/build_vocab_mapping.py \ --hidden-states-path ./cache/hidden_states/sharegpt_train_Qwen3-8B \ + --target-model-path Qwen/Qwen3-8B \ --max-length 2048 \ --draft-model-config configs/qwen3-8b-eagle3.json \ --output-path ./cache/vocab_mapping/qwen3-8b-32k.pt ``` -The draft config supplies `vocab_size` — the map's length, which must match the -model's `t2d` buffer — and the default `draft_vocab_size`. Both modes cache the -per-corpus token counts, reusing them only while the corpus fingerprint is -unchanged, so surveying sizes costs one pass: +The target model config supplies `vocab_size`, the full vocabulary indexed by +the mapping, while the EAGLE-family draft config supplies the default +`draft_vocab_size`. Both modes cache the per-corpus token counts, reusing them +only while the corpus fingerprint is unchanged, so surveying sizes costs one +pass: ```bash python scripts/build_vocab_mapping.py \ --hidden-states-path ./cache/hidden_states/sharegpt_train_Qwen3-8B \ + --target-model-path Qwen/Qwen3-8B \ --draft-model-config configs/qwen3-8b-eagle3.json \ --draft-vocab-size 16000,32000,64000 ``` @@ -390,9 +396,9 @@ python scripts/build_vocab_mapping.py \ A comma-separated list only reports the coverage each size would reach and writes nothing; omitting `--output-path` reports coverage for a single size. -Point `model.vocab_mapping_path` at the written file. A disaggregated run -requires this, because producer and consumer cannot derive one shared mapping -independently. +Point `model.vocab_mapping_path` at the written file. A disaggregated +EAGLE-family run requires this, because producer and consumer cannot derive one +shared mapping independently. `--hidden-states-path` is exact but reads every feature file serially, which is impractical for a large gzipped corpus; prefer `--data-path` there. diff --git a/scripts/build_vocab_mapping.py b/scripts/build_vocab_mapping.py index 7216e3489..6ba44273e 100644 --- a/scripts/build_vocab_mapping.py +++ b/scripts/build_vocab_mapping.py @@ -1,4 +1,4 @@ -"""Build a draft vocabulary mapping from prepared offline features. +"""Build an EAGLE-family draft vocabulary mapping from a training corpus. Training can derive this map on its own, but only for a colocated offline run, and only by reading every feature file first -- which for gzipped features means @@ -11,10 +11,9 @@ to iterate on. Changing K therefore never requires regenerating hidden states, and never requires a second pass over them. -The map is emitted at the *draft config's* ``vocab_size``, which is the length -the model's ``t2d`` buffer is registered with. Sizing it from the target config -instead would produce a file that silently fails to load whenever the target -declares ``padded_vocab_size``. +The map is emitted at the target model config's ``vocab_size``, matching the +target-token IDs and logits the mapping indexes. The draft config supplies only +the selected ``draft_vocab_size``. Two sources, same numbers. ``--hidden-states-path`` reads the prepared features, which is exact but serial -- for a large gzipped dataset it is not @@ -27,17 +26,19 @@ python scripts/build_vocab_mapping.py \ --data-path ./cache/dataset/train.jsonl \ + --target-model-path Qwen/Qwen3-8B \ --tokenizer-path Qwen/Qwen3-8B --chat-template qwen --max-length 4096 \ - --draft-model-config configs/qwen3.6-27b-dspark.json \ + --draft-model-config configs/qwen3-8b-eagle3.json \ --draft-vocab-size 16000,32000,48000,64000 Then write the chosen one, reusing the cached counts: python scripts/build_vocab_mapping.py \ --data-path ./cache/dataset/train.jsonl \ + --target-model-path Qwen/Qwen3-8B \ --tokenizer-path Qwen/Qwen3-8B --chat-template qwen --max-length 4096 \ - --draft-model-config configs/qwen3.6-27b-dspark-draftvocab64k.json \ - --output-path ./cache/vocab_mapping/qwen3.6-27b-k64000.pt + --draft-model-config configs/qwen3-8b-eagle3.json \ + --output-path ./cache/vocab_mapping/qwen3-8b-k32000.pt """ from __future__ import annotations @@ -56,7 +57,7 @@ def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description=( - "Derive t2d/d2t from prepared offline features, without " + "Derive an EAGLE-family t2d/d2t mapping, without " "regenerating them and without a second pass per vocabulary size." ) ) @@ -81,6 +82,16 @@ def build_parser() -> argparse.ArgumentParser: "will describe a different corpus than training sees." ), ) + parser.add_argument( + "--target-model-path", + required=True, + help="Target model/config path supplying the full target vocab_size.", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow custom code while loading the target model config.", + ) parser.add_argument( "--tokenizer-path", default=None, @@ -132,10 +143,7 @@ def build_parser() -> argparse.ArgumentParser: "--draft-model-config", type=Path, required=True, - help=( - "Draft config JSON. Supplies vocab_size (the map's length, matching " - "the model's t2d buffer) and the default draft_vocab_size." - ), + help="EAGLE-family draft config JSON supplying the default draft_vocab_size.", ) parser.add_argument( "--draft-vocab-size", @@ -284,7 +292,7 @@ def tally_loss_tokens(dataset, *, vocab_size: int, batch_size: int = 512) -> Cou flat = torch.cat(selected).long() if int(flat.min()) < 0 or int(flat.max()) >= vocab_size: raise ValueError( - f"token id {int(flat.max())} is outside the draft config's " + f"token id {int(flat.max())} is outside the target model's " f"vocab_size {vocab_size}" ) totals += torch.bincount(flat, minlength=vocab_size) @@ -367,9 +375,17 @@ def write_mapping( def main(argv=None) -> int: args = build_parser().parse_args(argv) + from transformers import AutoConfig + + target_config = AutoConfig.from_pretrained( + args.target_model_path, + trust_remote_code=args.trust_remote_code, + ) + target_text_config = getattr(target_config, "text_config", target_config) + vocab_size = int(target_text_config.vocab_size) + with open(args.draft_model_config, encoding="utf-8") as handle: draft_config = json.load(handle) - vocab_size = int(draft_config["vocab_size"]) if args.draft_vocab_size is not None: sizes = [int(item) for item in str(args.draft_vocab_size).split(",") if item] diff --git a/tests/test_scripts/test_build_vocab_mapping.py b/tests/test_scripts/test_build_vocab_mapping.py index 4e1575679..adefc42b4 100644 --- a/tests/test_scripts/test_build_vocab_mapping.py +++ b/tests/test_scripts/test_build_vocab_mapping.py @@ -49,12 +49,20 @@ def _write_features(directory: Path, *, compress: bool, seed: int = 0) -> None: def _write_draft_config(path: Path, draft_vocab_size=None) -> None: - payload = {"vocab_size": VOCAB_SIZE} + payload = {} if draft_vocab_size is not None: payload["draft_vocab_size"] = draft_vocab_size path.write_text(json.dumps(payload), encoding="utf-8") +def _write_target_config(directory: Path, vocab_size=VOCAB_SIZE) -> None: + directory.mkdir(exist_ok=True) + (directory / "config.json").write_text( + json.dumps({"model_type": "llama", "vocab_size": vocab_size}), + encoding="utf-8", + ) + + class BuildVocabMappingTest(unittest.TestCase): def setUp(self): self._tmp = tempfile.TemporaryDirectory() @@ -62,6 +70,8 @@ def setUp(self): self.features = self.root / "features" self.features.mkdir() self.config = self.root / "draft.json" + self.target = self.root / "target" + _write_target_config(self.target) def tearDown(self): self._tmp.cleanup() @@ -74,6 +84,8 @@ def _run(self, *extra, compress=False, draft_vocab_size=DRAFT_VOCAB_SIZE): [ "--hidden-states-path", str(self.features), + "--target-model-path", + str(self.target), "--draft-model-config", str(self.config), *extra, @@ -115,6 +127,16 @@ def test_gzipped_features_are_supported(self): self.assertEqual(tuple(mapping["t2d"].shape), (VOCAB_SIZE,)) self.assertEqual(tuple(mapping["d2t"].shape), (DRAFT_VOCAB_SIZE,)) + def test_full_vocab_size_comes_from_the_target_config(self): + target_vocab_size = VOCAB_SIZE + 32 + _write_target_config(self.target, target_vocab_size) + out = self.root / "mapping.pt" + + self.assertEqual(0, self._run("--output-path", str(out))) + + mapping = torch.load(out, map_location="cpu") + self.assertEqual(tuple(mapping["t2d"].shape), (target_vocab_size,)) + def test_counts_are_cached_and_reused(self): """Counting is the expensive half; changing K must not repeat it.""" cache = self.features / ".token_counts.pt" @@ -169,6 +191,8 @@ def test_data_path_requires_the_capture_settings(self): [ "--data-path", str(jsonl), + "--target-model-path", + str(self.target), "--draft-model-config", str(self.config), ] @@ -183,6 +207,8 @@ def test_the_two_sources_are_mutually_exclusive(self): str(self.features), "--data-path", str(self.root / "train.jsonl"), + "--target-model-path", + str(self.target), "--draft-model-config", str(self.config), ] @@ -191,7 +217,14 @@ def test_the_two_sources_are_mutually_exclusive(self): def test_a_source_is_required(self): _write_draft_config(self.config, DRAFT_VOCAB_SIZE) with self.assertRaises(SystemExit): - main(["--draft-model-config", str(self.config)]) + main( + [ + "--target-model-path", + str(self.target), + "--draft-model-config", + str(self.config), + ] + ) def test_dataset_counts_are_cached_under_the_dataset_cache_dir(self): """The JSONL route caches too, so surveying K stays a one-time cost.""" @@ -289,7 +322,7 @@ def test_tally_rejects_ids_outside_the_vocabulary(self): from scripts.build_vocab_mapping import tally_loss_tokens dataset = Dataset.from_dict({"input_ids": [[VOCAB_SIZE]], "loss_mask": [[1]]}) - with self.assertRaisesRegex(ValueError, "outside the draft config"): + with self.assertRaisesRegex(ValueError, "outside the target model"): tally_loss_tokens(dataset, vocab_size=VOCAB_SIZE) def test_coverage_ratio_matches_the_definition(self):