Initial functionality for HF Hub download

NatLibFi · juhoinkinen · Apr 23, 2024 · Feb 1, 2024 · Feb 5, 2024 · Feb 6, 2024
commit 3135114ce0fb9fd69cf17ce47525f4e03de5b722
diff --git a/annif/cli.py b/annif/cli.py
@@ -1,12 +1,12 @@
 """Definitions for command-line (Click) commands for invoking Annif
 operations and printing the results to console."""
 
-
 import collections
 import importlib
 import json
 import os.path
 import re
+import shutil
 import sys
 from fnmatch import fnmatch
 
@@ -606,36 +606,90 @@
    with the projects configuration to the specified Hugging Face Hub repository.
    An authentication token and commit message can be given with options.
    """
    projects = [
        proj
        for proj in annif.registry.get_projects(min_access=Access.private).values()
        if fnmatch(proj.project_id, project_ids_pattern)
    ]
    click.echo(f"Uploading project(s): {', '.join([p.project_id for p in projects])}")

    commit_message = (
        commit_message
        if commit_message is not None
        else f"Upload project(s) {project_ids_pattern} with Annif"
    )

    project_dirs = {p.datadir for p in projects}
    vocab_dirs = {p.vocab.datadir for p in projects}
    data_dirs = project_dirs.union(vocab_dirs)

    for data_dir in data_dirs:
        zip_path = data_dir.split(os.path.sep, 1)[1] + ".zip"  # TODO Check this
        fobj = cli_util.archive_dir(data_dir)
        cli_util.upload_to_hf_hub(fobj, zip_path, repo_id, token, commit_message)
        fobj.close()

    for project in projects:
        config_path = project.project_id + ".cfg"
        fobj = cli_util.write_config(project)
        cli_util.upload_to_hf_hub(fobj, config_path, repo_id, token, commit_message)
         fobj.close()
 
 
+@cli.command("download-projects")
+@click.argument("project_ids_pattern")
+@click.argument("repo_id")
+@click.option(
+    "--token",
+    help="""Authentication token, obtained from the Hugging Face Hub.
+    Will default to the stored token.""",
+)
+@click.option(
+    "--revision",
+    help="""
+    An optional Git revision id which can be a branch name, a tag, or a commit
+    hash.
+    """,
+)
+@cli_util.common_options
+def run_download_projects(project_ids_pattern, repo_id, token, revision):
+    """
+    Download selected projects from a Hugging Face Hub repository
+    \f
+    This command downloads the project and vocabulary archives and the
+    configuration files of the projects that match the given
+    `project_ids_pattern` from the specified Hugging Face Hub repository and
+    unzips the archives to `data/` directory and places the configuration files
+    to `projects.d/` directory. An authentication token and revision can
+    be given with options.
+    """
+
+    project_ids = cli_util.get_selected_project_ids_from_hf_hub(
+        project_ids_pattern, repo_id, token, revision
+    )
+    click.echo(f"Downding project(s): {', '.join(project_ids)}")
+
+    if not os.path.isdir("projects.d"):
+        os.mkdir("projects.d")
+    vocab_ids = set()
+    for project_id in project_ids:
+        project_zip_local_cache_path = cli_util.download_from_hf_hub(
+            f"projects/{project_id}.zip", repo_id, token, revision
+        )
+        cli_util.unzip(project_zip_local_cache_path)
+        local_config_cache_path = cli_util.download_from_hf_hub(
+            f"{project_id}.cfg", repo_id, token, revision
+        )
+        vocab_ids.add(cli_util.get_vocab_id(local_config_cache_path))
+        shutil.copy(local_config_cache_path, "projects.d")  # TODO Disallow overwrite
+
+    for vocab_id in vocab_ids:
+        vocab_zip_local_cache_path = cli_util.download_from_hf_hub(
+            f"vocabs/{vocab_id}.zip", repo_id, token, revision
+        )
+        cli_util.unzip(vocab_zip_local_cache_path)
+
+
 @cli.command("completion")
 @click.option("--bash", "shell", flag_value="bash")
 @click.option("--zsh", "shell", flag_value="zsh")

diff --git a/annif/cli_util.py b/annif/cli_util.py
@@ -1,21 +1,23 @@
 """Utility functions for Annif CLI commands"""
+
 from __future__ import annotations
 
 import collections
 import configparser
 import io
 import itertools
 import os
 import pathlib
 import sys
 import tempfile
 import zipfile
+from fnmatch import fnmatch
 from typing import TYPE_CHECKING
 
 import click
 import click_log
 from flask import current_app
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, hf_hub_download, list_repo_files
 from huggingface_hub.utils import HfHubHTTPError, HFValidationError
 
 import annif
@@ -238,49 +240,89 @@


 def _is_train_file(fname):
    train_file_patterns = ("-train", "tmp-")
    for pat in train_file_patterns:
        if pat in fname:
            return True
    return False


 def archive_dir(data_dir):
    fp = tempfile.TemporaryFile()
    path = pathlib.Path(data_dir)
    fpaths = [fpath for fpath in path.glob("**/*") if not _is_train_file(fpath.name)]
    with zipfile.ZipFile(fp, mode="w") as zfile:
        for fpath in fpaths:
            logger.debug(f"Adding {fpath}")
            zfile.write(fpath)
    fp.seek(0)
    return fp


 def write_config(project):
    fp = tempfile.TemporaryFile(mode="w+t")
    config = configparser.ConfigParser()
    config[project.project_id] = project.config
    config.write(fp)  # This needs tempfile in text mode
    fp.seek(0)
    # But for upload fobj needs to be in binary mode
    return io.BytesIO(fp.read().encode("utf8"))


 def upload_to_hf_hub(fileobj, filename, repo_id, token, commit_message):
    api = HfApi()
    try:
        api.upload_file(
            path_or_fileobj=fileobj,
            path_in_repo=filename,
            repo_id=repo_id,
            token=token,
            commit_message=commit_message,
        )
    except (HfHubHTTPError, HFValidationError) as err:
         raise OperationFailedException(str(err))
 
 
+def get_selected_project_ids_from_hf_hub(project_ids_pattern, repo_id, token, revision):
+    all_repo_file_paths = _list_files_in_hf_hub(repo_id, token, revision)
+    return [
+        path.rsplit(".zip")[0].split("projects/")[1]  # TODO Try-catch this
+        for path in all_repo_file_paths
+        if fnmatch(path, f"projects/{project_ids_pattern}.zip")
+    ]
+
+
+def _list_files_in_hf_hub(repo_id, token, revision):
+    return [
+        repofile
+        for repofile in list_repo_files(repo_id=repo_id, token=token, revision=revision)
+    ]
+
+
+def download_from_hf_hub(filename, repo_id, token, revision):
+    try:
+        return hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            token=token,
+            revision=revision,
+        )
+    except (HfHubHTTPError, HFValidationError) as err:
+        raise OperationFailedException(str(err))
+
+
+def unzip(source_path):
+    with zipfile.ZipFile(source_path, "r") as zfile:
+        zfile.extractall()  # TODO Disallow overwrite
+
+
+def get_vocab_id(config_path):
+    config = configparser.ConfigParser()
+    config.read(config_path)
+    section = config.sections()[0]
+    return config[section]["vocab"]
+
+
 def _get_completion_choices(
     param: Argument,
 ) -> dict[str, AnnifVocabulary] | dict[str, AnnifProject] | list: