feat/train-embeddings

autointent/_wrappers/embedder.py

@@ -10,17 +10,22 @@
 from functools import lru_cache
 from pathlib import Path
 from typing import TypedDict
+import tempfile
 
 import huggingface_hub
 import numpy as np
 import numpy.typing as npt
 import torch
 from appdirs import user_cache_dir
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments
 from sentence_transformers.similarity_functions import SimilarityFunction
+from sentence_transformers.losses import BatchAllTripletLoss
+from sentence_transformers.training_args import BatchSamplers
+from datasets import Dataset
+
 
 from autointent._hash import Hasher
-from autointent.configs import EmbedderConfig, TaskTypeEnum
+from autointent.configs import EmbedderConfig, TaskTypeEnum, EmbedderFineTuningConfig
 
 logger = logging.getLogger(__name__)
 
@@ -122,7 +127,41 @@ def _load_model(self) -> None:
                 similarity_fn_name=self.config.similarity_fn_name,
                 trust_remote_code=self.config.trust_remote_code,
             )
+    def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTuningConfig) -> None:
+        """Train the embedding model"""
+        self._load_model()
 
+        tr_ds = Dataset.from_dict({
+        "text": utterances,
+        "label": labels
+        })
+
+        loss = BatchAllTripletLoss(
+            model=self.embedding_model, 
+            margin=config.margin
+        )
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            args = SentenceTransformerTrainingArguments(
+                save_strategy="no",
+                output_dir=tmp_dir,
+                num_train_epochs=config.epoch_num,
+                per_device_train_batch_size=self.config.batch_size,
+                learning_rate=config.learning_rate,
+                warmup_ratio=config.warmup_ratio,
+                fp16=config.fp16,
+                bf16=config.bf16,
+                batch_sampler=BatchSamplers.NO_DUPLICATES,
+            )
+
+            trainer = SentenceTransformerTrainer(
+                model=self.embedding_model,
+                args=args,
+                train_dataset=tr_ds,
+                loss=loss,
+            )
+
+            trainer.train()
+
     def clear_ram(self) -> None:
         """Move the embedding model to CPU and delete it from memory."""
         if hasattr(self, "embedding_model"):

autointent/configs/_transformers.py

@@ -14,6 +14,13 @@ class TokenizerConfig(BaseModel):
     truncation: bool = True
     max_length: PositiveInt | None = Field(None, description="Maximum length of input sequences.")
 
+class EmbedderFineTuningConfig(BaseModel):
+    epoch_num: int
+    margin: float = Field(default=0.5)
+    learning_rate: float = Field(default=2e-5)
+    warmup_ratio: float = Field(default=0.1)
+    fp16: bool = Field(default=True)
+    bf16: bool = Field(default=False)
 
 class HFModelConfig(BaseModel):
     model_config = ConfigDict(extra="forbid")

pyproject.toml

@@ -29,7 +29,7 @@ classifiers=[
     'Framework :: Sphinx',
     'Typing :: Typed',
 ]
-requires-python = ">=3.10,<4.0"
+requires-python = ">=3.10,<3.13"
 dependencies = [
     "sentence-transformers (>=3,<4)",
     "scikit-learn (>=1.5,<2.0)",

tests/embedder/test_fine_tuning.py

@@ -0,0 +1,55 @@
+from autointent.context.data_handler import DataHandler
+from autointent._wrappers.embedder import Embedder
+from autointent.configs._transformers import HFModelConfig, EmbedderConfig, EmbedderFineTuningConfig
+import numpy as np
+import pytest
+
+def test_model_updates_after_training(dataset):
+    """Test that model weights actually change after training"""
+    data_handler = DataHandler(dataset)
+
+    hf_config = HFModelConfig(
+    model_name="intfloat/multilingual-e5-small",
+    batch_size=8,
+    trust_remote_code=True
+    )
+
+    embedder_config = EmbedderConfig(
+        **hf_config.model_dump(),
+        default_prompt="Represent this text for retrieval:",
+        query_prompt="Search query:",
+        passage_prompt="Document:",
+        similarity_fn_name="cosine",
+        use_cache=False,
+        freeze=False
+    )
+
+    train_config = EmbedderFineTuningConfig(
+        epoch_num = 1
+    )
+    embedder = Embedder(embedder_config)
+    embedder._load_model()
+
+    original_weights = [
+        param.data.detach().cpu().numpy().copy()
+        for param in embedder.embedding_model.parameters()
+        if param.requires_grad
+    ]
+    embedder.train(
+        utterances=data_handler.train_utterances(0)[:10],
+        labels=data_handler.train_labels(0)[:10],
+        config=train_config
+    )
+
+    trained_weights = [
+        param.data.detach().cpu().numpy().copy()
+        for param in embedder.embedding_model.parameters()
+        if param.requires_grad
+    ]
+
+    weights_changed = any(
+        not np.allclose(orig, trained, atol=1e-6)
+        for orig, trained in zip(original_weights, trained_weights)
+    )
+
+    assert weights_changed, "Model weights should change after training"