Merge branch 'dev' into feat/cross_encoder_refactor

Author: Dmitryv-2024

.github/workflows/build-docs.yaml

@@ -1,9 +1,12 @@
-name: Build and publish docs
+name: Build and Publish Multi-Version Docs
 
 on:
   push:
     branches:
-    - dev
+      - dev
+  release:
+    types:
+      - published
   pull_request:
     branches:
       - dev
@@ -17,13 +20,18 @@ permissions:
   contents: write
 
 jobs:
-  publish:
-    name: build and publish docs
+  build-docs:
+    name: Build Documentation
     runs-on: ubuntu-latest
+
     steps:
-      - uses: actions/checkout@v4
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          fetch-tags: true
 
-      - name: set up python 3.10
+      - name: Set up Python 3.10
         uses: actions/setup-python@v5
         with:
           python-version: "3.10"
@@ -36,37 +44,31 @@ jobs:
         run: |
           sudo apt install pandoc
 
-      - name: install dependencies
+      - name: Install dependencies
         run: |
           poetry install --with docs
 
-      - name: Test documentation
+      - name: Run tests
+        if: github.event_name != 'workflow_dispatch'
         run: |
+          echo "Testing documentation build..."
           make test-docs
 
-      - name: build documentation
+      - name: Build documentation
+        if: ${{ github.ref == 'refs/heads/dev' }} && github.event_name != 'workflow_dispatch'
         run: |
           make docs
 
-      - name: save branch name without slashes
-        env:
-          BRANCH_NAME: ${{ github.head_ref || github.ref_name }}
+      - name: build multiversion documentation
+        if: github.event_name == 'release' || github.event_name == 'workflow_dispatch'
         run: |
-          BRANCH_NAME=${{ env.BRANCH_NAME }}
-          BRANCH_NAME=${BRANCH_NAME////_}
-          echo BRANCH_NAME=${BRANCH_NAME} >> $GITHUB_ENV
-
-      - name: Upload artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: ${{ format('github-pages-for-branch-{0}', env.BRANCH_NAME) }}
-          path: docs/build/
-          retention-days: 3
+          make multi-version-docs
 
       - name: Deploy to GitHub Pages
-        uses: JamesIves/github-pages-deploy-action@v4.6.4
-        if: ${{ github.ref == 'refs/heads/dev' }}
+        uses: peaceiris/actions-gh-pages@v3
+        if: github.event_name == 'release' || github.event_name == 'workflow_dispatch'
         with:
-          branch: gh-pages
-          folder: docs/build/html/
-          single-commit: True
+          github_token: ${{ github.token }}
+          publish_dir: docs/build/html/versions
+          destination_dir: versions
+          keep_files: true

.github/workflows/ruff.yml

@@ -5,4 +5,6 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: astral-sh/ruff-action@v1
+      - uses: astral-sh/ruff-action@v2
+        with:
+          version: "0.8.4"

.gitignore

@@ -177,3 +177,5 @@ poetry.lock
 indexes_dirnames.json
 tests_logs
 tests/logs
+runs/
+vector_db*

Makefile

@@ -31,18 +31,22 @@ docs:
 	$(poetry) python -m sphinx build -b html docs/source docs/build/html
 
 .PHONY: test-docs
-test-docs: docs
+test-docs:
 	$(poetry) python -m sphinx build -b doctest docs/source docs/build/html
 
 .PHONY: serve-docs
-serve-docs: docs
+serve-docs:
 	$(poetry) python -m http.server -d docs/build/html 8333
 
+.PHONY: multi-version-docs
+multi-version-docs:
+	$(poetry) sphinx-multiversion docs/source docs/build/html
+
 .PHONY: clean-docs
 clean-docs:
 	rm -rf docs/build
 	rm -rf docs/source/autoapi
 	rm -rf docs/source/user_guides
 
 .PHONY: all
-all: lint
+all: lint

autointent/_callbacks/__init__.py

@@ -0,0 +1,34 @@
+from autointent._callbacks.base import OptimizerCallback
+from autointent._callbacks.callback_handler import CallbackHandler
+from autointent._callbacks.tensorboard import TensorBoardCallback
+from autointent._callbacks.wandb import WandbCallback
+
+REPORTERS = {cb.name: cb for cb in [WandbCallback, TensorBoardCallback]}
+
+
+def get_callbacks(reporters: list[str] | None) -> CallbackHandler:
+    """
+    Get the list of callbacks.
+
+    :param reporters: List of reporters to use.
+    :return: Callback handler.
+    """
+    if not reporters:
+        return CallbackHandler()
+
+    reporters_cb = []
+    for reporter in reporters:
+        if reporter not in REPORTERS:
+            msg = f"Reporter {reporter} not supported. Supported reporters {','.join(REPORTERS)}"
+            raise ValueError(msg)
+        reporters_cb.append(REPORTERS[reporter])
+    return CallbackHandler(callbacks=reporters_cb)
+
+
+__all__ = [
+    "CallbackHandler",
+    "OptimizerCallback",
+    "TensorBoardCallback",
+    "WandbCallback",
+    "get_callbacks",
+]

autointent/_callbacks/base.py

@@ -0,0 +1,59 @@
+"""Base class for reporters (W&B, TensorBoard, etc)."""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+
+class OptimizerCallback(ABC):
+    """Base class for reporters (W&B, TensorBoard, etc)."""
+
+    # Implementation inspired by TrainerCallback from HuggingFace Transformers. https://github.com/huggingface/transformers/blob/91b8ab18b778ae9e2f8191866e018cd1dc7097be/src/transformers/trainer_callback.py#L260
+    name: str
+
+    @abstractmethod
+    def __init__(self) -> None:
+        pass
+
+    @abstractmethod
+    def start_run(self, run_name: str, dirpath: Path) -> None:
+        """
+        Start a new run.
+
+        :param run_name: Name of the run.
+        :param dirpath: Path to the directory where the logs will be saved.
+        """
+
+    @abstractmethod
+    def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None:
+        """
+        Start a new module.
+
+        :param module_name: Name of the module.
+        :param num: Number of the module.
+        :param module_kwargs: Module parameters.
+        """
+
+    @abstractmethod
+    def log_value(self, **kwargs: dict[str, Any]) -> None:
+        """
+        Log data.
+
+        :param kwargs: Data to log.
+        """
+
+    @abstractmethod
+    def end_module(self) -> None:
+        """End a module."""
+
+    @abstractmethod
+    def end_run(self) -> None:
+        """End a run."""
+
+    @abstractmethod
+    def log_final_metrics(self, metrics: dict[str, Any]) -> None:
+        """
+        Log final metrics.
+
+        :param metrics: Final metrics.
+        """

autointent/_callbacks/callback_handler.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+from typing import Any
+
+from autointent._callbacks.base import OptimizerCallback
+
+
+class CallbackHandler(OptimizerCallback):
+    """Internal class that just calls the list of callbacks in order."""
+
+    callbacks: list[OptimizerCallback]
+
+    def __init__(self, callbacks: list[type[OptimizerCallback]] | None = None) -> None:
+        """Initialize the callback handler."""
+        if not callbacks:
+            self.callbacks = []
+            return
+
+        self.callbacks = [cb() for cb in callbacks]
+
+    def start_run(self, run_name: str, dirpath: Path) -> None:
+        """
+        Start a new run.
+
+        :param run_name: Name of the run.
+        :param dirpath: Path to the directory where the logs will be saved.
+        """
+        self.call_events("start_run", run_name=run_name, dirpath=dirpath)
+
+    def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None:
+        """
+        Start a new module.
+
+        :param module_name: Name of the module.
+        :param num: Number of the module.
+        :param module_kwargs: Module parameters.
+        """
+        self.call_events("start_module", module_name=module_name, num=num, module_kwargs=module_kwargs)
+
+    def log_value(self, **kwargs: dict[str, Any]) -> None:
+        """
+        Log data.
+
+        :param kwargs: Data to log.
+        """
+        self.call_events("log_value", **kwargs)
+
+    def end_module(self) -> None:
+        """End a module."""
+        self.call_events("end_module")
+
+    def end_run(self) -> None:
+        """End a run."""
+        self.call_events("end_run")
+
+    def log_final_metrics(self, metrics: dict[str, Any]) -> None:
+        """
+        Log final metrics.
+
+        :param metrics: Final metrics.
+        """
+        self.call_events("log_final_metrics", metrics=metrics)
+
+    def call_events(self, event: str, **kwargs: Any) -> None:  # noqa: ANN401
+        for callback in self.callbacks:
+            getattr(callback, event)(**kwargs)

autointent/_callbacks/tensorboard.py

@@ -0,0 +1,105 @@
+from pathlib import Path
+from typing import Any
+
+from autointent._callbacks.base import OptimizerCallback
+
+
+class TensorBoardCallback(OptimizerCallback):
+    """
+    TensorBoard callback.
+
+    This callback logs the optimization process to TensorBoard.
+    """
+
+    name = "tensorboard"
+
+    def __init__(self) -> None:
+        """Initialize the callback."""
+        try:
+            from torch.utils.tensorboard import SummaryWriter  # type: ignore[attr-defined]
+
+            self.writer = SummaryWriter
+        except ImportError:
+            try:
+                from tensorboardX import SummaryWriter  # type: ignore[no-redef]
+
+                self.writer = SummaryWriter
+            except ImportError:
+                msg = (
+                    "TensorBoardCallback requires tensorboard to be installed. Either update your PyTorch version or"
+                    " install tensorboardX."
+                )
+                raise ImportError(msg) from None
+
+    def start_run(self, run_name: str, dirpath: Path) -> None:
+        """
+        Start a new run.
+
+        :param run_name: Name of the run.
+        :param dirpath: Path to the directory where the logs will be saved.
+        """
+        self.run_name = run_name
+        self.dirpath = dirpath
+
+    def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]) -> None:
+        """
+        Start a new module.
+
+        :param module_name: Name of the module.
+        :param num: Number of the module.
+        :param module_kwargs: Module parameters.
+        """
+        module_run_name = f"{self.run_name}_{module_name}_{num}"
+        log_dir = Path(self.dirpath) / module_run_name
+        self.module_writer = self.writer(log_dir=log_dir)  # type: ignore[no-untyped-call]
+
+        self.module_writer.add_text("module_info", f"Starting module {module_name}_{num}")  # type: ignore[no-untyped-call]
+        for key, value in module_kwargs.items():
+            self.module_writer.add_text(f"module_params/{key}", str(value))  # type: ignore[no-untyped-call]
+
+    def log_value(self, **kwargs: dict[str, Any]) -> None:
+        """
+        Log data.
+
+        :param kwargs: Data to log.
+        """
+        if self.module_writer is None:
+            msg = "start_run must be called before log_value."
+            raise RuntimeError(msg)
+
+        for key, value in kwargs.items():
+            if isinstance(value, int | float):
+                self.module_writer.add_scalar(key, value)
+            else:
+                self.module_writer.add_text(key, str(value))  # type: ignore[no-untyped-call]
+
+    def log_final_metrics(self, metrics: dict[str, Any]) -> None:
+        """
+        Log final metrics.
+
+        :param metrics: Final metrics.
+        """
+        if self.module_writer is None:
+            msg = "start_run must be called before log_final_metrics."
+            raise RuntimeError(msg)
+
+        log_dir = Path(self.dirpath) / "final_metrics"
+        self.module_writer = self.writer(log_dir=log_dir)  # type: ignore[no-untyped-call]
+
+        for key, value in metrics.items():
+            if isinstance(value, int | float):
+                self.module_writer.add_scalar(key, value)  # type: ignore[no-untyped-call]
+            else:
+                self.module_writer.add_text(key, str(value))  # type: ignore[no-untyped-call]
+
+    def end_module(self) -> None:
+        """End a module."""
+        if self.module_writer is None:
+            msg = "start_run must be called before end_module."
+            raise RuntimeError(msg)
+
+        self.module_writer.add_text("module_info", "Ending module")  # type: ignore[no-untyped-call]
+        self.module_writer.close()  # type: ignore[no-untyped-call]
+
+    def end_run(self) -> None:
+        pass