"""Hugging Face discovery and pinned downloads for the embedding viewer.""" from dataclasses import dataclass from hashlib import sha256 from pathlib import Path, PurePosixPath from urllib.parse import unquote, urlsplit import io import json import os import re import shutil from threading import Lock import time CACHE_ROOT = Path(__file__).resolve().parent / '.model-cache' os.environ.setdefault('HF_XET_CACHE', str(CACHE_ROOT / 'xet')) from huggingface_hub import HfApi, hf_hub_download from huggingface_hub.utils import validate_repo_id from tqdm.auto import tqdm from embeddings import embedding_name from app_environment import hosted_mode, hosted_download_limit HOSTED_DOWNLOAD_LOCK = Lock() DISK_RESERVE_BYTES = 2_000_000_000 @dataclass(frozen=True) class HubLocation: repo_id: str revision: str = 'main' filename: str | None = None @dataclass(frozen=True) class HubCatalog: repo_id: str commit: str files: dict[str, int | None] choices: tuple[str, ...] preferred: str | None = None @dataclass(frozen=True) class DownloadPlan: repo_id: str commit: str selection: str files: tuple[str, ...] weight_files: tuple[str, ...] model_path: Path tokenizer_path: Path | None directory: Path size: int | None embeddings_only: bool = False file_sizes: tuple[int | None, ...] = () def safe_filename(name): if not isinstance(name, str) or not name or '\\' in name or ':' in name or '\x00' in name: raise ValueError('Invalid repository file path.') path = PurePosixPath(name) if path.is_absolute() or any(part in ('', '.', '..') for part in name.split('/')): raise ValueError('Invalid repository file path.') return name def parse_location(value, revision=''): value = value.strip() filename, url_revision = None, None if '://' in value: url = urlsplit(value) if url.scheme != 'https' or url.netloc.lower() != 'huggingface.co': raise ValueError('Use a huggingface.co model URL or a repository ID such as owner/model.') pieces = url.path.strip('/').split('/') if len(pieces) >= 2 and pieces[0] in ('datasets', 'spaces'): raise ValueError('Choose a model repository, rather than a dataset or Space.') repo_parts = 1 if len(pieces) < 2 or pieces[1] in ('tree', 'blob', 'resolve') else 2 repo_id = '/'.join(pieces[:repo_parts]) if len(pieces) > repo_parts: tail = pieces[repo_parts:] if tail[0] not in ('tree', 'blob', 'resolve') or len(tail) < 2: raise ValueError('Use the model repository URL, or a tree/blob/resolve file URL.') url_revision = unquote(tail[1]) if tail[0] in ('blob', 'resolve') and len(tail) > 2: filename = safe_filename(unquote('/'.join(tail[2:]))) else: repo_id = value validate_repo_id(repo_id) return HubLocation(repo_id, revision.strip() or url_revision or 'main', filename) def discover_models(value, revision='', token=None): location = parse_location(value, revision) info = HfApi().model_info(location.repo_id, revision=location.revision, files_metadata=True, token=token or False, timeout=20) if not info.sha: raise ValueError('The repository did not return a commit revision.') files = {s.rfilename: s.size for s in info.siblings or []} choices = [] for name in files: if name.endswith('.safetensors.index.json'): choices.append(name) elif name.endswith('.safetensors') and not re.search(r'-\d{5}-of-\d{5}\.safetensors$', name): choices.append(name) elif name.lower().endswith('.gguf') and not re.search(r'-\d{5}-of-\d{5}\.gguf$', name, re.I): choices.append(name) choices.sort(key=lambda n: (not n.endswith('.index.json'), n.lower())) if location.filename is not None and location.filename not in choices: raise ValueError('The linked file is not supported here. Use the repository URL to choose GGUF or safetensors weights.') return HubCatalog(location.repo_id, info.sha, files, tuple(choices), location.filename if location.filename in choices else None) def selection_directory(catalog, selection, cache_root=CACHE_ROOT): identity = f'{catalog.repo_id}@{catalog.commit}/{selection}' directory = Path(cache_root).resolve() / sha256(identity.encode()).hexdigest()[:24] return directory def build_plan(catalog, selection, token=None, cache_root=CACHE_ROOT): if selection not in catalog.choices: raise ValueError('Choose one of the compatible files from this repository.') safe_filename(selection) directory = selection_directory(catalog, selection, cache_root) files = [selection] weights = [selection] embeddings_only = False if selection.endswith('.index.json'): # The small index is fetched before large files so the plan can show its size. index_path = hf_hub_download(catalog.repo_id, selection, revision=catalog.commit, local_dir=directory, token=token or False) with Path(index_path).open('r', encoding='utf-8') as source: index = json.load(source) mapping = index.get('weight_map') if not isinstance(mapping, dict) or not mapping: raise ValueError('The safetensors index has no weight map.') relevant = [value for key, value in mapping.items() if embedding_name(key)] selected_shards = relevant or list(mapping.values()) embeddings_only = bool(relevant) parent = PurePosixPath(selection).parent weights = [] for value in selected_shards: safe_filename(value) name = str(parent / value) safe_filename(name) if not name.endswith('.safetensors') or name not in catalog.files: raise ValueError('The model index refers to a missing or unsupported weights file.') if name not in weights: weights.append(name) files.extend(weights) parent = PurePosixPath(selection).parent tokenizer = next((name for name in (str(parent / 'tokenizer.json'), 'tokenizer.json') if name in catalog.files), None) if tokenizer: files.append(tokenizer) files = tuple(dict.fromkeys(files)) sizes = [catalog.files.get(name) for name in files] size = sum(sizes) if all(value is not None for value in sizes) else None model_path = directory if selection.endswith('.index.json') else directory / selection return DownloadPlan(catalog.repo_id, catalog.commit, selection, files, tuple(weights), model_path, directory / tokenizer if tokenizer else None, directory, size, embeddings_only, tuple(sizes)) def check_hosted_download(plan): """Enforce the per-model limit and leave room for the running Space. Complete cached files need no additional disk space. Incomplete downloads are counted at their full size, making this estimate conservative. """ limit = hosted_download_limit() if limit is None: return if plan.size is None or plan.size > limit: raise ValueError(f'This Space supports models up to {format_size(limit)} with a known download size. ' 'Choose smaller weights or embedding shards, or run the app locally.') remaining = plan.size if len(plan.file_sizes) == len(plan.files): for name, size in zip(plan.files, plan.file_sizes): safe_filename(name) path = plan.directory / name if size is not None and path.is_file() and path.stat().st_size == size: remaining -= size ancestor = plan.directory.resolve() while not ancestor.exists(): ancestor = ancestor.parent free = shutil.disk_usage(ancestor).free if remaining and remaining + DISK_RESERVE_BYTES > free: raise ValueError(f'Not enough free storage on this Space: {format_size(remaining)} still needs downloading, ' f'but only {format_size(free)} is free. We reserve 2 GB for the app. ' 'Choose smaller weights or try again after the Space cache is cleared.') def download_model(plan, token=None, progress=None): """Progress callbacks may run on SDK threads; callers should use a queue.""" locked = hosted_mode() if locked and not HOSTED_DOWNLOAD_LOCK.acquire(blocking=False): raise ValueError('Another model is downloading on this Space. Please try again when it finishes.') try: # Recheck inside the lock: another visitor may have consumed disk space # since this session displayed the download plan. check_hosted_download(plan) return _download_model(plan, token, progress) finally: if locked: HOSTED_DOWNLOAD_LOCK.release() def _download_model(plan, token=None, progress=None): total_files = len(plan.files) for index, filename in enumerate(plan.files): safe_filename(filename) def report(fraction, detail='', file_index=index, name=filename): if progress: progress((file_index + max(0., min(1., fraction))) / total_files, f'{file_index + 1}/{total_files} · {name}' + (f' · {detail}' if detail else '')) class DownloadProgress(tqdm): def __init__(self, *args, **kwargs): kwargs.pop('name', None) kwargs['file'] = io.StringIO() kwargs['disable'] = False self._last_report = 0. self._progress_callback = report super().__init__(*args, **kwargs) def update(self, n=1): changed = super().update(n) now = time.monotonic() if now - self._last_report > .15 and self.total: self._last_report = now self._progress_callback(self.n / self.total, f'{format_size(self.n)} / {format_size(self.total)}') return changed def close(self): if getattr(self, 'total', None): self._progress_callback(self.n / self.total, f'{format_size(self.n)} / {format_size(self.total)}') super().close() report(0.) hf_hub_download(plan.repo_id, filename, revision=plan.commit, token=token or False, local_dir=plan.directory, tqdm_class=DownloadProgress) report(1.) if not plan.model_path.exists() or not all((plan.directory / name).is_file() for name in plan.weight_files): raise ValueError('The download finished without a usable model file.') return plan.model_path, plan.tokenizer_path def format_size(value): if value is None: return 'size unavailable' for unit in ('B', 'KB', 'MB', 'GB', 'TB'): if value < 1000 or unit == 'TB': return f'{value:.1f} {unit}' value /= 1000 def friendly_hub_error(error, token=None): from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError, RevisionNotFoundError if isinstance(error, GatedRepoError): return 'This model is gated. Request or accept access on its Hugging Face page, then use a read token from the approved account.' if isinstance(error, RevisionNotFoundError): return 'That revision was not found. Check the branch, tag or commit.' if isinstance(error, RepositoryNotFoundError): return 'The repository was not found or is private. Check its ID and, if needed, provide a read token.' message = str(error) if token: message = message.replace(token, '[redacted]') return re.sub(r'hf_[A-Za-z0-9]+', '[redacted]', message)