Blame for hub_models.py

"""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)