Merge branch 'refs/heads/dev' into add_dspy

Author: Samoed

.github/workflows/test-presets.yaml

@@ -0,0 +1,38 @@
+name: test presets
+
+on:
+  push:
+    branches:
+      - dev
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ ubuntu-latest ]
+        python-version: [ "3.10", "3.11", "3.12" ]
+        include:
+          - os: windows-latest
+            python-version: "3.10"
+
+    steps:
+    - name: Checkout code
+      uses: actions/checkout@v4
+
+    - name: Setup Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: "pip"
+
+    - name: Install dependencies
+      run: |
+          pip install .
+          pip install pytest pytest-asyncio
+
+    - name: Run tests
+      run: |
+        pytest tests/pipeline/test_presets.py

README.md

@@ -30,7 +30,7 @@ Example of building an intent classifier in a couple of lines of code:
 from autointent import Pipeline, Dataset
 
 dataset = Dataset.from_json(path_to_json)
-pipeline = Pipeline.default_optimizer(multilabel=False)
+pipeline = Pipeline.from_preset("light")
 pipeline.fit(dataset)
-pipeline.predict(["show me my latest recent transactions"])
+pipeline.predict(["show me my latest transactions"])
 ```

autointent/__init__.py

@@ -7,6 +7,7 @@
 from ._dataset import Dataset
 from ._hash import Hasher
 from .context import Context, load_dataset
+from ._optimization_config import OptimizationConfig
 from ._pipeline import Pipeline
 
 
@@ -15,6 +16,7 @@
     "Dataset",
     "Embedder",
     "Hasher",
+    "OptimizationConfig",
     "Pipeline",
     "Ranker",
     "VectorIndex",

autointent/_callbacks/__init__.py

@@ -7,15 +7,17 @@
 
 REPORTERS = {cb.name: cb for cb in [WandbCallback, TensorBoardCallback]}
 
-REPORTERS_NAMES = list(REPORTERS.keys())
+REPORTERS_NAMES = Literal[tuple(REPORTERS.keys())]  # type: ignore[valid-type]
 
 
 def get_callbacks(reporters: list[str] | None) -> CallbackHandler:
-    """
-    Get the list of callbacks.
+    """Get the list of callbacks.
+
+    Args:
+        reporters: List of reporters to use.
 
-    :param reporters: List of reporters to use.
-    :return: Callback handler.
+    Returns:
+        CallbackHandler: Callback handler.
     """
     if not reporters:
         return CallbackHandler()

autointent/_callbacks/base.py

@@ -17,37 +17,37 @@ def __init__(self) -> None:
 
     @abstractmethod
     def start_run(self, run_name: str, dirpath: Path) -> None:
-        """
-        Start a new run.
+        """Start a new run.
 
-        :param run_name: Name of the run.
-        :param dirpath: Path to the directory where the logs will be saved.
+        Args:
+            run_name: Name of the run.
+            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.
+        """Start a new module.
 
-        :param module_name: Name of the module.
-        :param num: Number of the module.
-        :param module_kwargs: Module parameters.
+        Args:
+            module_name: Name of the module.
+            num: Number of the module.
+            module_kwargs: Module parameters.
         """
 
     @abstractmethod
     def log_value(self, **kwargs: dict[str, Any]) -> None:
-        """
-        Log data.
+        """Log data.
 
-        :param kwargs: Data to log.
+        Args:
+            kwargs: Data to log.
         """
 
     @abstractmethod
     def log_metrics(self, metrics: dict[str, Any]) -> None:
-        """
-        Log metrics during training.
+        """Log metrics during training.
 
-        :param metrics: Metrics to log.
+        Args:
+            metrics: Metrics to log.
         """
 
     @abstractmethod
@@ -60,8 +60,8 @@ def end_run(self) -> None:
 
     @abstractmethod
     def log_final_metrics(self, metrics: dict[str, Any]) -> None:
-        """
-        Log final metrics.
+        """Log final metrics.
 
-        :param metrics: Final metrics.
+        Args:
+            metrics: Final metrics.
         """

autointent/_callbacks/callback_handler.py

@@ -10,45 +10,49 @@ class CallbackHandler(OptimizerCallback):
     callbacks: list[OptimizerCallback]
 
     def __init__(self, callbacks: list[type[OptimizerCallback]] | None = None) -> None:
-        """Initialize the callback handler."""
+        """Initialize the callback handler.
+
+        Args:
+            callbacks: List of callback classes.
+        """
         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.
+        """Start a new run.
 
-        :param run_name: Name of the run.
-        :param dirpath: Path to the directory where the logs will be saved.
+        Args:
+            run_name: Name of the run.
+            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.
+        """Start a new module.
 
-        :param module_name: Name of the module.
-        :param num: Number of the module.
-        :param module_kwargs: Module parameters.
+        Args:
+            module_name: Name of the module.
+            num: Number of the module.
+            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.
+        """Log data.
 
-        :param kwargs: Data to log.
+        Args:
+            kwargs: Data to log.
         """
         self.call_events("log_value", **kwargs)
 
     def log_metrics(self, metrics: dict[str, Any]) -> None:
-        """
-        Log metrics during training.
+        """Log metrics during training.
 
-        :param metrics: Metrics to log.
+        Args:
+            metrics: Metrics to log.
         """
         self.call_events("log_metrics", metrics=metrics)
 
@@ -61,13 +65,19 @@ def end_run(self) -> None:
         self.call_events("end_run")
 
     def log_final_metrics(self, metrics: dict[str, Any]) -> None:
-        """
-        Log final metrics.
+        """Log final metrics.
 
-        :param metrics: Final metrics.
+        Args:
+            metrics: Final metrics.
         """
         self.call_events("log_final_metrics", metrics=metrics)
 
     def call_events(self, event: str, **kwargs: Any) -> None:  # noqa: ANN401
+        """Call events for all callbacks.
+
+        Args:
+            event: Event name.
+            kwargs: Event parameters.
+        """
         for callback in self.callbacks:
             getattr(callback, event)(**kwargs)

autointent/_callbacks/tensorboard.py

@@ -5,16 +5,16 @@
 
 
 class TensorBoardCallback(OptimizerCallback):
-    """
-    TensorBoard callback.
-
-    This callback logs the optimization process to TensorBoard.
-    """
+    """TensorBoard callback for logging the optimization process."""
 
     name = "tensorboard"
 
     def __init__(self) -> None:
-        """Initialize the callback."""
+        """Initializes the TensorBoard callback.
+
+        Attempts to import `torch.utils.tensorboard` first. If unavailable, tries to import `tensorboardX`.
+        Raises an ImportError if neither are installed.
+        """
         try:
             from torch.utils.tensorboard import SummaryWriter  # type: ignore[attr-defined]
 
@@ -32,22 +32,22 @@ def __init__(self) -> None:
                 raise ImportError(msg) from None
 
     def start_run(self, run_name: str, dirpath: Path) -> None:
-        """
-        Start a new run.
+        """Starts a new run and sets the directory for storing logs.
 
-        :param run_name: Name of the run.
-        :param dirpath: Path to the directory where the logs will be saved.
+        Args:
+            run_name: Name of the run.
+            dirpath: Path to the directory where 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.
+        """Starts a new module and initializes a TensorBoard writer for it.
 
-        :param module_name: Name of the module.
-        :param num: Number of the module.
-        :param module_kwargs: Module parameters.
+        Args:
+            module_name: Name of the module.
+            num: Identifier number of the module.
+            module_kwargs: Dictionary containing module parameters.
         """
         module_run_name = f"{self.run_name}_{module_name}_{num}"
         log_dir = Path(self.dirpath) / module_run_name
@@ -57,43 +57,38 @@ def start_module(self, module_name: str, num: int, module_kwargs: dict[str, Any]
         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.
+    def log_value(self, **kwargs: dict[str, int | float | Any]) -> None:
+        """Logs scalar or text values.
 
-        :param kwargs: Data to log.
+        Args:
+            **kwargs: Key-value pairs of data to log. Scalars will be logged as numerical values, others as text.
         """
-        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_metrics(self, metrics: dict[str, Any]) -> None:
-        """
-        Log metrics during training.
+        """Logs training metrics.
 
-        :param metrics: Metrics to log.
+        Args:
+            metrics: Dictionary of metrics to log.
         """
-        if self.module_writer is None:
-            msg = "start_run must be called before log_value."
-            raise RuntimeError(msg)
-
         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 log_final_metrics(self, metrics: dict[str, Any]) -> None:
-        """
-        Log final metrics.
+        """Logs final metrics at the end of training.
+
+        Args:
+            metrics: Dictionary of final metrics.
 
-        :param metrics: Final metrics.
+        Raises:
+            RuntimeError: If `start_run` has not been called before logging final metrics.
         """
         if self.module_writer is None:
             msg = "start_run must be called before log_final_metrics."
@@ -109,7 +104,11 @@ def log_final_metrics(self, metrics: dict[str, Any]) -> None:
                 self.module_writer.add_text(key, str(value))  # type: ignore[no-untyped-call]
 
     def end_module(self) -> None:
-        """End a module."""
+        """Ends the current module and closes the TensorBoard writer.
+
+        Raises:
+            RuntimeError: If `start_run` has not been called before ending the module.
+        """
         if self.module_writer is None:
             msg = "start_run must be called before end_module."
             raise RuntimeError(msg)
@@ -118,4 +117,4 @@ def end_module(self) -> None:
         self.module_writer.close()  # type: ignore[no-untyped-call]
 
     def end_run(self) -> None:
-        pass
+        """Ends the current run. This method is currently a placeholder."""