integrate embeddings fine-tuning into Embedding modules

Author: voorhs

autointent/_wrappers/embedder.py

@@ -10,6 +10,7 @@
 import tempfile
 from functools import lru_cache
 from pathlib import Path
+from uuid import uuid4
 
 import huggingface_hub
 import numpy as np
@@ -26,6 +27,7 @@
 
 from autointent._hash import Hasher
 from autointent.configs import EmbedderConfig, EmbedderFineTuningConfig, TaskTypeEnum
+from autointent.custom_types import ListOfLabels
 
 logger = logging.getLogger(__name__)
 
@@ -72,7 +74,9 @@ class Embedder:
     """
 
     _metadata_dict_name: str = "metadata.json"
+    _weights_dir_name: str = "sentence_transformer"
     _dump_dir: Path | None = None
+    _trained: bool = False
 
     def __init__(self, embedder_config: EmbedderConfig) -> None:
         """Initialize the Embedder.
@@ -89,7 +93,7 @@ def _get_hash(self) -> int:
             The hash value of the Embedder.
         """
         hasher = Hasher()
-        if self.config.freeze:
+        if not Path(self.config.model_name).exists():
             commit_hash = _get_latest_commit_hash(self.config.model_name)
             hasher.update(commit_hash)
         else:
@@ -113,8 +117,22 @@ def _load_model(self) -> SentenceTransformer:
             res = self.embedding_model
         return res
 
-    def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTuningConfig) -> None:
+    def train(self, utterances: list[str], labels: ListOfLabels, config: EmbedderFineTuningConfig) -> None:
         """Train the embedding model."""
+        if len(utterances) != len(labels):
+            msg = f"Utterances and labels lists lengths mismatch: {len(utterances)=} != {len(labels)=}"
+            raise ValueError(msg)
+
+        if len(labels) == 0:
+            msg = "Empty data"
+            raise ValueError(msg)
+
+        # TODO support multi-label data
+        if isinstance(labels[0], list):
+            msg = "Multi-label data is not supported for embeddings fine-tuning for now"
+            logger.warning(msg)
+            return
+
         self._load_model()
         if config.early_stopping:
             x_train, x_val, y_train, y_val = train_test_split(utterances, labels, test_size=0.1, random_state=42)
@@ -131,8 +149,7 @@ def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTu
                 output_dir=tmp_dir,
                 num_train_epochs=config.epoch_num,
                 per_device_train_batch_size=config.batch_size,
-                per_device_eval_batch_size=8,
-                eval_steps=1,
+                per_device_eval_batch_size=config.batch_size,
                 learning_rate=config.learning_rate,
                 warmup_ratio=config.warmup_ratio,
                 fp16=config.fp16,
@@ -143,9 +160,9 @@ def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTu
                 eval_strategy="epoch",
                 greater_is_better=False,
             )
-            callback: list[TrainerCallback] = []
+            callbacks: list[TrainerCallback] = []
             if config.early_stopping:
-                callback.append(
+                callbacks.append(
                     EarlyStoppingCallback(
                         early_stopping_patience=config.early_stopping,
                         early_stopping_threshold=config.early_stopping_threshold,
@@ -157,11 +174,18 @@ def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTu
                 train_dataset=tr_ds,
                 eval_dataset=val_ds,
                 loss=loss,
-                callbacks=callback,
+                callbacks=callbacks,
             )
 
             trainer.train()
 
+        # use temporary path for re-usage
+        model_path = str(Path(tempfile.mkdtemp("autointent_embedders")) / str(uuid4()))
+        self.embedding_model.save(model_path)
+        self.config.model_name = model_path
+
+        self._trained = True
+
     def clear_ram(self) -> None:
         """Move the embedding model to CPU and delete it from memory."""
         if hasattr(self, "embedding_model"):
@@ -182,6 +206,11 @@ def dump(self, path: Path) -> None:
         Args:
             path: Path to the directory where the model will be saved.
         """
+        if self._trained:
+            model_path = str((path / self._weights_dir_name).resolve())
+            self.embedding_model.save(model_path, create_model_card=False)
+            self.config.model_name = model_path
+
         self._dump_dir = path
         path.mkdir(parents=True, exist_ok=True)
         with (path / self._metadata_dict_name).open("w") as file:

autointent/_wrappers/vector_index/vector_index.py

@@ -38,15 +38,15 @@ class VectorIndex:
     embedder: Embedder
     index: BaseIndexBackend
 
-    def __init__(self, embedder_config: EmbedderConfig, config: VectorIndexConfig) -> None:
+    def __init__(self, embedder_config: EmbedderConfig | Embedder, config: VectorIndexConfig) -> None:
         """Initialize the VectorIndex with an embedding model.
 
         Args:
             embedder_config: Configuration for the embedding model.
             config: settings for vector index.
             backend: vector index backend to use.
         """
-        self.embedder = Embedder(embedder_config)
+        self.embedder = embedder_config if isinstance(embedder_config, Embedder) else Embedder(embedder_config)
         self.config = config
 
     def _init_index(self, vector_size: int) -> BaseIndexBackend:

autointent/configs/_transformers.py

@@ -2,7 +2,7 @@
 from typing import Any, Literal
 
 from pydantic import BaseModel, ConfigDict, Field, PositiveInt
-from typing_extensions import Self
+from typing_extensions import Self, assert_never
 
 from autointent.custom_types import FloatFromZeroToOne
 from autointent.metrics import SCORING_METRICS_MULTICLASS, SCORING_METRICS_MULTILABEL
@@ -26,6 +26,16 @@ class EmbedderFineTuningConfig(BaseModel):
     fp16: bool = Field(default=False)
     bf16: bool = Field(default=False)
 
+    @classmethod
+    def from_search_config(cls, values: dict[str, Any] | BaseModel | None) -> Self | None:
+        if isinstance(values, BaseModel):
+            return cls(**values.model_dump())
+        if isinstance(values, dict):
+            return cls(**values)
+        if values is None:
+            return None
+        assert_never(values)
+
 
 class HFModelConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")
@@ -54,7 +64,7 @@ def from_search_config(cls, values: dict[str, Any] | str | BaseModel | None) ->
         if values is None:
             return cls()
         if isinstance(values, BaseModel):
-            return values  # type: ignore[return-value]
+            return cls(**values.model_dump())
         if isinstance(values, str):
             return cls(model_name=values)
         return cls(**values)
@@ -85,7 +95,6 @@ class EmbedderConfig(HFModelConfig):
         "cosine", description="Name of the similarity function to use."
     )
     use_cache: bool = Field(True, description="Whether to use embeddings caching.")
-    freeze: bool = Field(True, description="Whether to freeze the model parameters.")
 
     def get_prompt_config(self) -> dict[str, str] | None:
         """Get the prompt config for the given prompt type.
@@ -174,5 +183,7 @@ def from_search_config(cls, values: dict[str, Any] | BaseModel | None) -> Self:
         if values is None:
             return cls()
         if isinstance(values, BaseModel):
-            return values  # type: ignore[return-value]
-        return cls(**values)
+            return cls(**values.model_dump())
+        if isinstance(values, dict):
+            return cls(**values)
+        assert_never(values)

autointent/context/optimization_info/_optimization_info.py

@@ -20,7 +20,7 @@
 from autointent.configs import EmbedderConfig, InferenceNodeConfig
 from autointent.custom_types import NodeType
 
-from ._data_models import Artifact, Artifacts, EmbeddingArtifact, ScorerArtifact, Trial, Trials
+from ._data_models import Artifacts, EmbeddingArtifact, ScorerArtifact, Trial, Trials
 
 if TYPE_CHECKING:
     from autointent.modules.base import BaseModule
@@ -95,7 +95,6 @@ def log_module_optimization(
         metric_value: float,
         metric_name: str,
         metrics: dict[str, float],
-        artifact: Artifact,
         module_dump_dir: str | None,
         module: "BaseModule",
     ) -> None:
@@ -108,7 +107,6 @@ def log_module_optimization(
             metric_value: Metric value achieved by the module.
             metric_name: Name of the evaluation metric.
             metrics: Dictionary of metric names and their values.
-            artifact: Artifact generated by the module.
             module_dump_dir: Directory where the module is dumped.
             module: The module instance, if available.
         """
@@ -117,7 +115,7 @@ def log_module_optimization(
             self.modules.add_module(node_type, module)
             if module_dump_dir is not None:
                 module.dump(module_dump_dir)
-            self.artifacts.add_artifact(node_type, artifact)
+            self.artifacts.add_artifact(node_type, module.get_assets())
 
             if old_best_metric_value_idx is not None:
                 prev_best_dump = self.trials.get_trials(node_type)[old_best_metric_value_idx].module_dump_dir

autointent/modules/embedding/_logreg.py

@@ -10,7 +10,7 @@
 from sklearn.preprocessing import LabelEncoder
 
 from autointent import Context, Embedder
-from autointent.configs import EmbedderConfig, TaskTypeEnum
+from autointent.configs import EmbedderConfig, EmbedderFineTuningConfig, TaskTypeEnum
 from autointent.context.optimization_info import EmbeddingArtifact
 from autointent.custom_types import ListOfLabels
 from autointent.metrics import SCORING_METRICS_MULTICLASS, SCORING_METRICS_MULTILABEL
@@ -26,6 +26,7 @@ class LogregAimedEmbedding(BaseEmbedding):
     Args:
         embedder_config: Config of the embedder used for creating embeddings
         cv: Number of folds used in LogisticRegressionCV
+        ft_config: settings for fine-tuning embeddings
 
     Examples:
     --------
@@ -52,9 +53,11 @@ def __init__(
         self,
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
         cv: PositiveInt = 3,
+        ft_config: EmbedderFineTuningConfig | dict[str, Any] | None = None,
     ) -> None:
-        self.embedder_config = EmbedderConfig.from_search_config(embedder_config)
+        self._embedder = Embedder(EmbedderConfig.from_search_config(embedder_config))
         self.cv = cv
+        self.ft_config = EmbedderFineTuningConfig.from_search_config(ft_config)
 
         if self.cv < 0 or not isinstance(self.cv, int):
             msg = "`cv` argument of `LogregAimedEmbedding` must be a positive int"
@@ -65,6 +68,7 @@ def from_context(
         cls,
         context: Context,
         embedder_config: EmbedderConfig | str | None = None,
+        ft_config: EmbedderFineTuningConfig | dict[str, Any] | None = None,
         cv: PositiveInt = 3,
     ) -> "LogregAimedEmbedding":
         """Create a LogregAimedEmbedding instance using a Context object.
@@ -73,10 +77,12 @@ def from_context(
             context: Context containing configurations and utilities
             cv: Number of folds used in LogisticRegressionCV
             embedder_config: Config of the embedder to use
+            ft_config: settings for fine-tuning embeddings
         """
         return cls(
             cv=cv,
             embedder_config=embedder_config,
+            ft_config=ft_config,
         )
 
     def clear_cache(self) -> None:
@@ -93,9 +99,9 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         """
         self._validate_task(labels)
 
-        self._embedder = Embedder(
-            self.embedder_config,
-        )
+        if self.ft_config is not None:
+            self._embedder.train(utterances=utterances, labels=labels, config=self.ft_config)
+
         embeddings = self._embedder.embed(utterances, TaskTypeEnum.classification)
 
         if self._multilabel:
@@ -153,7 +159,7 @@ def get_assets(self) -> EmbeddingArtifact:
         Returns:
             EmbeddingArtifact object containing embedder information
         """
-        return EmbeddingArtifact(config=self.embedder_config)
+        return EmbeddingArtifact(config=self._embedder.config)
 
     def predict(self, utterances: list[str]) -> NDArray[np.float64]:
         """Predict probabilities for input utterances.

autointent/modules/embedding/_retrieval.py

@@ -4,8 +4,13 @@
 
 from pydantic import PositiveInt
 
-from autointent import Context, VectorIndex
-from autointent.configs import EmbedderConfig, VectorIndexConfig, get_default_vector_index_config
+from autointent import Context, Embedder, VectorIndex
+from autointent.configs import (
+    EmbedderConfig,
+    EmbedderFineTuningConfig,
+    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
@@ -21,6 +26,7 @@ class RetrievalAimedEmbedding(BaseEmbedding):
     Args:
         k: Number of nearest neighbors to retrieve
         embedder_config: Config of the embedder used for creating embeddings
+        ft_config: settings for fine-tuning embeddings
 
     Examples:
     --------
@@ -49,11 +55,12 @@ def __init__(
         embedder_config: EmbedderConfig | str | dict[str, Any] | None = None,
         vector_index_config: VectorIndexConfig | None = None,
         k: PositiveInt = 10,
+        ft_config: EmbedderFineTuningConfig | dict[str, Any] | None = None,
     ) -> None:
         self.k = k
-        embedder_config = EmbedderConfig.from_search_config(embedder_config)
-        self.embedder_config = embedder_config
+        self._embedder = Embedder(EmbedderConfig.from_search_config(embedder_config))
         self.vector_index_config = vector_index_config or get_default_vector_index_config()
+        self.ft_config = EmbedderFineTuningConfig.from_search_config(ft_config)
 
         if self.k < 0 or not isinstance(self.k, int):
             msg = "`k` argument of `RetrievalAimedEmbedding` must be a positive int"
@@ -65,18 +72,21 @@ def from_context(
         context: Context,
         embedder_config: EmbedderConfig | str | None = None,
         k: PositiveInt = 10,
+        ft_config: EmbedderFineTuningConfig | dict[str, Any] | None = None,
     ) -> "RetrievalAimedEmbedding":
         """Create an instance using a Context object.
 
         Args:
             context: The context containing configurations and utilities
             k: Number of nearest neighbors to retrieve
             embedder_config: Config of the embedder to use
+            ft_config: settings for fine-tuning embeddings
         """
         return cls(
             k=k,
             embedder_config=embedder_config,
             vector_index_config=context.vector_index_config,
+            ft_config=ft_config,
         )
 
     def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
@@ -88,7 +98,10 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         """
         self._validate_task(labels)
 
-        self._vector_index = VectorIndex(self.embedder_config, config=self.vector_index_config)
+        if self.ft_config is not None:
+            self._embedder.train(utterances=utterances, labels=labels, config=self.ft_config)
+
+        self._vector_index = VectorIndex(self._embedder, config=self.vector_index_config)
         self._vector_index.add(utterances, labels)
 
     def score_ho(self, context: Context, metrics: list[str]) -> dict[str, float]:
@@ -134,7 +147,7 @@ def get_assets(self) -> EmbeddingArtifact:
         Returns:
             A EmbeddingArtifact object containing embedder information
         """
-        return EmbeddingArtifact(config=self.embedder_config)
+        return EmbeddingArtifact(config=self._embedder.config)
 
     def clear_cache(self) -> None:
         """Clear cached data in memory used by the vector index."""

autointent/nodes/_node_optimizer.py

@@ -156,7 +156,6 @@ def objective(
             metric_value=target_metric,
             metric_name=self.target_metric,
             metrics=quality_metrics,
-            artifact=module.get_assets(),  # retriever name / scores / predictions
             module_dump_dir=self.get_module_dump_dir(context, module_name, self._counter),
             module=module,
         )