Spaces:
Running
Running
| """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 | |
| class HubLocation: | |
| repo_id: str | |
| revision: str = 'main' | |
| filename: str | None = None | |
| class HubCatalog: | |
| repo_id: str | |
| commit: str | |
| files: dict[str, int | None] | |
| choices: tuple[str, ...] | |
| preferred: str | None = None | |
| 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) | |