Implement annif upload and annif download commands for Hugging Face Hub integration

annif/cli.py

@@ -1,13 +1,14 @@
 """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
 
 import click
 import click_log
@@ -583,6 +584,112 @@
     click.echo("---")
 
 
+@cli.command("upload-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(
+    "--commit-message",
+    help="""The summary / title / first line of the generated commit.""",
+)
+@cli_util.common_options
+def run_upload_projects(project_ids_pattern, repo_id, token, commit_message):
+    """
+    Upload selected projects to a Hugging Face Hub repository
+    \f
+    This command zips the project directories and vocabularies of the projects
+    that match the given `project_ids_pattern`, and uploads the archives along
+    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")

annif/cli_util.py

@@ -1,18 +1,27 @@
 """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, hf_hub_download, list_repo_files
+from huggingface_hub.utils import HfHubHTTPError, HFValidationError
 
 import annif
-from annif.exception import ConfigurationException
+from annif.exception import ConfigurationException, OperationFailedException
 from annif.project import Access
 
 if TYPE_CHECKING:
@@ -230,6 +239,90 @@
     return list(itertools.product(limits, thresholds))
 
 
+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:

docs/source/commands.rst

@@ -66,6 +66,13 @@ Project administration
 
    N/A
 
+.. click:: annif.cli:run_upload_projects
+   :prog: annif upload-projects
+
+**REST equivalent**
+
+   N/A
+
 ****************************
 Subject index administration
 ****************************

pyproject.toml

@@ -48,6 +48,7 @@ python-dateutil = "2.8.*"
 tomli = { version = "2.0.*", python = "<3.11" }
 simplemma = "0.9.*"
 jsonschema = "4.17.*"
+huggingface-hub = "0.20.*"
 
 fasttext-wheel = {version = "0.9.2", optional = true}
 voikko = {version = "0.5.*", optional = true}