Add models configs (#124)

* init

* init

* fix name

* part fix loader

* fix tests

* fix tests

* lint

* fix typing

* update docs

* fix more docs

* lint

* remove commented code

* fix config

* fix last config

* fix default configs

* add dict type to init

* fix docstring

* lint

* resolve part conversations

* fix rerank scorer

* add prompts to encode

* add descriptions to schemes

* fix tests

* fix imports

* fix import config

Author: Samoed

autointent/_datafiles/default-multiclass-config.yaml

@@ -4,7 +4,7 @@
   search_space:
     - module_name: retrieval
       k: [10]
-      embedder_name:
+      embedder_config:
         - avsolatorio/GIST-small-Embedding-v0
         - sergeyzh/rubert-tiny-turbo
 - node_type: scoring
@@ -15,7 +15,7 @@
       weights: ["uniform", "distance", "closest"]
     - module_name: linear
     - module_name: dnnc
-      cross_encoder_name:
+      cross_encoder_config:
         - cross-encoder/ms-marco-MiniLM-L-6-v2
       k: [1, 3, 5, 10]
 - node_type: decision

autointent/_datafiles/default-multilabel-config.yaml

@@ -4,7 +4,7 @@
   search_space:
     - module_name: retrieval
       k: [10]
-      embedder_name:
+      embedder_config:
         - deepvk/USER-bge-m3
 - node_type: scoring
   target_metric: scoring_roc_auc

autointent/_datafiles/inference-config-example.yaml

@@ -2,7 +2,7 @@
   module_name: retrieval
   module_config:
     k: 10
-    model_name: sergeyzh/rubert-tiny-turbo
+    model_config: sergeyzh/rubert-tiny-turbo
   load_path: .
 - node_type: scoring
   module_name: knn

autointent/_embedder.py

@@ -17,6 +17,7 @@
 from sentence_transformers import SentenceTransformer
 
 from ._hash import Hasher
+from .schemas import EmbedderConfig, TaskTypeEnum
 
 
 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,22 @@ 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: Config of embedder.
         """
-        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.embedding_model = SentenceTransformer(str(model_name_or_path), device=device)
+        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_config = embedder_config
+
+        self.embedding_model = SentenceTransformer(
+            self.model_name, device=self.device, prompts=embedder_config.get_prompt_config()
+        )
 
         self.logger = logging.getLogger(__name__)
 
@@ -132,9 +125,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,18 +135,21 @@ 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]:
+    def embed(self, utterances: list[str], task_type: TaskTypeEnum | None = None) -> npt.NDArray[np.float32]:
         """
         Calculate embeddings for a list of utterances.
 
         :param utterances: List of input texts to calculate embeddings for.
+        :param task_type: Type of task for which embeddings are calculated.
         :return: A numpy array of embeddings.
         """
         if self.use_cache:
@@ -183,6 +177,7 @@ def embed(self, utterances: list[str]) -> npt.NDArray[np.float32]:
             convert_to_numpy=True,
             batch_size=self.batch_size,
             normalize_embeddings=True,
+            prompt_name=self.embedding_config.get_prompt_type(task_type),
         )
 
         if self.use_cache:

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, ValidationScheme
 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)
@@ -138,8 +132,8 @@ def fit(
         context = Context()
         context.set_dataset(dataset, scheme, n_folds)
         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)
+        context.configure_vector_index(self.vector_index_config)
+
         self.validate_modules(dataset)
         self._fit(context)
 

autointent/_ranker.py

@@ -20,14 +20,15 @@
 from torch import nn
 
 from autointent.custom_types import ListOfLabels
+from autointent.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
 
@@ -105,32 +106,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 | str | dict[str, Any],
         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: Config of 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_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
+        self.cross_encoder = st.CrossEncoder(
+            self.cross_encoder_config.model_name,
+            trust_remote_code=True,
+            device=self.cross_encoder_config.device,
+            max_length=self.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
 
-        if classifier_head is not None or train_classifier:
+        if classifier_head is not None or self.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)
@@ -150,10 +146,16 @@ 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(),
+                )
+            )
 
         # 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)
 
         res = np.concatenate(self._activations_list, axis=0)
         self._activations_list.clear()
@@ -223,8 +225,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]
@@ -247,11 +249,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:
@@ -272,7 +274,16 @@ 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,
+        )
 
     def clear_ram(self) -> None:
         self.cross_encoder.model.cpu()

autointent/_vector_index.py

@@ -16,11 +16,12 @@
 
 from autointent import Embedder
 from autointent.custom_types import ListOfLabels
+from autointent.schemas import EmbedderConfig, TaskTypeEnum
 
 
 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: Config 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] = []
@@ -81,7 +64,7 @@ def add(self, texts: list[str], labels: ListOfLabels) -> None:
         :param labels: List of labels corresponding to the texts.
         """
         self.logger.debug("Adding embeddings to vector index %s", self.embedder.model_name)
-        embeddings = self.embedder.embed(texts)
+        embeddings = self.embedder.embed(texts, TaskTypeEnum.passage)
 
         if not hasattr(self, "index"):
             self.index = faiss.IndexFlatIP(embeddings.shape[1])
@@ -120,7 +103,7 @@ def _search_by_text(self, texts: list[str], k: int) -> list[list[dict[str, Any]]
         :param k: Number of nearest neighbors to return.
         :return: List of search results for each query.
         """
-        query_embedding: npt.NDArray[np.float64] = self.embedder.embed(texts)  # type: ignore[assignment]
+        query_embedding: npt.NDArray[np.float64] = self.embedder.embed(texts, TaskTypeEnum.query)  # type: ignore[assignment]
         return self._search_by_embedding(query_embedding, k)
 
     def _search_by_embedding(self, embedding: npt.NDArray[Any], k: int) -> list[list[dict[str, Any]]]:
@@ -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: