Feat/cnn scorer (#230)

* added first version of CNNScorer

* Update optimizer_config.schema.json

* fix ruff

* fix ruff and mypy

* fix ruff

* fix furr for cnn.py

* added dump, removed cnn_kwargs

* Revert "added dump, removed cnn_kwargs"

This reverts commit efd9362.

* added test_cnn, added dump, fixed cnn_kwargs

* fix ruff

* fix ruff test_cnn

* Update cnn.py

* fixed from_context

* mypy fix

* fixes

* more fixes

* fix of kernel sizes

* Update test_cnn.py

* fix ruff

* fix test

* fix mypy

* added dump load test

* Update test_cnn.py

* Update test_cnn.py

* Update test_cnn.py

* try to fix dumper

* Update _dump_tools.py

* added dump-load

* fix

* fix ruff

* Update _dump_tools.py

* Update _dump_tools.py

* change order in dump

* Update _dump_tools.py

* Update _dump_tools.py

* fix dumper

* fix ruff

* Update _dump_tools.py

* add get_implicit_initialization_params, fix mypy partially

* Update textcnn.py

* Update textcnn.py

* Update textcnn.py

* fixed test

* fix mypy

* Update tensorboard.py

* Update _dump_tools.py

* add `_wrappers` submodule

* refactor cnn dumping

* codestyle

* bug fix

* move vocabulary management to textcnn from cnnscorer

* codestyle

* remove containers

* stylistic changes

* add device independence

* trigger ci

* add proper file encoding everywhere

---------

Co-authored-by: riapush <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

autointent/__init__.py

@@ -1,9 +1,7 @@
 """This is AutoIntent API reference."""
 
 from ._logging import setup_logging
-from ._ranker import Ranker
-from ._embedder import Embedder
-from ._vector_index import VectorIndex
+from ._wrappers import Ranker, Embedder, VectorIndex
 from ._dataset import Dataset
 from ._hash import Hasher
 from .context import Context, load_dataset

autointent/_dump_tools.py

@@ -10,6 +10,7 @@
 from peft import PeftModel
 from pydantic import BaseModel
 from sklearn.base import BaseEstimator
+from torch import nn
 from transformers import (  # type: ignore[attr-defined]
     AutoModelForSequenceClassification,
     AutoTokenizer,
@@ -19,14 +20,15 @@
 )
 
 from autointent import Embedder, Ranker, VectorIndex
+from autointent._wrappers import BaseTorchModule
 from autointent.configs import CrossEncoderConfig, EmbedderConfig
 from autointent.context.optimization_info import Artifact
 from autointent.schemas import TagsList
 
 ModuleSimpleAttributes = None | str | int | float | bool | list  # type: ignore[type-arg]
 
 ModuleAttributes: TypeAlias = (
-    ModuleSimpleAttributes | TagsList | np.ndarray | Embedder | VectorIndex | BaseEstimator | Ranker  # type: ignore[type-arg]
+    ModuleSimpleAttributes | TagsList | np.ndarray | Embedder | VectorIndex | BaseEstimator | Ranker | nn.Module  # type: ignore[type-arg]
 )
 
 logger = logging.getLogger(__name__)
@@ -43,6 +45,7 @@ class Dumper:
     pydantic_models: str = "pydantic"
     hf_models = "hf_models"
     hf_tokenizers = "hf_tokenizers"
+    torch_models = "torch_models"
     ptuning_models = "ptuning_models"
 
     @staticmethod
@@ -62,6 +65,7 @@ def make_subdirectories(path: Path, exists_ok: bool = False) -> None:
             path / Dumper.pydantic_models,
             path / Dumper.hf_models,
             path / Dumper.hf_tokenizers,
+            path / Dumper.torch_models,
             path / Dumper.ptuning_models,
         ]
         for subdir in subdirectories:
@@ -139,6 +143,20 @@ def dump(obj: Any, path: Path, exists_ok: bool = False, exclude: list[type[Any]]
                 except Exception as e:
                     msg = f"Error dumping HF model {key}: {e}"
                     logger.exception(msg)
+            elif isinstance(val, BaseTorchModule):
+                model_path = path / Dumper.torch_models / key
+                model_path.mkdir(parents=True, exist_ok=True)
+                try:
+                    class_info = {
+                        "module": val.__class__.__module__,
+                        "name": val.__class__.__name__,
+                    }
+                    with (model_path / "class_info.json").open("w") as f:
+                        json.dump(class_info, f)
+                    val.dump(model_path)
+                except Exception as e:
+                    msg = f"Error dumping torch model {key}: {e}"
+                    logger.exception(msg)
             elif isinstance(val, PreTrainedTokenizer | PreTrainedTokenizerFast):
                 tokenizer_path = path / Dumper.hf_tokenizers / key
                 tokenizer_path.mkdir(parents=True, exist_ok=True)
@@ -174,6 +192,7 @@ def load(  # noqa: C901, PLR0912, PLR0915
         pydantic_models: dict[str, Any] = {}
         hf_models: dict[str, Any] = {}
         hf_tokenizers: dict[str, Any] = {}
+        torch_models: dict[str, Any] = {}
 
         for child in path.iterdir():
             if child.name == Dumper.tags:
@@ -248,6 +267,18 @@ def load(  # noqa: C901, PLR0912, PLR0915
                     except Exception as e:  # noqa: PERF203
                         msg = f"Error loading HF tokenizer {tokenizer_dir.name}: {e}"
                         logger.exception(msg)
+            elif child.name == Dumper.torch_models:
+                try:
+                    for model_dir in child.iterdir():
+                        with (model_dir / "class_info.json").open("r") as f:
+                            class_info = json.load(f)
+                        module = importlib.import_module(class_info["module"])
+                        model_class: BaseTorchModule = getattr(module, class_info["name"])
+                        model = model_class.load(model_dir)
+                        torch_models[model_dir.name] = model
+                except Exception as e:
+                    msg = f"Error loading torch model {model_dir.name}: {e}"
+                    logger.exception(msg)
             else:
                 msg = f"Found unexpected child {child}"
                 logger.error(msg)
@@ -263,4 +294,5 @@ def load(  # noqa: C901, PLR0912, PLR0915
             | pydantic_models
             | hf_models
             | hf_tokenizers
+            | torch_models
         )

autointent/_utils.py

@@ -2,6 +2,8 @@
 
 from typing import TypeVar
 
+import torch
+
 T = TypeVar("T")
 
 
@@ -14,3 +16,12 @@ def _funcs_to_dict(*funcs: T) -> dict[str, T]:
         Dictionary of functions
     """
     return {func.__name__: func for func in funcs}  # type: ignore[attr-defined]
+
+
+def detect_device() -> str:
+    """Automatically detects CUDA, MPS and CPU."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.mps.is_available():
+        return "mps"
+    return "cpu"

autointent/_wrappers/__init__.py

@@ -0,0 +1,6 @@
+from .ranker import Ranker
+from .embedder import Embedder
+from .vector_index import VectorIndex
+from .base_torch_module import BaseTorchModule
+
+__all__ = ["BaseTorchModule", "Embedder", "Ranker", "VectorIndex"]

autointent/_wrappers/base_torch_module.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+import torch
+from torch import nn
+from typing_extensions import Self
+
+
+class BaseTorchModule(nn.Module, ABC):
+    @abstractmethod
+    def dump(self, path: Path) -> None:
+        """Dump torch module to disk.
+
+        This method encapsulates all the logic of dumping module's weights and
+        hyperparameters required for initialization from disk and nice inference.
+
+        Args:
+            path: path in file system
+        """
+
+    @classmethod
+    @abstractmethod
+    def load(cls, path: Path, device: str | None = None) -> Self:
+        """Load torch module from disk.
+
+        This method loads all weights and hyperparameters required for
+        initialization from disk and inference.
+
+        Args:
+            path: path in file system
+            device: torch notation for CPU, CUDA, MPS, etc. By default, it is inferred automatically.
+        """
+
+    @property
+    def device(self) -> torch.device:
+        """Torch device object where this module resides."""
+        return next(self.parameters()).device

autointent/_embedder.py autointent/_wrappers/embedder.pyautointent/_embedder.py renamed to autointent/_wrappers/embedder.py

@@ -19,8 +19,8 @@
 from sentence_transformers import SentenceTransformer
 from sentence_transformers.similarity_functions import SimilarityFunction
 
-from ._hash import Hasher
-from .configs import EmbedderConfig, TaskTypeEnum
+from autointent._hash import Hasher
+from autointent.configs import EmbedderConfig, TaskTypeEnum
 
 logger = logging.getLogger(__name__)
 

autointent/_ranker.py autointent/_wrappers/ranker.pyautointent/_ranker.py renamed to autointent/_wrappers/ranker.py

@@ -14,7 +14,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from autointent import Embedder
+from autointent._wrappers import Embedder
 from autointent.configs import EmbedderConfig, TaskTypeEnum, TokenizerConfig
 from autointent.custom_types import ListOfLabels
 

autointent/_vector_index.py autointent/_wrappers/vector_index.pyautointent/_vector_index.py renamed to autointent/_wrappers/vector_index.py

@@ -15,6 +15,7 @@
 from .scoring import (
     BERTLoRAScorer,
     BertScorer,
+    CNNScorer,
     DescriptionScorer,
     DNNCScorer,
     KNNScorer,
@@ -48,6 +49,7 @@ def _create_modules_dict(modules: list[type[T]]) -> dict[str, type[T]]:
         SklearnScorer,
         MLKnnScorer,
         BertScorer,
+        CNNScorer,
         BERTLoRAScorer,
         PTuningScorer,
     ]

autointent/modules/__init__.py

@@ -1,4 +1,5 @@
 from ._bert import BertScorer
+from ._cnn import CNNScorer
 from ._description import DescriptionScorer
 from ._dnnc import DNNCScorer
 from ._knn import KNNScorer, RerankScorer
@@ -11,6 +12,7 @@
 __all__ = [
     "BERTLoRAScorer",
     "BertScorer",
+    "CNNScorer",
     "DNNCScorer",
     "DescriptionScorer",
     "KNNScorer",