Vector index and embeddings caching (#89)

Author: voorhs

autointent/_embedder.py

@@ -60,7 +60,7 @@ def __init__(
         device: str = "cpu",
         batch_size: int = 32,
         max_length: int | None = None,
-        use_cache: bool = False,
+        use_cache: bool = True,
     ) -> None:
         """
         Initialize the Embedder.
@@ -69,7 +69,7 @@ def __init__(
         :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 embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
+        :param use_cache: Flag indicating whether to cache intermediate embeddings.
         """
         self.model_name = model_name
         self.device = device

autointent/configs/_optimization_cli.py

@@ -109,7 +109,7 @@ class EmbedderConfig:
     """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"""
-    use_cache: bool = False
+    use_cache: bool = True
     """Flag indicating whether to cache embeddings for reuse, improving performance in repeated operations."""
     device: str = "cpu"
     """Device to use for the vector index. Can be 'cpu', 'cuda', 'cuda:0', 'mps', etc."""

autointent/context/vector_index_client/_vector_index.py

@@ -31,7 +31,7 @@ def __init__(
         embedder_device: str,
         embedder_batch_size: int = 32,
         embedder_max_length: int | None = None,
-        embedder_use_cache: bool = False,
+        embedder_use_cache: bool = True,
     ) -> None:
         """
         Initialize the vector index.
@@ -121,7 +121,7 @@ def _search_by_embedding(self, embedding: npt.NDArray[Any], k: int) -> list[list
             msg = "`embedding` should be a 2D array of shape (n_queries, dim_size)"
             raise ValueError(msg)
 
-        cos_sim, indices = self.index.search(embedding, k)
+        cos_sim, indices = self.index.search(embedding, k)  # TODO add caching similar to Embedder.embed() caching
         distances = 1 - cos_sim
 
         results = []

autointent/context/vector_index_client/_vector_index_client.py

@@ -32,7 +32,7 @@ def __init__(
         db_dir: str | Path | None,
         embedder_batch_size: int = 32,
         embedder_max_length: int | None = None,
-        embedder_use_cache: bool = False,
+        embedder_use_cache: bool = True,
     ) -> None:
         """
         Initialize the VectorIndexClient.

autointent/modules/embedding/_retrieval.py

@@ -72,7 +72,7 @@ def __init__(
         embedder_device: str = "cpu",
         batch_size: int = 32,
         max_length: int | None = None,
-        embedder_use_cache: bool = False,
+        embedder_use_cache: bool = True,
     ) -> None:
         """
         Initialize the RetrievalEmbedding.

autointent/modules/scoring/_description/description.py

@@ -10,19 +10,17 @@
 from sklearn.metrics.pairwise import cosine_similarity
 
 from autointent import Context, Embedder
-from autointent.context.vector_index_client import VectorIndex, VectorIndexClient
 from autointent.custom_types import LabelType
 from autointent.modules.abc import ScoringModule
 
 
 class DescriptionScorerDumpMetadata(TypedDict):
     """Metadata for dumping the state of a DescriptionScorer."""
 
-    db_dir: str
     n_classes: int
     multilabel: bool
-    batch_size: int
-    max_length: int | None
+    embedder_batch_size: int
+    embedder_max_length: int | None
 
 
 class DescriptionScorer(ScoringModule):
@@ -34,46 +32,40 @@ class DescriptionScorer(ScoringModule):
 
     :ivar weights_file_name: Filename for saving the description vectors (`description_vectors.npy`).
     :ivar embedder: The embedder used to generate embeddings for utterances and descriptions.
-    :ivar precomputed_embeddings: Flag indicating whether precomputed embeddings are used.
     :ivar embedding_model_subdir: Directory for storing the embedder's model files.
-    :ivar _vector_index: Internal vector index used when embeddings are precomputed.
-    :ivar db_dir: Directory path where the vector database is stored.
     :ivar name: Name of the scorer, defaults to "description".
 
     """
 
     weights_file_name: str = "description_vectors.npy"
     embedder: Embedder
-    precomputed_embeddings: bool = False
     embedding_model_subdir: str = "embedding_model"
-    _vector_index: VectorIndex
-    db_dir: str
     name = "description"
 
     def __init__(
         self,
         embedder_name: str,
         temperature: float = 1.0,
         embedder_device: str = "cpu",
-        batch_size: int = 32,
-        max_length: int | None = None,
-        embedder_use_cache: bool = False,
+        embedder_batch_size: int = 32,
+        embedder_max_length: int | None = None,
+        embedder_use_cache: bool = True,
     ) -> None:
         """
         Initialize the DescriptionScorer.
 
         :param embedder_name: Name of the embedder model.
         :param temperature: Temperature parameter for scaling logits, defaults to 1.0.
         :param embedder_device: Device to run the embedder on, e.g., "cpu" or "cuda".
-        :param batch_size: Batch size for embedding generation, defaults to 32.
-        :param max_length: Maximum sequence length for embedding, defaults to None.
+        :param embedder_batch_size: Batch size for embedding generation, defaults to 32.
+        :param embedder_max_length: Maximum sequence length for embedding, defaults to None.
         :param embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
         """
         self.temperature = temperature
         self.embedder_device = embedder_device
         self.embedder_name = embedder_name
-        self.batch_size = batch_size
-        self.max_length = max_length
+        self.embedder_batch_size = embedder_batch_size
+        self.embedder_max_length = embedder_max_length
         self.embedder_use_cache = embedder_use_cache
 
     @classmethod
@@ -93,19 +85,15 @@ def from_context(
         """
         if embedder_name is None:
             embedder_name = context.optimization_info.get_best_embedder()
-            precomputed_embeddings = True
-        else:
-            precomputed_embeddings = context.vector_index_client.exists(embedder_name)
 
-        instance = cls(
+        return cls(
             temperature=temperature,
             embedder_device=context.get_device(),
             embedder_name=embedder_name,
             embedder_use_cache=context.get_use_cache(),
+            embedder_batch_size=context.get_batch_size(),
+            embedder_max_length=context.get_max_length(),
         )
-        instance.precomputed_embeddings = precomputed_embeddings
-        instance.db_dir = str(context.get_db_dir())
-        return instance
 
     def get_embedder_name(self) -> str:
         """
@@ -136,39 +124,22 @@ def fit(
             self.n_classes = len(set(labels))
             self.multilabel = False
 
-        if self.precomputed_embeddings:
-            # this happens only when LinearScorer is within Pipeline opimization after RetrievalNode optimization
-            vector_index_client = VectorIndexClient(
-                self.embedder_device,
-                self.db_dir,
-                self.batch_size,
-                self.max_length,
-                self.embedder_use_cache,
-            )
-            vector_index = vector_index_client.get_index(self.embedder_name)
-            features = vector_index.get_all_embeddings()
-            if len(features) != len(utterances):
-                msg = "Vector index mismatches provided utterances"
-                raise ValueError(msg)
-            embedder = vector_index.embedder
-        else:
-            embedder = Embedder(
-                device=self.embedder_device,
-                model_name=self.embedder_name,
-                batch_size=self.batch_size,
-                max_length=self.max_length,
-                use_cache=self.embedder_use_cache,
-            )
-            features = embedder.embed(utterances)
-
         if any(description is None for description in descriptions):
             error_text = (
                 "Some intent descriptions (label_description) are missing (None). "
                 "Please ensure all intents have descriptions."
             )
             raise ValueError(error_text)
 
-        self.description_vectors = embedder.embed([desc for desc in descriptions if desc])
+        embedder = Embedder(
+            device=self.embedder_device,
+            model_name=self.embedder_name,
+            batch_size=self.embedder_batch_size,
+            max_length=self.embedder_max_length,
+            use_cache=self.embedder_use_cache,
+        )
+
+        self.description_vectors = embedder.embed(descriptions)
         self.embedder = embedder
 
     def predict(self, utterances: list[str]) -> NDArray[np.float64]:
@@ -198,11 +169,10 @@ def dump(self, path: str) -> None:
         :param path: Path to the directory where assets will be dumped.
         """
         self.metadata = DescriptionScorerDumpMetadata(
-            db_dir=str(self.db_dir),
             n_classes=self.n_classes,
             multilabel=self.multilabel,
-            batch_size=self.batch_size,
-            max_length=self.max_length,
+            embedder_batch_size=self.embedder_batch_size,
+            embedder_max_length=self.embedder_max_length,
         )
 
         dump_dir = Path(path)
@@ -232,7 +202,7 @@ def load(self, path: str) -> None:
         self.embedder = Embedder(
             device=self.embedder_device,
             model_name=embedder_dir,
-            batch_size=self.metadata["batch_size"],
-            max_length=self.metadata["max_length"],
+            batch_size=self.metadata["embedder_batch_size"],
+            max_length=self.metadata["embedder_max_length"],
             use_cache=self.embedder_use_cache,
         )

autointent/modules/scoring/_dnnc/dnnc.py

@@ -51,7 +51,6 @@ class DNNCScorer(ScoringModule):
 
     :ivar crossencoder_subdir: Subdirectory for storing the cross-encoder model (`crossencoder`).
     :ivar model: The model used for scoring, which could be a `CrossEncoder` or a `CrossEncoderWithLogreg`.
-    :ivar prebuilt_index: Flag indicating whether a prebuilt vector index is used.
     :ivar _db_dir: Path to the database directory where the vector index is stored.
     :ivar name: Name of the scorer, defaults to "dnnc".
 
@@ -95,7 +94,6 @@ class DNNCScorer(ScoringModule):
 
     crossencoder_subdir: str = "crossencoder"
     model: CrossEncoder | CrossEncoderWithLogreg
-    prebuilt_index: bool = False
 
     def __init__(
         self,
@@ -107,7 +105,7 @@ def __init__(
         train_head: bool = False,
         batch_size: int = 32,
         max_length: int | None = None,
-        embedder_use_cache: bool = False,
+        embedder_use_cache: bool = True,
     ) -> None:
         """
         Initialize the DNNCScorer.
@@ -164,11 +162,8 @@ def from_context(
         """
         if embedder_name is None:
             embedder_name = context.optimization_info.get_best_embedder()
-            prebuilt_index = True
-        else:
-            prebuilt_index = context.vector_index_client.exists(embedder_name)
 
-        instance = cls(
+        return cls(
             cross_encoder_name=cross_encoder_name,
             embedder_name=embedder_name,
             k=k,
@@ -179,8 +174,6 @@ def from_context(
             max_length=context.get_max_length(),
             embedder_use_cache=context.get_use_cache(),
         )
-        instance.prebuilt_index = prebuilt_index
-        return instance
 
     def fit(self, utterances: list[str], labels: list[LabelType]) -> None:
         """
@@ -195,15 +188,7 @@ def fit(self, utterances: list[str], labels: list[LabelType]) -> None:
         self.model = CrossEncoder(self.cross_encoder_name, trust_remote_code=True, device=self.device)
 
         vector_index_client = VectorIndexClient(self.device, self.db_dir, embedder_use_cache=self.embedder_use_cache)
-
-        if self.prebuilt_index:
-            # this happens only when LinearScorer is within Pipeline opimization after RetrievalNode optimization
-            self.vector_index = vector_index_client.get_index(self.embedder_name)
-            if len(utterances) != len(self.vector_index.texts):
-                msg = "Vector index mismatches provided utterances"
-                raise ValueError(msg)
-        else:
-            self.vector_index = vector_index_client.create_index(self.embedder_name, utterances, labels)
+        self.vector_index = vector_index_client.create_index(self.embedder_name, utterances, labels)
 
         if self.train_head:
             model = CrossEncoderWithLogreg(self.model)