move batch_size to EmbedderFineTuningConfig

Author: k0lenk4

autointent/_wrappers/embedder.py

@@ -126,25 +126,20 @@ 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
-        })
+        tr_ds = Dataset.from_dict({"text": utterances, "label": labels})
 
-        loss = BatchAllTripletLoss(
-            model=self.embedding_model,
-            margin=config.margin
-        )
+        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,
+                per_device_train_batch_size=config.batch_size,
                 learning_rate=config.learning_rate,
                 warmup_ratio=config.warmup_ratio,
                 fp16=config.fp16,

autointent/configs/_transformers.py

@@ -14,14 +14,17 @@ class TokenizerConfig(BaseModel):
     truncation: bool = True
     max_length: PositiveInt | None = Field(None, description="Maximum length of input sequences.")
 
+
 class EmbedderFineTuningConfig(BaseModel):
     epoch_num: int
+    batch_size: 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")
     model_name: str = Field(

tests/embedder/test_fine_tuning.py

@@ -9,11 +9,7 @@ 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
-    )
+    hf_config = HFModelConfig(model_name="intfloat/multilingual-e5-small", batch_size=8, trust_remote_code=True)
 
     embedder_config = EmbedderConfig(
         **hf_config.model_dump(),
@@ -22,12 +18,10 @@ def test_model_updates_after_training(dataset):
         passage_prompt="Document:",
         similarity_fn_name="cosine",
         use_cache=False,
-        freeze=False
+        freeze=False,
     )
 
-    train_config = EmbedderFineTuningConfig(
-        epoch_num = 1
-    )
+    train_config = EmbedderFineTuningConfig(epoch_num=1)
     embedder = Embedder(embedder_config)
     embedder._load_model()
 
@@ -37,9 +31,7 @@ def test_model_updates_after_training(dataset):
         if param.requires_grad
     ]
     embedder.train(
-        utterances=data_handler.train_utterances(0)[:10],
-        labels=data_handler.train_labels(0)[:10],
-        config=train_config
+        utterances=data_handler.train_utterances(0)[:10], labels=data_handler.train_labels(0)[:10], config=train_config
     )
 
     trained_weights = [