File size: 4,880 Bytes
e65937c | 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 | """Shared utility functions for checkpoint conversion operations."""
from __future__ import annotations
import json
from pathlib import Path
import torch
from huggingface_hub import snapshot_download
from loguru import logger
from safetensors import safe_open
def download_checkpoint_from_hub(model_id: str, cache_dir: str | None = None) -> Path:
"""
Download a checkpoint from HuggingFace Hub.
:param model_id: HuggingFace model ID
:param cache_dir: Optional directory to cache downloads
:return: Local path to the downloaded checkpoint
:raises FileNotFoundError: If the checkpoint cannot be downloaded
:Example:
>>> path = download_checkpoint_from_hub("yuhuili/EAGLE-LLaMA3.1-Instruct-8B")
>>> print(path)
/home/user/.cache/huggingface/hub/models--yuhuili--EAGLE-LLaMA3.1-Instruct-8B/snapshots/...
"""
logger.info(f"Downloading checkpoint from HuggingFace: {model_id}")
try:
local_path = snapshot_download(
repo_id=model_id,
allow_patterns=["*.json", "*.safetensors", "*.bin", "*.index.json"],
cache_dir=cache_dir,
)
logger.debug(f"Downloaded to: {local_path}")
return Path(local_path)
except Exception as hf_exception:
logger.error(f"Failed to download checkpoint: {hf_exception}")
raise FileNotFoundError(f"Checkpoint not found: {model_id}") from hf_exception
def ensure_checkpoint_is_local(
checkpoint_path: str | Path, cache_dir: str | Path | None = None
) -> Path:
"""
Ensure we have a local copy of the checkpoint.
If the path exists locally, return it. Otherwise, treat it as a
HuggingFace model ID and download it.
:param checkpoint_path: Local path or HuggingFace model ID
:param cache_dir: Optional cache directory for downloads
:return: Path to local checkpoint directory
:Example:
>>> # Local path - returned as-is
>>> local = ensure_checkpoint_is_local("./my_checkpoint")
>>> # HuggingFace ID - downloaded first
>>> downloaded = ensure_checkpoint_is_local(
... "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
... )
"""
checkpoint_path = Path(checkpoint_path)
if checkpoint_path.exists():
logger.debug(f"Using local checkpoint: {checkpoint_path}")
return checkpoint_path
return download_checkpoint_from_hub(
model_id=str(checkpoint_path), cache_dir=str(cache_dir) if cache_dir else None
)
def load_checkpoint_config(checkpoint_dir: Path) -> dict:
"""
Load the config.json from a checkpoint directory.
:param checkpoint_dir: Path to checkpoint directory
:return: Config dictionary
:raises FileNotFoundError: If config.json is not found
:Example:
>>> config = load_checkpoint_config(Path("./checkpoint"))
>>> print(config["model_type"])
llama
"""
config_path = checkpoint_dir / "config.json"
if not config_path.exists():
raise FileNotFoundError(f"No config.json found at {checkpoint_dir}")
logger.debug(f"Loading config from: {config_path}")
with config_path.open() as f:
return json.load(f)
def load_checkpoint_weights(checkpoint_dir: Path) -> dict[str, torch.Tensor]:
"""
Load model weights from a checkpoint directory.
Supports both safetensors and PyTorch bin formats.
:param checkpoint_dir: Path to checkpoint directory
:return: Dictionary mapping weight names to tensors
:raises FileNotFoundError: If no weights are found
:raises NotImplementedError: If checkpoint is sharded
:Example:
>>> weights = load_checkpoint_weights(Path("./checkpoint"))
>>> print(f"Loaded {len(weights)} weights")
Loaded 50 weights
"""
weights = {}
safetensors_path = checkpoint_dir / "model.safetensors"
if safetensors_path.exists():
logger.debug(f"Loading safetensors weights from: {safetensors_path}")
with safe_open(safetensors_path, framework="pt") as f:
# safetensors requires iterating over keys() method
for key in f.keys(): # noqa: SIM118
weights[key] = f.get_tensor(key)
return weights
pytorch_path = checkpoint_dir / "pytorch_model.bin"
if pytorch_path.exists():
logger.debug(f"Loading PyTorch weights from: {pytorch_path}")
return torch.load(pytorch_path, map_location="cpu")
index_paths = [
checkpoint_dir / "model.safetensors.index.json",
checkpoint_dir / "pytorch_model.bin.index.json",
]
for index_path in index_paths:
if index_path.exists():
raise NotImplementedError(
f"Sharded checkpoint detected: {index_path}. "
"Please use a single-file checkpoint."
)
raise FileNotFoundError(f"No weights found at {checkpoint_dir}")
|