change usage of vector index

Author: voorhs

autointent/_pipeline/_pipeline.py

@@ -18,13 +18,10 @@
     HPOConfig,
     InferenceNodeConfig,
     LoggingConfig,
+    VectorIndexConfig,
+    get_default_vector_index_config,
 )
-from autointent.custom_types import (
-    ListOfGenericLabels,
-    NodeType,
-    SearchSpacePreset,
-    SearchSpaceValidationMode,
-)
+from autointent.custom_types import ListOfGenericLabels, NodeType, SearchSpacePreset, SearchSpaceValidationMode
 from autointent.metrics import DECISION_METRICS, DICISION_METRICS_MULTILABEL
 from autointent.nodes import InferenceNode, NodeOptimizer
 from autointent.utils import load_preset, load_search_space
@@ -64,11 +61,19 @@ def __init__(
             self.data_config = DataConfig()
             self.transformer_config = HFModelConfig()
             self.hpo_config = HPOConfig()
+            self.vector_index_config = get_default_vector_index_config()
         elif not isinstance(nodes[0], InferenceNode):
             assert_never(nodes)
 
     def set_config(
-        self, config: LoggingConfig | EmbedderConfig | CrossEncoderConfig | DataConfig | HFModelConfig | HPOConfig
+        self,
+        config: LoggingConfig
+        | EmbedderConfig
+        | CrossEncoderConfig
+        | DataConfig
+        | HFModelConfig
+        | HPOConfig
+        | VectorIndexConfig,
     ) -> None:
         """Set the configuration for the pipeline.
 
@@ -87,6 +92,8 @@ def set_config(
             self.transformer_config = config
         elif isinstance(config, HPOConfig):
             self.hpo_config = config
+        elif isinstance(config, VectorIndexConfig):
+            self.vector_index_config = config
         else:
             assert_never(config)
 
@@ -203,6 +210,7 @@ def fit(
         context.configure_transformer(self.cross_encoder_config)
         context.configure_transformer(self.transformer_config)
         context.configure_hpo(self.hpo_config)
+        context.configure_vector_index(self.vector_index_config)
 
         self.validate_modules(dataset, mode=incompatible_search_space)
 

autointent/_wrappers/vector_index/opensearch.py

@@ -107,7 +107,6 @@ def add(self, embeddings: NDArray[Any], documents: list[Document]) -> None:
                 }
             )
 
-
         self._init_index()
 
         # Use bulk API for efficient indexing

autointent/configs/__init__.py

@@ -11,7 +11,7 @@
     TaskTypeEnum,
     TokenizerConfig,
 )
-from ._vector_index import FaissConfig, OpenSearchConfig, VectorIndexConfig
+from ._vector_index import FaissConfig, OpenSearchConfig, VectorIndexConfig, get_default_vector_index_config
 
 __all__ = [
     "CrossEncoderConfig",
@@ -29,4 +29,5 @@
     "TorchTrainingConfig",
     "VectorIndexConfig",
     "VocabConfig",
+    "get_default_vector_index_config",
 ]

autointent/configs/_vector_index.py

@@ -18,3 +18,7 @@ class OpenSearchConfig(VectorIndexConfig):
     hosts: list[OpenSearchHost]
     index_name: str | None = None
     kwargs: dict[str, Any] = Field(default_factory=dict)  # TODO define set of options
+
+
+def get_default_vector_index_config() -> VectorIndexConfig:
+    return FaissConfig()

autointent/context/_context.py

@@ -8,7 +8,15 @@
 
 from autointent import Dataset
 from autointent._callbacks import CallbackHandler, get_callbacks
-from autointent.configs import CrossEncoderConfig, DataConfig, EmbedderConfig, HFModelConfig, HPOConfig, LoggingConfig
+from autointent.configs import (
+    CrossEncoderConfig,
+    DataConfig,
+    EmbedderConfig,
+    HFModelConfig,
+    HPOConfig,
+    LoggingConfig,
+    VectorIndexConfig,
+)
 
 from .data_handler import DataHandler
 from .optimization_info import OptimizationInfo
@@ -74,6 +82,12 @@ def configure_hpo(self, config: HPOConfig) -> None:
         else:
             assert_never(config)
 
+    def configure_vector_index(self, config: VectorIndexConfig) -> None:
+        if isinstance(config, VectorIndexConfig):
+            self.vector_index_config = config
+        else:
+            assert_never(config)
+
     def set_dataset(self, dataset: Dataset, config: DataConfig) -> None:
         """Set the datasets for training, validation and testing.
 

autointent/modules/embedding/_retrieval.py

@@ -5,7 +5,7 @@
 from pydantic import PositiveInt
 
 from autointent import Context, VectorIndex
-from autointent.configs import EmbedderConfig
+from autointent.configs import EmbedderConfig, VectorIndexConfig, get_default_vector_index_config
 from autointent.context.optimization_info import EmbeddingArtifact
 from autointent.custom_types import ListOfLabels
 from autointent.metrics import RETRIEVAL_METRICS_MULTICLASS, RETRIEVAL_METRICS_MULTILABEL
@@ -47,11 +47,13 @@ class RetrievalAimedEmbedding(BaseEmbedding):
     def __init__(
         self,
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
+        vector_index_config: VectorIndexConfig | None = None,
         k: PositiveInt = 10,
     ) -> None:
         self.k = k
         embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.embedder_config = embedder_config
+        self.vector_index_config = vector_index_config or get_default_vector_index_config()
 
         if self.k < 0 or not isinstance(self.k, int):
             msg = "`k` argument of `RetrievalAimedEmbedding` must be a positive int"
@@ -74,6 +76,7 @@ def from_context(
         return cls(
             k=k,
             embedder_config=embedder_config,
+            vector_index_config=context.vector_index_config,
         )
 
     def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
@@ -85,9 +88,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         """
         self._validate_task(labels)
 
-        self._vector_index = VectorIndex(
-            self.embedder_config,
-        )
+        self._vector_index = VectorIndex(self.embedder_config, config=self.vector_index_config)
         self._vector_index.add(utterances, labels)
 
     def score_ho(self, context: Context, metrics: list[str]) -> dict[str, float]:

autointent/modules/scoring/_dnnc/dnnc.py

@@ -9,7 +9,7 @@
 from pydantic import PositiveInt
 
 from autointent import Context, Ranker, VectorIndex
-from autointent.configs import CrossEncoderConfig, EmbedderConfig
+from autointent.configs import CrossEncoderConfig, EmbedderConfig, VectorIndexConfig, get_default_vector_index_config
 from autointent.custom_types import ListOfLabels
 from autointent.modules.base import BaseScorer
 
@@ -64,10 +64,12 @@ def __init__(
         k: PositiveInt = 5,
         cross_encoder_config: CrossEncoderConfig | str | dict[str, Any] | None = None,
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
+        vector_index_config: VectorIndexConfig | None = None,
     ) -> None:
         self.cross_encoder_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
         self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.k = k
+        self.vector_index_config = vector_index_config or get_default_vector_index_config()
 
         if self.k < 0 or not isinstance(self.k, int):
             msg = "`k` argument of `DNNCScorer` must be a positive int"
@@ -99,6 +101,7 @@ def from_context(
             k=k,
             embedder_config=embedder_config,
             cross_encoder_config=cross_encoder_config,
+            vector_index_config=context.vector_index_config,
         )
 
     def get_implicit_initialization_params(self) -> dict[str, Any]:
@@ -119,7 +122,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         """
         self._validate_task(labels)
 
-        self._vector_index = VectorIndex(self.embedder_config)
+        self._vector_index = VectorIndex(self.embedder_config, config=self.vector_index_config)
         self._vector_index.add(utterances, labels)
 
         self._cross_encoder = Ranker(self.cross_encoder_config, output_range="sigmoid")

autointent/modules/scoring/_knn/knn.py

@@ -7,7 +7,7 @@
 from pydantic import PositiveInt
 
 from autointent import Context, VectorIndex
-from autointent.configs import EmbedderConfig
+from autointent.configs import EmbedderConfig, VectorIndexConfig, get_default_vector_index_config
 from autointent.custom_types import ListOfLabels, WeightType
 from autointent.modules.base import BaseScorer
 
@@ -59,10 +59,12 @@ def __init__(
         k: PositiveInt = 5,
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
         weights: WeightType = "distance",
+        vector_index_config: VectorIndexConfig | None = None,
     ) -> None:
         self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.k = k
         self.weights = weights
+        self.vector_index_config = vector_index_config or get_default_vector_index_config()
 
         if self.k < 0 or not isinstance(self.k, int):
             msg = "`k` argument of `KNNScorer` must be a positive int"
@@ -92,9 +94,7 @@ def from_context(
             embedder_config = context.resolve_embedder()
 
         return cls(
-            embedder_config=embedder_config,
-            k=k,
-            weights=weights,
+            embedder_config=embedder_config, k=k, weights=weights, vector_index_config=context.vector_index_config
         )
 
     def get_implicit_initialization_params(self) -> dict[str, Any]:
@@ -113,7 +113,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         """
         self._validate_task(labels)
 
-        self._vector_index = VectorIndex(self.embedder_config)
+        self._vector_index = VectorIndex(self.embedder_config, config=self.vector_index_config)
         self._vector_index.add(utterances, labels)
 
     def predict(self, utterances: list[str]) -> npt.NDArray[Any]:

autointent/modules/scoring/_knn/rerank_scorer.py

@@ -7,7 +7,7 @@
 from pydantic import PositiveInt
 
 from autointent import Context, Ranker
-from autointent.configs import CrossEncoderConfig, EmbedderConfig
+from autointent.configs import CrossEncoderConfig, EmbedderConfig, VectorIndexConfig, get_default_vector_index_config
 from autointent.custom_types import ListOfLabels, WeightType
 
 from .knn import KNNScorer
@@ -42,6 +42,7 @@ def __init__(
         m: PositiveInt | None = None,
         cross_encoder_config: CrossEncoderConfig | str | dict[str, Any] | None = None,
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
+        vector_index_config: VectorIndexConfig | None = None,
     ) -> None:
         super().__init__(
             embedder_config=embedder_config,
@@ -50,6 +51,7 @@ def __init__(
         )
 
         self.cross_encoder_config = CrossEncoderConfig.from_search_config(cross_encoder_config)
+        self.vector_index_config = vector_index_config or get_default_vector_index_config()
 
         self.m = k if m is None else m
         self.use_cross_encoder_scores = use_cross_encoder_scores
@@ -94,6 +96,7 @@ def from_context(
             use_cross_encoder_scores=use_cross_encoder_scores,
             embedder_config=embedder_config,
             cross_encoder_config=cross_encoder_config,
+            vector_index_config=context.vector_index_config,
         )
 
     def get_implicit_initialization_params(self) -> dict[str, Any]:

autointent/modules/scoring/_mlknn/mlknn.py

@@ -8,7 +8,7 @@
 from typing_extensions import assert_never
 
 from autointent import Context, VectorIndex
-from autointent.configs import EmbedderConfig
+from autointent.configs import EmbedderConfig, VectorIndexConfig, get_default_vector_index_config
 from autointent.custom_types import ListOfLabels
 from autointent.modules.base import BaseScorer
 
@@ -67,11 +67,13 @@ def __init__(
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
         s: float = 1.0,
         ignore_first_neighbours: int = 0,
+        vector_index_config: VectorIndexConfig | None = None,
     ) -> None:
         self.k = k
         self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
         self.s = s
         self.ignore_first_neighbours = ignore_first_neighbours
+        self.vector_index_config = vector_index_config or get_default_vector_index_config()
 
         if self.k < 0 or not isinstance(self.k, int):
             msg = "`k` argument of `MLKnnScorer` must be a positive int"
@@ -109,6 +111,7 @@ def from_context(
             embedder_config=embedder_config,
             s=s,
             ignore_first_neighbours=ignore_first_neighbours,
+            vector_index_config=context.vector_index_config,
         )
 
     def get_implicit_initialization_params(self) -> dict[str, Any]:
@@ -127,15 +130,7 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         """
         self._validate_task(labels)
 
-        self._vector_index = VectorIndex(
-            EmbedderConfig(
-                model_name=self.embedder_config.model_name,
-                device=self.embedder_config.device,
-                batch_size=self.embedder_config.batch_size,
-                tokenizer_config=self.embedder_config.tokenizer_config,
-                use_cache=self.embedder_config.use_cache,
-            ),
-        )
+        self._vector_index = VectorIndex(embedder_config=self.embedder_config, config=self.vector_index_config)
         self._vector_index.add(utterances, labels)
 
         self._features = self._vector_index.get_all_embeddings()