init

Author: Samoed

autointent/_embedder.py

@@ -17,6 +17,7 @@
 from sentence_transformers import SentenceTransformer
 
 from ._hash import Hasher
+from .schemas._schemas import EmbedderConfig
 
 
 def get_embeddings_path(filename: str) -> Path:
@@ -40,7 +41,7 @@ class EmbedderDumpMetadata(TypedDict):
 
     model_name_or_path: str
     """Name of the hugging face model or a local path to sentence transformers dump."""
-    device: str
+    device: str | None
     """Torch notation for CPU or CUDA."""
     batch_size: int
     """Batch size used for embedding calculations."""
@@ -61,30 +62,19 @@ class Embedder:
     metadata_dict_name: str = "metadata.json"
     dump_dir: Path | None = None
 
-    def __init__(
-        self,
-        model_name_or_path: str | Path,
-        device: str = "cpu",
-        batch_size: int = 32,
-        max_length: int | None = None,
-        use_cache: bool = True,
-    ) -> None:
+    def __init__(self, embedder_config: EmbedderConfig) -> None:
         """
         Initialize the Embedder.
 
-        :param model_name_or_path: Path to a local model directory or a Hugging Face model name.
-        :param device: Device to run the model on (e.g., "cpu", "cuda").
-        :param batch_size: Batch size for embedding calculations.
-        :param max_length: Maximum sequence length for the embedding model.
-        :param use_cache: Flag indicating whether to cache intermediate embeddings.
+        :param embedder_config: Path to a local model directory or a Hugging Face model name.
         """
-        self.model_name = model_name_or_path
-        self.device = device
-        self.batch_size = batch_size
-        self.max_length = max_length
-        self.use_cache = use_cache
+        self.model_name = embedder_config.model_name
+        self.device = embedder_config.device
+        self.batch_size = embedder_config.batch_size
+        self.max_length = embedder_config.max_length
+        self.use_cache = embedder_config.use_cache
 
-        self.embedding_model = SentenceTransformer(str(model_name_or_path), device=device)
+        self.embedding_model = SentenceTransformer(self.model_name, device=self.device)
 
         self.logger = logging.getLogger(__name__)
 
@@ -132,9 +122,7 @@ def dump(self, path: Path) -> None:
             json.dump(metadata, file, indent=4)
 
     @classmethod
-    def load(
-        cls, path: Path | str, batch_size: int | None = None, use_cache: bool | None = None, device: str | None = None
-    ) -> "Embedder":
+    def load(cls, path: Path | str) -> "Embedder":
         """
         Load the embedding model and metadata from disk.
 
@@ -144,11 +132,13 @@ def load(
             metadata: EmbedderDumpMetadata = json.load(file)
 
         return cls(
-            model_name_or_path=metadata["model_name_or_path"],
-            device=device or metadata["device"],
-            batch_size=batch_size or metadata["batch_size"],
-            max_length=metadata["max_length"],
-            use_cache=use_cache or metadata["use_cache"],
+            EmbedderConfig(
+                model_name=metadata["model_name_or_path"],
+                device=metadata["device"],
+                batch_size=metadata["batch_size"],
+                max_length=metadata["max_length"],
+                use_cache=metadata["use_cache"],
+            )
         )
 
     def embed(self, utterances: list[str]) -> npt.NDArray[np.float32]:
@@ -189,4 +179,4 @@ def embed(self, utterances: list[str]) -> npt.NDArray[np.float32]:
             embeddings_path.parent.mkdir(parents=True, exist_ok=True)
             np.save(embeddings_path, embeddings)
 
-        return embeddings
+        return embeddings  # type: ignore[return-value]

autointent/_pipeline/_pipeline.py

@@ -9,7 +9,7 @@
 import yaml
 
 from autointent import Context, Dataset
-from autointent.configs import CrossEncoderConfig, EmbedderConfig, InferenceNodeConfig, LoggingConfig, VectorIndexConfig
+from autointent.configs import InferenceNodeConfig, LoggingConfig, VectorIndexConfig
 from autointent.custom_types import ListOfGenericLabels, NodeType
 from autointent.metrics import PREDICTION_METRICS_MULTILABEL
 from autointent.nodes import InferenceNode, NodeOptimizer
@@ -43,13 +43,11 @@ def __init__(
         if isinstance(nodes[0], NodeOptimizer):
             self.logging_config = LoggingConfig(dump_dir=None)
             self.vector_index_config = VectorIndexConfig()
-            self.embedder_config = EmbedderConfig()
-            self.cross_encoder_config = CrossEncoderConfig()
         elif not isinstance(nodes[0], InferenceNode):
             msg = "Pipeline should be initialized with list of NodeOptimizers or InferenceNodes"
             raise TypeError(msg)
 
-    def set_config(self, config: LoggingConfig | VectorIndexConfig | EmbedderConfig | CrossEncoderConfig) -> None:
+    def set_config(self, config: LoggingConfig | VectorIndexConfig) -> None:
         """
         Set configuration for the optimizer.
 
@@ -59,10 +57,6 @@ def set_config(self, config: LoggingConfig | VectorIndexConfig | EmbedderConfig
             self.logging_config = config
         elif isinstance(config, VectorIndexConfig):
             self.vector_index_config = config
-        elif isinstance(config, EmbedderConfig):
-            self.embedder_config = config
-        elif isinstance(config, CrossEncoderConfig):
-            self.cross_encoder_config = config
         else:
             msg = "unknown config type"
             raise TypeError(msg)
@@ -136,8 +130,6 @@ def fit(self, dataset: Dataset) -> Context:
         context = Context()
         context.set_dataset(dataset)
         context.configure_logging(self.logging_config)
-        context.configure_vector_index(self.vector_index_config, self.embedder_config)
-        context.configure_cross_encoder(self.cross_encoder_config)
 
         self._fit(context)
 

autointent/_ranker.py

@@ -19,14 +19,15 @@
 from torch import nn
 
 from autointent.custom_types import ListOfLabels
+from autointent.schemas._schemas import CrossEncoderConfig
 
 logger = logging.getLogger(__name__)
 
 
 class CrossEncoderMetadata(TypedDict):
     model_name: str
     train_classifier: bool
-    device: str
+    device: str | None
     max_length: int | None
     batch_size: int
 
@@ -104,32 +105,27 @@ class Ranker:
 
     def __init__(
         self,
-        model_name: str,
-        device: str = "cpu",
-        train_classifier: bool = False,
-        batch_size: int = 326,
-        max_length: int | None = None,
+        cross_encoder_config: CrossEncoderConfig,
         classifier_head: LogisticRegressionCV | None = None,
     ) -> None:
         """
         Initialize the Ranker.
 
-        :param model: The cross-encoder hugging face model name to use.
-        :param device: Device to run operations on, e.g., "cpu" or "cuda".
-        :param train_classifier: Whether to train a custom classifier, defaults to False.
-        :param batch_size: Batch size for processing text pairs, defaults to 326.
+        :param cross_encoder_config: The cross-encoder hugging face model name to use.
         :param max_length (int, optional): Max length for input sequences for the cross encoder.
         :param classifier_head (LogisticRegressionCV, optional): Classifier (to be used in restore procedure mainly).
         """
-        self.model_name = model_name
-        self.device = device
-        self.cross_encoder = st.CrossEncoder(model_name, trust_remote_code=True, device=device, max_length=max_length)  # type: ignore[arg-type]
+        self.cross_encoder = st.CrossEncoder(
+            cross_encoder_config.model_name,
+            trust_remote_code=True,
+            device=cross_encoder_config.device,  # type: ignore[arg-type]
+            max_length=cross_encoder_config.max_length,  # type: ignore[arg-type]
+        )
         self.train_classifier = False
-        self.batch_size = batch_size
-        self.max_length = max_length
         self._clf = classifier_head
+        self.cross_encoder_config = cross_encoder_config
 
-        if classifier_head is not None or train_classifier:
+        if classifier_head is not None or cross_encoder_config.train_head:
             self.train_classifier = True
             self._activations_list: list[npt.NDArray[Any]] = []
             self._hook_handler = self.cross_encoder.model.classifier.register_forward_hook(self._classifier_hook)
@@ -149,10 +145,14 @@ def _get_features_or_predictions(self, pairs: list[tuple[str, str]]) -> npt.NDAr
         :return: Numpy array of extracted features.
         """
         if not self.train_classifier:
-            return np.array(self.cross_encoder.predict(pairs, batch_size=self.batch_size, activation_fct=nn.Sigmoid()))
+            return np.array(
+                self.cross_encoder.predict(
+                    pairs, batch_size=self.cross_encoder_config.batch_size, activation_fct=nn.Sigmoid()  # type: ignore[arg-type]
+                )
+            )
 
         # put the data through, features will be taken in the hook
-        self.cross_encoder.predict(pairs, batch_size=self.batch_size)
+        self.cross_encoder.predict(pairs, batch_size=self.cross_encoder_config.batch_size)  # type: ignore[arg-type]
 
         res = np.concatenate(self._activations_list, axis=0)
         self._activations_list.clear()
@@ -222,8 +222,8 @@ def rank(
         Rank documents according to meaning closeness to the query.
 
         :param query: The reference document.
-        :query_docs: List of documents to rank
-        :top_k: how many document to return
+        :param query_docs: List of documents to rank
+        :param top_k: how many document to return
         :return: array of dictionaries of ranked items.
         """
         query_doc_pairs = [(query, doc) for doc in query_docs]
@@ -246,11 +246,11 @@ def save(self, path: str) -> None:
         dump_dir.mkdir(parents=True)
 
         metadata = CrossEncoderMetadata(
-            model_name=self.model_name,
+            model_name=self.cross_encoder_config.model_name,
             train_classifier=self.train_classifier,
-            device=self.device,
-            max_length=self.max_length,
-            batch_size=self.batch_size,
+            device=self.cross_encoder_config.device,
+            max_length=self.cross_encoder_config.max_length,
+            batch_size=self.cross_encoder_config.batch_size,
         )
 
         with (dump_dir / self.metadata_file_name).open("w") as file:
@@ -271,4 +271,13 @@ def load(cls, path: Path) -> "Ranker":
         with (path / cls.metadata_file_name).open() as file:
             metadata: CrossEncoderMetadata = json.load(file)
 
-        return cls(**metadata, classifier_head=clf)
+        return cls(
+            CrossEncoderConfig(
+                model_name=metadata["model_name"],
+                device=metadata["device"],
+                max_length=metadata["max_length"],
+                batch_size=metadata["batch_size"],
+                train_head=metadata["train_classifier"],
+            ),
+            classifier_head=clf,
+        )

autointent/_vector_index.py

@@ -16,11 +16,12 @@
 
 from autointent import Embedder
 from autointent.custom_types import ListOfLabels
+from autointent.schemas._schemas import EmbedderConfig
 
 
 class VectorIndexMetadata(TypedDict):
     embedder_model_name: str
-    embedder_device: str
+    embedder_device: str | None
     embedder_batch_size: int
     embedder_max_length: int | None
     embedder_use_cache: bool
@@ -42,31 +43,13 @@ class VectorIndex:
     _data_file = "data.json"
     _meta_data_file = "metadata.json"
 
-    def __init__(
-        self,
-        embedder_model_name: str,
-        embedder_device: str,
-        embedder_batch_size: int = 32,
-        embedder_max_length: int | None = None,
-        embedder_use_cache: bool = True,
-    ) -> None:
+    def __init__(self, embedder_config: EmbedderConfig) -> None:
         """
         Initialize the vector index.
 
-        :param embedder_model_name: Name of the embedding model to use.
-        :param embedder_device: Device for running the embedding model (e.g., "cpu", "cuda").
-        :param embedder_batch_size: Batch size for the embedder.
-        :param embedder_max_length: Maximum sequence length for the embedder.
-        :param embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
-        """
-        self.embedder = Embedder(
-            model_name_or_path=embedder_model_name,
-            batch_size=embedder_batch_size,
-            device=embedder_device,
-            max_length=embedder_max_length,
-            use_cache=embedder_use_cache,
-        )
-        self.embedder_device = embedder_device
+        :param embedder_config: Name of the embedding model to use.
+        """
+        self.embedder = Embedder(embedder_config)
 
         self.labels: ListOfLabels = []  # (n_samples,) or (n_samples, n_classes)
         self.texts: list[str] = []
@@ -233,11 +216,13 @@ def load(
             metadata: VectorIndexMetadata = json.load(file)
 
         instance = cls(
-            embedder_model_name=metadata["embedder_model_name"],
-            embedder_device=embedder_device or metadata["embedder_device"],
-            embedder_batch_size=embedder_batch_size or metadata["embedder_batch_size"],
-            embedder_max_length=metadata["embedder_max_length"],
-            embedder_use_cache=embedder_use_cache or metadata["embedder_use_cache"],
+            EmbedderConfig(
+                model_name=metadata["embedder_model_name"],
+                device=embedder_device or metadata["embedder_device"],
+                batch_size=embedder_batch_size or metadata["embedder_batch_size"],
+                max_length=metadata["embedder_max_length"],
+                use_cache=embedder_use_cache or metadata["embedder_use_cache"],
+            )
         )
 
         with (dir_path / cls._data_file).open() as file:

autointent/configs/__init__.py

@@ -2,18 +2,14 @@
 
 from ._inference_node import InferenceNodeConfig
 from ._optimization import (
-    CrossEncoderConfig,
     DataConfig,
-    EmbedderConfig,
     LoggingConfig,
     TaskConfig,
     VectorIndexConfig,
 )
 
 __all__ = [
-    "CrossEncoderConfig",
     "DataConfig",
-    "EmbedderConfig",
     "InferenceNodeConfig",
     "InferenceNodeConfig",
     "LoggingConfig",

autointent/configs/_optimization.py

@@ -55,44 +55,3 @@ class VectorIndexConfig(BaseModel):
 
     save_db: bool = False
     """Whether to save the vector index database or not"""
-
-
-class TransformerConfig(BaseModel):
-    """
-    Base class for configuration for the transformer.
-
-    Transformer is used under the hood in :py:class:`autointent.Embedder` and :py:class:`autointent.Ranker`.
-    """
-
-    batch_size: int = 32
-    """Batch size for the embedder"""
-    max_length: int | None = None
-    """Max length for the embedder. If None, the max length will be taken from model config"""
-    device: str = "cpu"
-    """Device to use for the vector index. Can be 'cpu', 'cuda', 'cuda:0', 'mps', etc."""
-
-
-class EmbedderConfig(TransformerConfig):
-    """
-    Configuration for the embedder.
-
-    The embedder is used to embed the data before training the model. These parameters
-    will be applied to the embedder used in the optimization process in vector db.
-    Only one model can be used globally.
-    """
-
-    use_cache: bool = True
-    """Whether to cache embeddings for reuse, improving performance in repeated operations."""
-
-
-class CrossEncoderConfig(TransformerConfig):
-    """
-    Configuration for the embedder.
-
-    The embedder is used to embed the data before training the model. These parameters
-    will be applied to the embedder used in the optimization process in vector db.
-    Only one model can be used globally.
-    """
-
-    train_head: bool = False
-    """Whether to train the ranking head of a cross encoder."""