work on review

Author: SeBorgey

autointent/modules/__init__.py

@@ -13,14 +13,14 @@
 from .embedding import LogregAimedEmbedding, RetrievalAimedEmbedding
 from .regex import SimpleRegex
 from .scoring import (
+    BertScorer,
     DescriptionScorer,
     DNNCScorer,
     KNNScorer,
     LinearScorer,
     MLKnnScorer,
     RerankScorer,
     SklearnScorer,
-    TransformerScorer,
 )
 
 T = TypeVar("T", bound=BaseModule)
@@ -45,7 +45,7 @@ def _create_modules_dict(modules: list[type[T]]) -> dict[str, type[T]]:
         RerankScorer,
         SklearnScorer,
         MLKnnScorer,
-        TransformerScorer,
+        BertScorer,
     ]
 )
 

autointent/modules/scoring/__init__.py

@@ -4,15 +4,15 @@
 from ._linear import LinearScorer
 from ._mlknn import MLKnnScorer
 from ._sklearn import SklearnScorer
-from ._transformer import TransformerScorer
+from ._transformer import BertScorer
 
 __all__ = [
+    "BertScorer",
     "DNNCScorer",
     "DescriptionScorer",
     "KNNScorer",
     "LinearScorer",
     "MLKnnScorer",
     "RerankScorer",
     "SklearnScorer",
-    "TransformerScorer"
 ]

autointent/modules/scoring/_transformer.py

@@ -1,4 +1,4 @@
-"""TransformerScorer class for transformer-based classification."""
+"""BertScorer class for transformer-based classification."""
 
 import tempfile
 from typing import Any
@@ -21,7 +21,21 @@
 from autointent.modules.base import BaseScorer
 
 
-class TransformerScorer(BaseScorer):
+class TokenizerConfig:
+    """Configuration for tokenizer parameters."""
+
+    def __init__(
+        self,
+        max_length: int = 128,
+        padding: str = "max_length",
+        truncation: bool = True,
+    ) -> None:
+        self.max_length = max_length
+        self.padding = padding
+        self.truncation = truncation
+
+
+class BertScorer(BaseScorer):
     name = "transformer"
     supports_multiclass = True
     supports_multilabel = True
@@ -36,26 +50,46 @@ def __init__(
         batch_size: int = 8,
         learning_rate: float = 5e-5,
         seed: int = 0,
+        tokenizer_config: TokenizerConfig | None = None,
     ) -> None:
         self.model_config = EmbedderConfig.from_search_config(model_config)
         self.num_train_epochs = num_train_epochs
         self.batch_size = batch_size
         self.learning_rate = learning_rate
         self.seed = seed
+        self.tokenizer_config = tokenizer_config or TokenizerConfig()
+        self._multilabel = False
 
     @classmethod
     def from_context(
         cls,
         context: Context,
         model_config: EmbedderConfig | str | None = None,
-    ) -> "TransformerScorer":
+        num_train_epochs: int = 3,
+        batch_size: int = 8,
+        learning_rate: float = 5e-5,
+        seed: int = 0,
+        tokenizer_config: TokenizerConfig | None = None,
+    ) -> "BertScorer":
         if model_config is None:
             model_config = context.resolve_embedder()
-        return cls(model_config=model_config)
+        return cls(
+            model_config=model_config,
+            num_train_epochs=num_train_epochs,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            seed=seed,
+            tokenizer_config=tokenizer_config,
+        )
 
     def get_embedder_config(self) -> dict[str, Any]:
         return self.model_config.model_dump()
 
+    def _validate_task(self, labels: ListOfLabels) -> None:
+        """Validate the task and set _multilabel flag."""
+        super()._validate_task(labels)
+        self._multilabel = isinstance(labels[0], list)
+
     def fit(
         self,
         utterances: list[str],
@@ -67,7 +101,7 @@ def fit(
         self._validate_task(labels)
 
         if self._multilabel:
-            labels_array = np.array(labels) if not isinstance(labels, np.ndarray) else labels
+            labels_array = np.array(labels)
             num_labels = labels_array.shape[1]
         else:
             num_labels = len(set(labels))
@@ -76,8 +110,15 @@ def fit(
         self._tokenizer = AutoTokenizer.from_pretrained(model_name)
         self._model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
 
-        def tokenize_function(examples: dict[str, Any]) -> dict[str, Any]:
-            return self._tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
+        use_cpu = hasattr(self.model_config, "device") and self.model_config.device == "cpu"
+
+        def tokenize_function(examples: dict[str, Any]) -> dict[str, Any]:  # type: ignore[no-any-return]
+            return self._tokenizer(
+                examples["text"],
+                padding=self.tokenizer_config.padding,
+                truncation=self.tokenizer_config.truncation,
+                max_length=self.tokenizer_config.max_length,
+            )
 
         dataset = Dataset.from_dict({"text": utterances, "labels": labels})
         tokenized_dataset = dataset.map(tokenize_function, batched=True)
@@ -90,8 +131,10 @@ def tokenize_function(examples: dict[str, Any]) -> dict[str, Any]:
                 learning_rate=self.learning_rate,
                 seed=self.seed,
                 save_strategy="no",
-                logging_strategy="no",
-                report_to="none",
+                logging_strategy="steps",
+                logging_steps=10,
+                report_to="wandb",
+                use_cpu=use_cpu,
             )
 
             trainer = Trainer(
@@ -111,7 +154,9 @@ def predict(self, utterances: list[str]) -> npt.NDArray[Any]:
             msg = "Model is not trained. Call fit() first."
             raise RuntimeError(msg)
 
-        inputs = self._tokenizer(utterances, padding=True, truncation=True, max_length=128, return_tensors="pt")
+        inputs = self._tokenizer(
+            utterances, padding=True, truncation=True, max_length=self.tokenizer_config.max_length, return_tensors="pt"
+        )
 
         with torch.no_grad():
             outputs = self._model(**inputs)
@@ -121,7 +166,6 @@ def predict(self, utterances: list[str]) -> npt.NDArray[Any]:
             return torch.sigmoid(logits).numpy()
         return torch.softmax(logits, dim=1).numpy()
 
-
     def clear_cache(self) -> None:
         if hasattr(self, "_model"):
             del self._model

tests/modules/scoring/test_transformer.py

@@ -2,13 +2,14 @@
 import pytest
 
 from autointent.context.data_handler import DataHandler
-from autointent.modules import TransformerScorer
+from autointent.modules import BertScorer
 
 
-def test_base_transformer(dataset):
+def test_bert_prediction(dataset):
+    """Test that the transformer model can fit and make predictions."""
     data_handler = DataHandler(dataset)
 
-    scorer = TransformerScorer(model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
+    scorer = BertScorer(model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
 
     scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
 
@@ -22,25 +23,45 @@ def test_base_transformer(dataset):
 
     predictions = scorer.predict(test_data)
 
+    # Verify prediction shape
     assert predictions.shape[0] == len(test_data)
     assert predictions.shape[1] == len(set(data_handler.train_labels(0)))
 
+    # Verify predictions are probabilities
     assert 0.0 <= np.min(predictions) <= np.max(predictions) <= 1.0
 
+    # Verify probabilities sum to 1 for multiclass
     if not scorer._multilabel:
         for pred_row in predictions:
             np.testing.assert_almost_equal(np.sum(pred_row), 1.0, decimal=5)
 
+    # Test metadata function if available
     if hasattr(scorer, "predict_with_metadata"):
         predictions, metadata = scorer.predict_with_metadata(test_data)
         assert len(predictions) == len(test_data)
         assert metadata is None
-    else:
-        pytest.skip("predict_with_metadata not implemented in TransformerScorer")
 
+
+def test_bert_cache_clearing(dataset):
+    """Test that the transformer model properly handles cache clearing."""
+    data_handler = DataHandler(dataset)
+
+    scorer = BertScorer(model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
+
+    scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
+
+    test_data = ["test text"]
+
+    # Should work before clearing cache
+    scorer.predict(test_data)
+
+    # Clear the cache
     scorer.clear_cache()
+
+    # Verify model and tokenizer are removed
     assert not hasattr(scorer, "_model") or scorer._model is None
     assert not hasattr(scorer, "_tokenizer") or scorer._tokenizer is None
 
+    # Should raise exception after clearing cache
     with pytest.raises(RuntimeError):
         scorer.predict(test_data)