Spaces:
Running on L40S
Running on L40S
File size: 7,482 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | # -----------------------------------------------------------------------------
# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
# -----------------------------------------------------------------------------
"""Helpers for consistent HuggingFace cache-only loading in tokenizer jobs."""
from __future__ import annotations
import atexit
import json
import shutil
import tempfile
from pathlib import Path
from typing import Any
NEMOTRON_VISION_TOKEN_REMAP: dict[str, str] = {
"<SPECIAL_20>": "<|vision_start|>",
"<SPECIAL_21>": "<|vision_end|>",
"<SPECIAL_22>": "<|image_pad|>",
}
_PATCHED_NEMOTRON_SNAPSHOT_CACHE: dict[tuple[str, str], str] = {}
_PATCHED_NEMOTRON_SNAPSHOT_DIRS: set[Path] = set()
def _cleanup_patched_nemotron_snapshots() -> None:
"""Delete any process-local writable Nemotron tokenizer snapshots."""
for snapshot_dir in list(_PATCHED_NEMOTRON_SNAPSHOT_DIRS):
shutil.rmtree(snapshot_dir, ignore_errors=True)
_PATCHED_NEMOTRON_SNAPSHOT_DIRS.clear()
_PATCHED_NEMOTRON_SNAPSHOT_CACHE.clear()
atexit.register(_cleanup_patched_nemotron_snapshots)
def _iter_snapshot_candidates(model_root: Path) -> list[Path]:
"""Return snapshot candidates in cache-preferred order for one HF model root."""
snapshots_root = model_root / "snapshots"
refs_root = model_root / "refs"
candidate_names: list[str] = []
seen_names: set[str] = set()
preferred_ref_paths = [refs_root / "main", refs_root / "master"]
for ref_path in preferred_ref_paths:
if not ref_path.is_file():
continue
ref_value = ref_path.read_text(encoding="utf-8").strip()
if ref_value and ref_value not in seen_names:
candidate_names.append(ref_value)
seen_names.add(ref_value)
if refs_root.is_dir():
for ref_path in sorted(path for path in refs_root.rglob("*") if path.is_file()):
ref_value = ref_path.read_text(encoding="utf-8").strip()
if ref_value and ref_value not in seen_names:
candidate_names.append(ref_value)
seen_names.add(ref_value)
if snapshots_root.is_dir():
for snapshot_path in sorted(path for path in snapshots_root.iterdir() if path.is_dir()):
snapshot_name = snapshot_path.name
if snapshot_name not in seen_names:
candidate_names.append(snapshot_name)
seen_names.add(snapshot_name)
return [snapshots_root / candidate_name for candidate_name in candidate_names]
def resolve_hf_snapshot_path(
model_name: str,
cache_dir: str | None,
required_files: tuple[str, ...] = ("tokenizer_config.json",),
) -> str | None:
"""Find a cached HuggingFace snapshot directory for one model.
Args:
model_name: HuggingFace model identifier such as ``Qwen/Qwen3-0.6B``.
cache_dir: Root HF cache directory, typically ``HF_HOME``.
required_files: Files that must exist inside the snapshot directory.
Returns:
The first matching snapshot directory, or ``None`` if not found.
"""
if cache_dir is None:
return None
safe_name = model_name.replace("/", "--")
for hub_prefix in ("hub", ""):
model_root = (
Path(cache_dir) / hub_prefix / f"models--{safe_name}"
if hub_prefix
else Path(cache_dir) / f"models--{safe_name}"
)
snapshots_root = model_root / "snapshots"
if not snapshots_root.is_dir():
continue
for snapshot_path in _iter_snapshot_candidates(model_root):
if not snapshot_path.is_dir():
continue
if all((snapshot_path / filename).exists() for filename in required_files):
return str(snapshot_path)
return None
def load_auto_tokenizer_from_cache(
model_name: str,
cache_dir: str | None,
required_files: tuple[str, ...] = ("tokenizer_config.json",),
**kwargs: Any,
) -> Any:
"""Load a tokenizer from the local HF cache when available.
When ``cache_dir`` is set but no local snapshot exists, this still enforces
``local_files_only=True`` so tokenizer workers fail fast instead of trying to
reach the public HuggingFace API from restricted clusters.
"""
from transformers import AutoTokenizer
local_snapshot = resolve_hf_snapshot_path(model_name, cache_dir, required_files=required_files)
if local_snapshot is not None:
return AutoTokenizer.from_pretrained(local_snapshot, local_files_only=True, **kwargs)
return AutoTokenizer.from_pretrained(
model_name,
cache_dir=cache_dir,
local_files_only=cache_dir is not None,
**kwargs,
)
def _patch_nemotron_tokenizer_snapshot_in_place(snapshot_dir: Path) -> None:
"""Rename reserved Nemotron placeholder tokens to tokenizer-visible vision tokens."""
tokenizer_json_path = snapshot_dir / "tokenizer.json"
if tokenizer_json_path.exists():
with tokenizer_json_path.open(encoding="utf-8") as f:
tokenizer_data = json.load(f)
for entry in tokenizer_data.get("added_tokens", []):
content = entry.get("content")
if content in NEMOTRON_VISION_TOKEN_REMAP:
entry["content"] = NEMOTRON_VISION_TOKEN_REMAP[content]
vocab = tokenizer_data.get("model", {}).get("vocab", {})
for old_name, new_name in NEMOTRON_VISION_TOKEN_REMAP.items():
if old_name in vocab:
vocab[new_name] = vocab.pop(old_name)
with tokenizer_json_path.open("w", encoding="utf-8") as f:
json.dump(tokenizer_data, f)
tokenizer_config_path = snapshot_dir / "tokenizer_config.json"
if tokenizer_config_path.exists():
with tokenizer_config_path.open(encoding="utf-8") as f:
tokenizer_config = json.load(f)
for entry in tokenizer_config.get("added_tokens_decoder", {}).values():
content = entry.get("content")
if content in NEMOTRON_VISION_TOKEN_REMAP:
entry["content"] = NEMOTRON_VISION_TOKEN_REMAP[content]
with tokenizer_config_path.open("w", encoding="utf-8") as f:
json.dump(tokenizer_config, f)
def prepare_nemotron_tokenizer_snapshot(
model_name: str,
cache_dir: str | None,
required_files: tuple[str, ...] = ("tokenizer_config.json", "tokenizer.json"),
) -> str | None:
"""Return a writable copied snapshot with Nemotron vision tokens remapped."""
local_snapshot = resolve_hf_snapshot_path(model_name, cache_dir, required_files=required_files)
if local_snapshot is None:
return None
source_snapshot = Path(local_snapshot)
cache_key = (model_name, str(source_snapshot.resolve()))
cached_snapshot = _PATCHED_NEMOTRON_SNAPSHOT_CACHE.get(cache_key)
if cached_snapshot is not None and Path(cached_snapshot).is_dir():
return cached_snapshot
patched_root = Path(
tempfile.mkdtemp(
prefix=(f"imaginaire4_nemotron_tokenizer_{model_name.replace('/', '--')}_{source_snapshot.name}_"),
dir=tempfile.gettempdir(),
)
)
shutil.copytree(source_snapshot, patched_root, dirs_exist_ok=True)
_patch_nemotron_tokenizer_snapshot_in_place(patched_root)
patched_snapshot = str(patched_root)
_PATCHED_NEMOTRON_SNAPSHOT_CACHE[cache_key] = patched_snapshot
_PATCHED_NEMOTRON_SNAPSHOT_DIRS.add(patched_root)
return patched_snapshot
|