correct some syntax isues

Author: k0lenk4

autointent/_wrappers/embedder.py

@@ -7,25 +7,24 @@
 import json
 import logging
 import shutil
+import tempfile
 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 datasets import Dataset
 from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, SentenceTransformerTrainingArguments
-from sentence_transformers.similarity_functions import SimilarityFunction
 from sentence_transformers.losses import BatchAllTripletLoss
+from sentence_transformers.similarity_functions import SimilarityFunction
 from sentence_transformers.training_args import BatchSamplers
-from datasets import Dataset
-
 
 from autointent._hash import Hasher
-from autointent.configs import EmbedderConfig, TaskTypeEnum, EmbedderFineTuningConfig
+from autointent.configs import EmbedderConfig, EmbedderFineTuningConfig, TaskTypeEnum
 
 logger = logging.getLogger(__name__)
 
@@ -128,7 +127,7 @@ def _load_model(self) -> None:
                 trust_remote_code=self.config.trust_remote_code,
             )
     def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTuningConfig) -> None:
-        """Train the embedding model"""
+        """Train the embedding model."""
         self._load_model()
 
         tr_ds = Dataset.from_dict({
@@ -137,7 +136,7 @@ def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTu
         })
 
         loss = BatchAllTripletLoss(
-            model=self.embedding_model, 
+            model=self.embedding_model,
             margin=config.margin
         )
         with tempfile.TemporaryDirectory() as tmp_dir:
@@ -159,9 +158,9 @@ def train(self, utterances: list[str], labels: list[int], config: EmbedderFineTu
                 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/__init__.py

@@ -6,8 +6,8 @@
 from ._transformers import (
     CrossEncoderConfig,
     EarlyStoppingConfig,
-    EmbedderFineTuningConfig,
     EmbedderConfig,
+    EmbedderFineTuningConfig,
     HFModelConfig,
     TaskTypeEnum,
     TokenizerConfig,
@@ -17,8 +17,8 @@
     "CrossEncoderConfig",
     "DataConfig",
     "EarlyStoppingConfig",
-    "EmbedderFineTuningConfig",
     "EmbedderConfig",
+    "EmbedderFineTuningConfig",
     "HFModelConfig",
     "HPOConfig",
     "InferenceNodeConfig",

tests/embedder/test_fine_tuning.py

@@ -1,8 +1,10 @@
-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
 
+from autointent._wrappers.embedder import Embedder
+from autointent.configs._transformers import EmbedderConfig, EmbedderFineTuningConfig, HFModelConfig
+from autointent.context.data_handler import DataHandler
+
+
 def test_model_updates_after_training(dataset):
     """Test that model weights actually change after training"""
     data_handler = DataHandler(dataset)
@@ -48,7 +50,6 @@ def test_model_updates_after_training(dataset):
 
     weights_changed = any(
         not np.allclose(orig, trained, atol=1e-6)
-        for orig, trained in zip(original_weights, trained_weights)
+        for orig, trained in zip(original_weights, trained_weights, strict=True)
     )
-    
     assert weights_changed, "Model weights should change after training"