deleted kwargs and local savings, added config

Author: k0lenk4

autointent/_wrappers/embedder.py

@@ -10,21 +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, SentenceTransformerTrainer, SentenceTransformerTrainingArguments, InputExample
+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__)
 
@@ -126,7 +127,7 @@ 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], **kwargs) -> None:
+    def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTuningConfig) -> None:
         """Train the embedding model"""
         self._load_model()
 
@@ -137,31 +138,29 @@ def train(self, utterances: list[str], labels: list[int], **kwargs) -> None:
 
         loss = BatchAllTripletLoss(
             model=self.embedding_model, 
-            margin=kwargs.get("margin", 0.5)
-        )
-        
-        args = SentenceTransformerTrainingArguments(
-            save_strategy="no",
-            output_dir=kwargs['out_dir'],
-            num_train_epochs=kwargs['epoch_num'],
-            per_device_train_batch_size=self.config.batch_size,
-            learning_rate=kwargs.get("learning_rate", 2e-5),
-            warmup_ratio=kwargs.get("warmup_ratio", 0.1),
-            fp16=kwargs.get("fp16", True),
-            bf16=kwargs.get("bf16", False),
-            batch_sampler=BatchSamplers.NO_DUPLICATES,
+            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()
-    
-        self.embedding_model.save(kwargs['out_dir'])
+            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."""

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")