# ----------------------------------------------------------------------------- # 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] = { "": "<|vision_start|>", "": "<|vision_end|>", "": "<|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