From 9b13a1db0b146260cfdaa0044ffcb1f97bed842f Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 12:47:28 +0200 Subject: [PATCH 1/8] feat: add support to specify the tokenizer version --- scripts/data.py | 16 +++++----------- src/post_training/config.py | 4 ++++ src/post_training/methods/common.py | 14 +++++++++++--- 3 files changed, 20 insertions(+), 14 deletions(-) diff --git a/scripts/data.py b/scripts/data.py index 1f0f49f..7ab9929 100755 --- a/scripts/data.py +++ b/scripts/data.py @@ -29,6 +29,8 @@ from pathlib import Path from typing import Any +from post_training.methods.common import build_tokenizer + _PROJECT_ROOT = Path(__file__).resolve().parent.parent if str(_PROJECT_ROOT / "src") not in sys.path: sys.path.insert(0, str(_PROJECT_ROOT / "src")) @@ -110,9 +112,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None: # ── Lazy imports (heavy) ──────────────────────────────────────── from datasets import load_dataset - from transformers import AutoTokenizer - from post_training.chat_templates.registry import get_chat_template from post_training.data.transforms import get_transform n = args.num_samples @@ -151,9 +151,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None: # ── Formatted via chat template ───────────────────────────── if args.show_formatted: _print_header(f"FORMATTED (chat_template='{config.data.chat_template}')") - tokenizer = AutoTokenizer.from_pretrained(config.model.name_or_path) - template_str = get_chat_template(config.data.chat_template) - tokenizer.chat_template = template_str + tokenizer = build_tokenizer(config) for i in range(min(n, len(ds))): row = ds[i] @@ -174,9 +172,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None: # ── Tokens ────────────────────────────────────────────────── if args.show_tokens: _print_header("TOKENS") - tokenizer = AutoTokenizer.from_pretrained(config.model.name_or_path) - template_str = get_chat_template(config.data.chat_template) - tokenizer.chat_template = template_str + tokenizer = build_tokenizer(config) for i in range(min(n, len(ds))): row = ds[i] @@ -208,9 +204,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None: print(f" Total rows : {len(ds)}") print(f" Columns : {ds.column_names}") - tokenizer = AutoTokenizer.from_pretrained(config.model.name_or_path) - template_str = get_chat_template(config.data.chat_template) - tokenizer.chat_template = template_str + tokenizer = build_tokenizer(config) # Sample up to 200 rows for token-length stats. sample_size = min(200, len(ds)) diff --git a/src/post_training/config.py b/src/post_training/config.py index fb65058..2e626ab 100644 --- a/src/post_training/config.py +++ b/src/post_training/config.py @@ -29,6 +29,10 @@ class ModelConfig: dtype: str = "bfloat16" revision: str | None = None + # Tokenizer to use. If None, defaults to the model's tokenizer. + tokenizer_name_or_path: str | None = None + tokenizer_revision: str | None = None + @dataclass class LRSchedulerKwargs: diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py index 60540d8..33a19da 100644 --- a/src/post_training/methods/common.py +++ b/src/post_training/methods/common.py @@ -42,12 +42,20 @@ def resolve_torch_dtype(name: str) -> torch.dtype: def build_tokenizer(config: PostTrainingConfig) -> AutoTokenizer: """Load the tokenizer and apply the configured chat template.""" - tokenizer = AutoTokenizer.from_pretrained( - config.model.name_or_path, revision=config.model.revision - ) + + if config.model.tokenizer_name_or_path is not None: + tokenizer_name_or_path = config.model.tokenizer_name_or_path + revision = config.model.tokenizer_revision + else: + tokenizer_name_or_path = config.model.name_or_path + revision = config.model.revision + + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=revision) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token + logger.info("Using tokenizer '%s' (revision=%s)", tokenizer_name_or_path, revision) + template_str = get_chat_template(config.data.chat_template) tokenizer.chat_template = template_str logger.info("Chat template set to '%s'.", config.data.chat_template) From 41d81ee8bb1b689a1a2db164f0e5f6677d3a8896 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 13:51:50 +0200 Subject: [PATCH 2/8] fix(data): keep the build_tokenizer import lazy in scripts/data.py --- scripts/data.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/scripts/data.py b/scripts/data.py index 7ab9929..9c4933e 100755 --- a/scripts/data.py +++ b/scripts/data.py @@ -29,8 +29,6 @@ from pathlib import Path from typing import Any -from post_training.methods.common import build_tokenizer - _PROJECT_ROOT = Path(__file__).resolve().parent.parent if str(_PROJECT_ROOT / "src") not in sys.path: sys.path.insert(0, str(_PROJECT_ROOT / "src")) @@ -114,6 +112,7 @@ def _run_inspect(args: argparse.Namespace, cli_overrides: list[str]) -> None: from datasets import load_dataset from post_training.data.transforms import get_transform + from post_training.methods.common import build_tokenizer n = args.num_samples From 731f684bb58957677d5de4389338e2e7e3605dda Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 13:52:40 +0200 Subject: [PATCH 3/8] fix(config): apply model.tokenizer_revision without a separate tokenizer path It was silently ignored when tokenizer_name_or_path was unset. Resolution now lives in ModelConfig.resolve_tokenizer(). --- src/post_training/config.py | 19 ++++++++++++++++++- src/post_training/methods/common.py | 11 ++--------- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/src/post_training/config.py b/src/post_training/config.py index 2e626ab..bd3be4e 100644 --- a/src/post_training/config.py +++ b/src/post_training/config.py @@ -29,10 +29,27 @@ class ModelConfig: dtype: str = "bfloat16" revision: str | None = None - # Tokenizer to use. If None, defaults to the model's tokenizer. + # Tokenizer source. If ``tokenizer_name_or_path`` is None the tokenizer + # comes from the model repo, and ``tokenizer_revision`` then pins it to a + # revision other than ``revision``. tokenizer_name_or_path: str | None = None tokenizer_revision: str | None = None + def resolve_tokenizer(self) -> tuple[str, str | None]: + """Return the ``(name_or_path, revision)`` pair to load the tokenizer from. + + ``tokenizer_name_or_path`` points at a different repo, so ``revision`` + does not apply to it: such a tokenizer is pinned by + ``tokenizer_revision`` alone. Without ``tokenizer_name_or_path`` the + tokenizer comes from the model repo, where ``tokenizer_revision`` + overrides ``revision``. + """ + if self.tokenizer_name_or_path is not None: + return self.tokenizer_name_or_path, self.tokenizer_revision + if self.tokenizer_revision is not None: + return self.name_or_path, self.tokenizer_revision + return self.name_or_path, self.revision + @dataclass class LRSchedulerKwargs: diff --git a/src/post_training/methods/common.py b/src/post_training/methods/common.py index 33a19da..213260b 100644 --- a/src/post_training/methods/common.py +++ b/src/post_training/methods/common.py @@ -42,20 +42,13 @@ def resolve_torch_dtype(name: str) -> torch.dtype: def build_tokenizer(config: PostTrainingConfig) -> AutoTokenizer: """Load the tokenizer and apply the configured chat template.""" - - if config.model.tokenizer_name_or_path is not None: - tokenizer_name_or_path = config.model.tokenizer_name_or_path - revision = config.model.tokenizer_revision - else: - tokenizer_name_or_path = config.model.name_or_path - revision = config.model.revision + tokenizer_name_or_path, revision = config.model.resolve_tokenizer() + logger.info("Loading tokenizer '%s' (revision=%s)", tokenizer_name_or_path, revision or "main") tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=revision) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - logger.info("Using tokenizer '%s' (revision=%s)", tokenizer_name_or_path, revision) - template_str = get_chat_template(config.data.chat_template) tokenizer.chat_template = template_str logger.info("Chat template set to '%s'.", config.data.chat_template) From 16e14f1f89ac5e7ae138cd167032592e5a6b73e0 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 13:53:10 +0200 Subject: [PATCH 4/8] feat(guardrails): show the resolved tokenizer in the submission review --- src/post_training/utils/guardrails.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/post_training/utils/guardrails.py b/src/post_training/utils/guardrails.py index cd6bf96..006348d 100644 --- a/src/post_training/utils/guardrails.py +++ b/src/post_training/utils/guardrails.py @@ -179,6 +179,13 @@ def run_guardrails(config: PostTrainingConfig, run_dir: Path, tokenize_only: boo _row("Attention impl", config.model.attn_implementation) _row("Dtype", config.model.dtype) _row("Revision", config.model.revision or "main") + tokenizer_name_or_path, tokenizer_revision = config.model.resolve_tokenizer() + _row( + "Tokenizer", + tokenizer_name_or_path, + warn=tokenizer_name_or_path != config.model.name_or_path, + ) + _row("Tokenizer revision", tokenizer_revision or "main") _row("Chat template", config.data.chat_template) _row("Data seed", str(config.data.seed)) for i, entry in enumerate(config.data.datasets): From bf3f05a6f5f41b1e36df7dd784e869b51b2b53c7 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 13:54:50 +0200 Subject: [PATCH 5/8] fix(prefetch): download the tokenizer so offline runs can load it --- scripts/submit.py | 1 + src/post_training/utils/prefetch.py | 51 +++++++++++++++++++++++++++-- 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/scripts/submit.py b/scripts/submit.py index f00137e..555241b 100644 --- a/scripts/submit.py +++ b/scripts/submit.py @@ -146,6 +146,7 @@ def main() -> None: # entirely — avoiding contention when every rank loads it at once. if prefetched is not None: config.model.name_or_path = prefetched.model + config.model.tokenizer_name_or_path = prefetched.tokenizer if prefetched.ref_model is not None: config.dpo.ref_model_name_or_path = prefetched.ref_model diff --git a/src/post_training/utils/prefetch.py b/src/post_training/utils/prefetch.py index 391bca4..ac8b2c5 100644 --- a/src/post_training/utils/prefetch.py +++ b/src/post_training/utils/prefetch.py @@ -31,16 +31,21 @@ class PrefetchedPaths: """Resolved local snapshot directories for the models in *config*. - Substituting these back into ``config.model.name_or_path`` (and, for DPO, + Substituting these back into ``config.model.name_or_path``, + ``config.model.tokenizer_name_or_path`` (and, for DPO, ``config.dpo.ref_model_name_or_path``) before training starts makes ``from_pretrained`` treat the model as a plain local directory, skipping huggingface_hub's cache-resolution/filelock path entirely — the source of spurious "does not appear to have a file named ..." errors when many ranks call ``from_pretrained`` on the same shared (e.g. Lustre) cache at once. + + ``tokenizer`` equals ``model`` when the tokenizer comes from the model + repo, because the model snapshot already contains the tokenizer files. """ model: str + tokenizer: str ref_model: str | None = None @@ -62,6 +67,38 @@ def _prefetch_model(name_or_path: str, revision: str | None = None) -> str: return local_path +# Files AutoTokenizer needs, so a separate tokenizer repo does not pull weights. +# ``config.json`` is included because AutoTokenizer falls back to its +# ``model_type`` when ``tokenizer_config.json`` names no tokenizer class. +_TOKENIZER_PATTERNS = [ + "tokenizer*", # tokenizer.json, tokenizer_config.json, tokenizer.model + "special_tokens_map.json", + "added_tokens.json", + "vocab.*", # vocab.json, vocab.txt + "merges.txt", + "chat_template.jinja", + "config.json", +] + + +def _prefetch_tokenizer(name_or_path: str, revision: str | None = None) -> str: + """Ensure the tokenizer files of *name_or_path* are cached; return their directory.""" + if _is_local(name_or_path): + logger.info("Tokenizer '%s' is a local path, skipping download.", name_or_path) + return name_or_path + logger.info( + "Downloading tokenizer '%s' (revision=%s) to HF cache...", name_or_path, revision or "main" + ) + local_path = snapshot_download( + repo_id=name_or_path, + repo_type="model", + revision=revision, + allow_patterns=_TOKENIZER_PATTERNS, + ) + logger.info("Tokenizer '%s' cached at '%s'.", name_or_path, local_path) + return local_path + + def _prefetch_dataset(entry: DatasetEntry) -> None: if _is_local(entry.path): logger.info("Dataset '%s' is a local path, skipping download.", entry.name) @@ -87,6 +124,16 @@ def prefetch_assets(config: PostTrainingConfig) -> PrefetchedPaths: model_path = _prefetch_model(config.model.name_or_path, revision=config.model.revision) + tokenizer_name_or_path, tokenizer_revision = config.model.resolve_tokenizer() + if (tokenizer_name_or_path, tokenizer_revision) == ( + config.model.name_or_path, + config.model.revision, + ): + # The tokenizer files are part of the model snapshot fetched above. + tokenizer_path = model_path + else: + tokenizer_path = _prefetch_tokenizer(tokenizer_name_or_path, revision=tokenizer_revision) + ref_model_path = None if config.method == "dpo" and config.dpo.ref_model_name_or_path is not None: ref_model_path = _prefetch_model(config.dpo.ref_model_name_or_path) @@ -95,4 +142,4 @@ def prefetch_assets(config: PostTrainingConfig) -> PrefetchedPaths: _prefetch_dataset(entry) logger.info("All assets pre-fetched successfully.") - return PrefetchedPaths(model=model_path, ref_model=ref_model_path) + return PrefetchedPaths(model=model_path, tokenizer=tokenizer_path, ref_model=ref_model_path) From fe6c0745a3def9b9414b1ae05eeccd956a934d7e Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 14:22:32 +0200 Subject: [PATCH 6/8] docs: document model.revision and the tokenizer options --- README.md | 5 ++++- configs/trl/dpo.yaml | 3 +++ configs/trl/sft.yaml | 3 +++ 3 files changed, 10 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 98016a8..76c092d 100644 --- a/README.md +++ b/README.md @@ -329,7 +329,7 @@ You must specify exactly one determining factor for training duration in the `tr - **Offline**: `offline: true` Disables Hugging Face Hub / Weights & Biases network calls (essential for air-gapped nodes). - **Prefetch assets**: `prefetch_assets: true` (default) - When submitting via `scripts/submit.py`, pre-fetches models and datasets on the login node before the job is submitted to warm the shared Hugging Face caches (useful for faster startup and for `offline: true` air-gapped runs). Runs whenever `offline` or `prefetch_assets` is true — set `prefetch_assets: false` to skip this step (e.g., assets are already warm in the shared cache). + When submitting via `scripts/submit.py`, pre-fetches models and datasets on the login node before the job is submitted to warm the shared Hugging Face caches (useful for faster startup and for `offline: true` air-gapped runs). Runs whenever `offline` or `prefetch_assets` is true — set `prefetch_assets: false` to skip this step (e.g., assets are already warm in the shared cache). A tokenizer from a different repo or revision (`model.tokenizer_name_or_path` / `model.tokenizer_revision`) is fetched too, without its model weights. - **Debug**: `debug.enabled: true` Forces `report_to: none`, uses a separate output directory, and allows overwriting existing runs. - **Tokenize only**: `--tokenize-only` (CLI flag on `train.py` / `submit.py`) @@ -444,6 +444,9 @@ model: name_or_path: "allenai/Olmo-3-1025-7B" attn_implementation: "flash_attention_3" dtype: "bfloat16" + revision: null # pin the model repo to a branch, tag, or commit + tokenizer_name_or_path: null # null = take the tokenizer from the model repo + tokenizer_revision: null # pins the tokenizer (overrides revision above) # -- Training hyper-parameters ----------------------------------------------- training: diff --git a/configs/trl/dpo.yaml b/configs/trl/dpo.yaml index d24a611..d5d908d 100644 --- a/configs/trl/dpo.yaml +++ b/configs/trl/dpo.yaml @@ -26,6 +26,9 @@ model: name_or_path: "allenai/Olmo-3-1025-7B" attn_implementation: "flash_attention_3" dtype: "bfloat16" + revision: null # pin the model repo to a branch, tag, or commit + tokenizer_name_or_path: null # null = take the tokenizer from the model repo + tokenizer_revision: null # pins the tokenizer (overrides revision above) # -- Training hyper-parameters ----------------------------------------------- training: diff --git a/configs/trl/sft.yaml b/configs/trl/sft.yaml index 5699cf2..a361aee 100644 --- a/configs/trl/sft.yaml +++ b/configs/trl/sft.yaml @@ -35,6 +35,9 @@ model: name_or_path: "allenai/Olmo-3-1025-7B" attn_implementation: "flash_attention_3" dtype: "bfloat16" + revision: null # pin the model repo to a branch, tag, or commit + tokenizer_name_or_path: null # null = take the tokenizer from the model repo + tokenizer_revision: null # pins the tokenizer (overrides revision above) # -- Training hyper-parameters ----------------------------------------------- training: From bc76ab6cff2e3ad3abfb6777e89bb86ce45f66e3 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 14:23:25 +0200 Subject: [PATCH 7/8] test: cover tokenizer resolution and tokenizer prefetching --- tests/test_config.py | 29 ++++++++++++++++- tests/test_prefetch.py | 70 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 tests/test_prefetch.py diff --git a/tests/test_config.py b/tests/test_config.py index 8857f00..3033faf 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,7 +3,7 @@ import pytest import yaml -from post_training.config import PostTrainingConfig +from post_training.config import ModelConfig, PostTrainingConfig from post_training.methods.common import build_common_training_kwargs @@ -107,3 +107,30 @@ def test_deepspeed_old_style_config_path_rejected(tmp_path): with pytest.raises(ValueError, match="deepspeed.config_path is no longer supported"): PostTrainingConfig.load(config_path) + + +@pytest.mark.parametrize( + ("tokenizer_name_or_path", "tokenizer_revision", "expected"), + [ + # No tokenizer override: the tokenizer follows the model. + (None, None, ("org/model", "r1")), + # A tokenizer revision alone pins the tokenizer inside the model repo. + (None, "r2", ("org/model", "r2")), + # A separate repo ignores model.revision, which pins a different repo. + ("org/tokenizer", None, ("org/tokenizer", None)), + ("org/tokenizer", "r2", ("org/tokenizer", "r2")), + ], +) +def test_resolve_tokenizer(tokenizer_name_or_path, tokenizer_revision, expected): + model = ModelConfig( + name_or_path="org/model", + revision="r1", + tokenizer_name_or_path=tokenizer_name_or_path, + tokenizer_revision=tokenizer_revision, + ) + + assert model.resolve_tokenizer() == expected + + +def test_resolve_tokenizer_without_any_revision(): + assert ModelConfig(name_or_path="org/model").resolve_tokenizer() == ("org/model", None) diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py new file mode 100644 index 0000000..fc82a89 --- /dev/null +++ b/tests/test_prefetch.py @@ -0,0 +1,70 @@ +"""Tests for which repos prefetch_assets downloads.""" + +import pytest + +from post_training.config import PostTrainingConfig +from post_training.utils import prefetch + + +@pytest.fixture +def downloads(monkeypatch): + """Record snapshot_download calls instead of contacting the Hub.""" + calls = [] + + def fake_snapshot_download(repo_id, repo_type=None, revision=None, allow_patterns=None): + calls.append((repo_id, revision, allow_patterns is not None)) + return f"/cache/{repo_id}@{revision or 'main'}" + + monkeypatch.setattr(prefetch, "snapshot_download", fake_snapshot_download) + monkeypatch.setattr(prefetch, "_prefetch_dataset", lambda entry: None) + return calls + + +def _config(**model_overrides) -> PostTrainingConfig: + config = PostTrainingConfig() + config.model.name_or_path = "org/model" + for key, value in model_overrides.items(): + setattr(config.model, key, value) + return config + + +def test_tokenizer_reuses_the_model_snapshot(downloads): + """The model snapshot already holds the tokenizer, so nothing extra is fetched.""" + paths = prefetch.prefetch_assets(_config(revision="r1")) + + assert downloads == [("org/model", "r1", False)] + assert paths.tokenizer == paths.model + + +def test_tokenizer_revision_is_fetched_separately(downloads): + """A tokenizer pinned to another revision of the model repo needs its own snapshot.""" + paths = prefetch.prefetch_assets(_config(revision="r1", tokenizer_revision="r2")) + + assert downloads == [("org/model", "r1", False), ("org/model", "r2", True)] + assert paths.model == "/cache/org/model@r1" + assert paths.tokenizer == "/cache/org/model@r2" + + +def test_model_revision_is_ignored_when_tokenizer_name_is_specified(downloads): + """The model revision is ignored when a tokenizer revision is specified.""" + paths = prefetch.prefetch_assets(_config(revision="r1", tokenizer_name_or_path="org/tokenizer")) + + assert downloads == [("org/model", "r1", False), ("org/tokenizer", None, True)] + assert paths.model == "/cache/org/model@r1" + assert paths.tokenizer == "/cache/org/tokenizer@main" + + +def test_tokenizer_repo_is_fetched_without_weights(downloads): + """A tokenizer from another repo is fetched with allow_patterns, so no weights.""" + paths = prefetch.prefetch_assets(_config(tokenizer_name_or_path="org/tokenizer")) + + assert downloads == [("org/model", None, False), ("org/tokenizer", None, True)] + assert paths.tokenizer == "/cache/org/tokenizer@main" + + +def test_local_tokenizer_path_is_not_downloaded(downloads, tmp_path): + """An existing local tokenizer directory is used as-is.""" + paths = prefetch.prefetch_assets(_config(tokenizer_name_or_path=str(tmp_path))) + + assert downloads == [("org/model", None, False)] + assert paths.tokenizer == str(tmp_path) From 6f6747dd907746864a1189f242e1261a9215e3da Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 25 Aug 2026 16:53:10 +0200 Subject: [PATCH 8/8] fix(prefetch): skip weight files instead of allow-listing tokenizer files --- src/post_training/utils/prefetch.py | 36 ++++++++++++++++++----------- tests/test_prefetch.py | 26 ++++++++++++++------- 2 files changed, 41 insertions(+), 21 deletions(-) diff --git a/src/post_training/utils/prefetch.py b/src/post_training/utils/prefetch.py index ac8b2c5..3097114 100644 --- a/src/post_training/utils/prefetch.py +++ b/src/post_training/utils/prefetch.py @@ -67,22 +67,32 @@ def _prefetch_model(name_or_path: str, revision: str | None = None) -> str: return local_path -# Files AutoTokenizer needs, so a separate tokenizer repo does not pull weights. -# ``config.json`` is included because AutoTokenizer falls back to its -# ``model_type`` when ``tokenizer_config.json`` names no tokenizer class. -_TOKENIZER_PATTERNS = [ - "tokenizer*", # tokenizer.json, tokenizer_config.json, tokenizer.model - "special_tokens_map.json", - "added_tokens.json", - "vocab.*", # vocab.json, vocab.txt - "merges.txt", - "chat_template.jinja", - "config.json", +# Weight formats, skipped when a repo is fetched for its tokenizer alone. +# Denying the large blobs is safer than allowing a list of tokenizer files: +# every tokenizer format arrives, including the sentencepiece and legacy names +# (``spiece.model``, ``sentencepiece.bpe.model``, ``bpe.codes``, ``*.spm``) and +# the ``additional_chat_templates/`` directory. The worst case is a few extra +# small files, not a missing file on an offline compute node. +_WEIGHT_PATTERNS = [ + "*.safetensors", + "*.bin", + "*.pt", + "*.pth", + "*.ckpt", + "*.h5", + "*.msgpack", + "*.gguf", + "*.onnx", + "*.onnx_data", + "*.ot", # rust_model.ot + "*.tflite", + "*.mlmodel", + "*.npz", ] def _prefetch_tokenizer(name_or_path: str, revision: str | None = None) -> str: - """Ensure the tokenizer files of *name_or_path* are cached; return their directory.""" + """Cache *name_or_path* without its weights; return the local directory.""" if _is_local(name_or_path): logger.info("Tokenizer '%s' is a local path, skipping download.", name_or_path) return name_or_path @@ -93,7 +103,7 @@ def _prefetch_tokenizer(name_or_path: str, revision: str | None = None) -> str: repo_id=name_or_path, repo_type="model", revision=revision, - allow_patterns=_TOKENIZER_PATTERNS, + ignore_patterns=_WEIGHT_PATTERNS, ) logger.info("Tokenizer '%s' cached at '%s'.", name_or_path, local_path) return local_path diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py index fc82a89..a957c63 100644 --- a/tests/test_prefetch.py +++ b/tests/test_prefetch.py @@ -11,8 +11,8 @@ def downloads(monkeypatch): """Record snapshot_download calls instead of contacting the Hub.""" calls = [] - def fake_snapshot_download(repo_id, repo_type=None, revision=None, allow_patterns=None): - calls.append((repo_id, revision, allow_patterns is not None)) + def fake_snapshot_download(repo_id, repo_type=None, revision=None, ignore_patterns=None): + calls.append((repo_id, revision, ignore_patterns)) return f"/cache/{repo_id}@{revision or 'main'}" monkeypatch.setattr(prefetch, "snapshot_download", fake_snapshot_download) @@ -32,7 +32,7 @@ def test_tokenizer_reuses_the_model_snapshot(downloads): """The model snapshot already holds the tokenizer, so nothing extra is fetched.""" paths = prefetch.prefetch_assets(_config(revision="r1")) - assert downloads == [("org/model", "r1", False)] + assert downloads == [("org/model", "r1", None)] assert paths.tokenizer == paths.model @@ -40,7 +40,10 @@ def test_tokenizer_revision_is_fetched_separately(downloads): """A tokenizer pinned to another revision of the model repo needs its own snapshot.""" paths = prefetch.prefetch_assets(_config(revision="r1", tokenizer_revision="r2")) - assert downloads == [("org/model", "r1", False), ("org/model", "r2", True)] + assert downloads == [ + ("org/model", "r1", None), + ("org/model", "r2", prefetch._WEIGHT_PATTERNS), + ] assert paths.model == "/cache/org/model@r1" assert paths.tokenizer == "/cache/org/model@r2" @@ -49,16 +52,23 @@ def test_model_revision_is_ignored_when_tokenizer_name_is_specified(downloads): """The model revision is ignored when a tokenizer revision is specified.""" paths = prefetch.prefetch_assets(_config(revision="r1", tokenizer_name_or_path="org/tokenizer")) - assert downloads == [("org/model", "r1", False), ("org/tokenizer", None, True)] + assert downloads == [ + ("org/model", "r1", None), + ("org/tokenizer", None, prefetch._WEIGHT_PATTERNS), + ] assert paths.model == "/cache/org/model@r1" assert paths.tokenizer == "/cache/org/tokenizer@main" def test_tokenizer_repo_is_fetched_without_weights(downloads): - """A tokenizer from another repo is fetched with allow_patterns, so no weights.""" + """A tokenizer from another repo is fetched without the weight files.""" paths = prefetch.prefetch_assets(_config(tokenizer_name_or_path="org/tokenizer")) - assert downloads == [("org/model", None, False), ("org/tokenizer", None, True)] + assert downloads == [ + ("org/model", None, None), + ("org/tokenizer", None, prefetch._WEIGHT_PATTERNS), + ] + assert "*.safetensors" in prefetch._WEIGHT_PATTERNS assert paths.tokenizer == "/cache/org/tokenizer@main" @@ -66,5 +76,5 @@ def test_local_tokenizer_path_is_not_downloaded(downloads, tmp_path): """An existing local tokenizer directory is used as-is.""" paths = prefetch.prefetch_assets(_config(tokenizer_name_or_path=str(tmp_path))) - assert downloads == [("org/model", None, False)] + assert downloads == [("org/model", None, None)] assert paths.tokenizer == str(tmp_path)