Reimplement NN ensemble using PyTorch

.github/workflows/cicd.yml

@@ -96,7 +96,7 @@ jobs:
         # Selectively install the optional dependencies for some Python versions
         # For Python 3.10:
         if [[ ${{ matrix.python-version }} == '3.10' ]]; then
-          uv sync --extra nn --extra omikuji --extra yake --extra voikko --extra stwfsa;
+          uv sync --extra nn --extra torch-cpu --extra omikuji --extra yake --extra voikko --extra stwfsa;
         fi
         # For Python 3.11:
         if [[ ${{ matrix.python-version }} == '3.11' ]]; then
@@ -106,7 +106,7 @@ jobs:
         fi
         # For Python 3.12:
         if [[ ${{ matrix.python-version }} == '3.12' ]]; then
-          uv sync --extra nn --extra fasttext --extra yake --extra stwfsa --extra voikko --extra spacy;
+          uv sync --extra nn --extra torch-cpu --extra fasttext --extra yake --extra stwfsa --extra voikko --extra spacy;
           # download the small English pretrained spaCy model needed by spacy analyzer
           uv run python -m spacy download en_core_web_sm --upgrade-strategy only-if-needed
         fi

Dockerfile

@@ -3,7 +3,7 @@ FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim
 LABEL org.opencontainers.image.authors="[email protected]"
 SHELL ["/bin/bash", "-c"]
 
-ARG optional_dependencies="voikko fasttext nn omikuji yake spacy stwfsa"
+ARG optional_dependencies="voikko fasttext nn omikuji yake spacy stwfsa torch-cpu"
 
 # Install system dependencies needed at runtime:
 RUN apt-get update && apt-get upgrade -y && \

README.md

@@ -96,7 +96,9 @@ Create a virtual environment and install dependencies:
 
     uv sync
 
-By default development dependencies are included. Use option `--extra` to install dependencies for selected optional features (`--extra "extra1 extra2"` for multiple extras), or install all of them with `--all-extras`. By default the virtual environment directory is `.venv` under the project directory.
+By default development dependencies are included. Use the option `--extra` to install dependencies for selected optional features (`--extra foo --extra bar` for multiple extras). There are also special extras for selecting the PyTorch backend for those optional features (`nn`) that depend on PyTorch: `--extra torch-cpu` will use the CPU backend and `--extra torch-cu128` will use the CUDA 12.8 backend. For a full set of extras with CPU-only PyTorch, use `uv sync --group all --extra torch-cpu`.
+
+By default the virtual environment directory is `.venv` under the project directory.
 
 You can run Annif in one of two ways:
 

annif/__init__.py

@@ -22,8 +22,6 @@
 def create_flask_app(config_name: str | None = None) -> Flask:
     """Create a Flask app to be used by the CLI."""
 
-    _set_tensorflow_loglevel()
-
     app = Flask(__name__)
     config_name = _get_config_name(config_name)
     logger.debug(f"creating flask app with configuration {config_name}")
@@ -93,20 +91,3 @@ def _get_config_name(config_name: str | None) -> str:
         else:
             config_name = "annif.default_config.ProductionConfig"  # pragma: no cover
     return config_name
-
-
-def _set_tensorflow_loglevel():
-    """Set TensorFlow log level based on Annif log level (--verbosity/-v
-    option) using an environment variable. INFO messages by TF are shown only on
-    DEBUG (or NOTSET) level of Annif."""
-    annif_loglevel = logger.getEffectiveLevel()
-    tf_loglevel_mapping = {
-        0: "0",  # NOTSET
-        10: "0",  # DEBUG
-        20: "1",  # INFO
-        30: "1",  # WARNING
-        40: "2",  # ERROR
-        50: "3",  # CRITICAL
-    }
-    tf_loglevel = tf_loglevel_mapping[annif_loglevel]
-    os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", tf_loglevel)

annif/backend/__init__.py

@@ -48,9 +48,7 @@ def _nn_ensemble() -> Type[AnnifBackend]:
 
         return nn_ensemble.NNEnsembleBackend
     except ImportError:
-        raise ValueError(
-            "Keras and TensorFlow not available, cannot use " + "nn_ensemble backend"
-        )
+        raise ValueError("PyTorch not available, cannot use nn_ensemble backend")
 
 
 def _omikuji() -> Type[AnnifBackend]: