annif/hfh_util.py

"""Utility functions for interactions with Hugging Face Hub."""

import binascii
import configparser
import importlib
import io
import os
import pathlib
import shutil
import tempfile
import time
import zipfile
from fnmatch import fnmatch
from typing import Any

import click
from flask import current_app

import annif
from annif.exception import OperationFailedException
from annif.project import Access, AnnifProject

logger = annif.logger


def get_matching_projects(pattern: str) -> list[AnnifProject]:
    """
    Get projects that match the given pattern.
    """
    return [
        proj
        for proj in annif.registry.get_projects(min_access=Access.private).values()
        if fnmatch(proj.project_id, pattern)
    ]


def prepare_commits(projects: list[AnnifProject], repo_id: str) -> tuple[list, list]:
    """Prepare and pre-upload data and config commit operations for projects to a
    Hugging Face Hub repository."""
    from huggingface_hub import preupload_lfs_files

    fobjs, operations = [], []
    data_dirs = {p.datadir for p in projects}
    vocab_dirs = {p.vocab.datadir for p in projects}
    all_dirs = data_dirs.union(vocab_dirs)

    for data_dir in all_dirs:
        fobj, operation = _prepare_datadir_commit(data_dir)
        preupload_lfs_files(repo_id, additions=[operation])
        fobjs.append(fobj)
        operations.append(operation)

    for project in projects:
        fobj, operation = _prepare_config_commit(project)
        fobjs.append(fobj)
        operations.append(operation)

    return fobjs, operations


def _prepare_datadir_commit(data_dir: str) -> tuple[io.BufferedRandom, Any]:
    from huggingface_hub import CommitOperationAdd

    zip_repo_path = data_dir.split(os.path.sep, 1)[1] + ".zip"
    fobj = _archive_dir(data_dir)
    operation = CommitOperationAdd(path_in_repo=zip_repo_path, path_or_fileobj=fobj)
    return fobj, operation


def _prepare_config_commit(project: AnnifProject) -> tuple[io.BytesIO, Any]:
    from huggingface_hub import CommitOperationAdd

    config_repo_path = project.project_id + ".cfg"
    fobj = _get_project_config(project)
    operation = CommitOperationAdd(path_in_repo=config_repo_path, path_or_fileobj=fobj)
    return fobj, operation


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


def _archive_dir(data_dir: str) -> io.BufferedRandom:
    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:
        zfile.comment = bytes(
            f"Archived by Annif {importlib.metadata.version('annif')}",
            encoding="utf-8",
        )
        for fpath in fpaths:
            logger.debug(f"Adding {fpath}")
            arcname = os.path.join(*fpath.parts[1:])
            zfile.write(fpath, arcname=arcname)
    fp.seek(0)
    return fp


def _get_project_config(project: AnnifProject) -> io.BytesIO:
    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 get_matching_project_ids_from_hf_hub(
    project_ids_pattern: str, repo_id: str, token, revision: str
) -> list[str]:
    """Get project IDs of the projects in a Hugging Face Model Hub repository that match
    the given pattern."""
    all_repo_file_paths = _list_files_in_hf_hub(repo_id, token, revision)
    return [
        path.rsplit(".cfg")[0]
        for path in all_repo_file_paths
        if fnmatch(path, f"{project_ids_pattern}.cfg")
    ]


def _list_files_in_hf_hub(repo_id: str, token: str, revision: str) -> list[str]:
    from huggingface_hub import list_repo_files
    from huggingface_hub.utils import HfHubHTTPError, HFValidationError

    try:
        return [
            repofile
            for repofile in list_repo_files(
                repo_id=repo_id, token=token, revision=revision
            )
        ]
    except (HfHubHTTPError, HFValidationError) as err:
        raise OperationFailedException(str(err))


def download_from_hf_hub(
    filename: str, repo_id: str, token: str, revision: str
) -> list[str]:
    from huggingface_hub import hf_hub_download
    from huggingface_hub.utils import HfHubHTTPError, HFValidationError

    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_archive(src_path: str, force: bool) -> None:
    """Unzip a zip archive of projects and vocabularies to a directory, by
    default data/ under current directory."""
    datadir = current_app.config["DATADIR"]
    with zipfile.ZipFile(src_path, "r") as zfile:
        archive_comment = str(zfile.comment, encoding="utf-8")
        logger.debug(
            f'Extracting archive {src_path}; archive comment: "{archive_comment}"'
        )
        for member in zfile.infolist():
            _unzip_member(zfile, member, datadir, force)


def _unzip_member(
    zfile: zipfile.ZipFile, member: zipfile.ZipInfo, datadir: str, force: bool
) -> None:
    dest_path = os.path.join(datadir, member.filename)
    if os.path.exists(dest_path) and not force:
        _handle_existing_file(member, dest_path)
        return
    logger.debug(f"Unzipping to {dest_path}")
    zfile.extract(member, path=datadir)
    _restore_timestamps(member, dest_path)


def _handle_existing_file(member: zipfile.ZipInfo, dest_path: str) -> None:
    if _are_identical_member_and_file(member, dest_path):
        logger.debug(f"Skipping unzip to {dest_path}; already in place")
    else:
        click.echo(f"Not overwriting {dest_path} (use --force to override)")


def _are_identical_member_and_file(member: zipfile.ZipInfo, dest_path: str) -> bool:
    path_crc = _compute_crc32(dest_path)
    return path_crc == member.CRC


def _restore_timestamps(member: zipfile.ZipInfo, dest_path: str) -> None:
    date_time = time.mktime(member.date_time + (0, 0, -1))
    os.utime(dest_path, (date_time, date_time))


def copy_project_config(src_path: str, force: bool) -> None:
    """Copy a given project configuration file to projects.d/ directory."""
    project_configs_dest_dir = "projects.d"
    os.makedirs(project_configs_dest_dir, exist_ok=True)

    dest_path = os.path.join(project_configs_dest_dir, os.path.basename(src_path))
    if os.path.exists(dest_path) and not force:
        if _are_identical_files(src_path, dest_path):
            logger.debug(f"Skipping copy to {dest_path}; already in place")
        else:
            click.echo(f"Not overwriting {dest_path} (use --force to override)")
    else:
        logger.debug(f"Copying to {dest_path}")
        shutil.copy(src_path, dest_path)


def _are_identical_files(src_path: str, dest_path: str) -> bool:
    src_crc32 = _compute_crc32(src_path)
    dest_crc32 = _compute_crc32(dest_path)
    return src_crc32 == dest_crc32


def _compute_crc32(path: str) -> int:
    if os.path.isdir(path):
        return 0

    size = 1024 * 1024 * 10  # 10 MiB chunks
    with open(path, "rb") as fp:
        crcval = 0
        while chunk := fp.read(size):
            crcval = binascii.crc32(chunk, crcval)
    return crcval


def get_vocab_id_from_config(config_path: str) -> str:
    """Get the vocabulary ID from a configuration file."""
    config = configparser.ConfigParser()
    config.read(config_path)
    section = config.sections()[0]
    return config[section]["vocab"]