Create hub_models.py

hub_models.py

@@ -0,0 +1,219 @@
+"""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 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
+
+
+@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
+
+
+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)
+
+
+def download_model(plan, token=None, progress=None):
+    """Progress callbacks may run on SDK threads; callers should use a queue."""
+    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)