rename bert model

Author: Samoed

autointent/modules/scoring/_bert.py

@@ -31,14 +31,14 @@ class BertScorer(BaseScorer):
 
     def __init__(
         self,
-        model_config: HFModelConfig | str | dict[str, Any] | None = None,
+        classification_model_config: HFModelConfig | str | dict[str, Any] | None = None,
         num_train_epochs: int = 3,
         batch_size: int = 8,
         learning_rate: float = 5e-5,
         seed: int = 0,
         report_to: REPORTERS_NAMES | None = None,  # type: ignore  # noqa: PGH003
     ) -> None:
-        self.model_config = HFModelConfig.from_search_config(model_config)
+        self.classification_model_config = HFModelConfig.from_search_config(classification_model_config)
         self.num_train_epochs = num_train_epochs
         self.batch_size = batch_size
         self.learning_rate = learning_rate
@@ -49,19 +49,19 @@ def __init__(
     def from_context(
         cls,
         context: Context,
-        model_config: HFModelConfig | str | dict[str, Any] | None = None,
+        classification_model_config: HFModelConfig | str | dict[str, Any] | None = None,
         num_train_epochs: int = 3,
         batch_size: int = 8,
         learning_rate: float = 5e-5,
         seed: int = 0,
     ) -> "BertScorer":
-        if model_config is None:
-            model_config = context.resolve_embedder()
+        if classification_model_config is None:
+            classification_model_config = context.resolve_embedder()
 
         report_to = context.logging_config.report_to
 
         return cls(
-            model_config=model_config,
+            classification_model_config=classification_model_config,
             num_train_epochs=num_train_epochs,
             batch_size=batch_size,
             learning_rate=learning_rate,
@@ -70,7 +70,7 @@ def from_context(
         )
 
     def get_embedder_config(self) -> dict[str, Any]:
-        return self.model_config.model_dump()
+        return self.classification_model_config.model_dump()
 
     def fit(
         self,
@@ -81,7 +81,7 @@ def fit(
             self.clear_cache()
         self._validate_task(labels)
 
-        model_name = self.model_config.model_name
+        model_name = self.classification_model_config.model_name
         self._tokenizer = AutoTokenizer.from_pretrained(model_name)
 
         label2id = {i: i for i in range(self._n_classes)}
@@ -95,11 +95,11 @@ def fit(
             problem_type="multi_label_classification" if self._multilabel else "single_label_classification",
         )
 
-        use_cpu = self.model_config.device == "cpu"
+        use_cpu = self.classification_model_config.device == "cpu"
 
         def tokenize_function(examples: dict[str, Any]) -> dict[str, Any]:
             return self._tokenizer(  # type: ignore[no-any-return]
-                examples["text"], return_tensors="pt", **self.model_config.tokenizer_config.model_dump()
+                examples["text"], return_tensors="pt", **self.classification_model_config.tokenizer_config.model_dump()
             )
 
         dataset = Dataset.from_dict({"text": utterances, "labels": labels})
@@ -148,7 +148,7 @@ def predict(self, utterances: list[str]) -> npt.NDArray[Any]:
         all_predictions = []
         for i in range(0, len(utterances), self.batch_size):
             batch = utterances[i : i + self.batch_size]
-            inputs = self._tokenizer(batch, return_tensors="pt", **self.model_config.tokenizer_config.model_dump())
+            inputs = self._tokenizer(batch, return_tensors="pt", **self.classification_model_config.tokenizer_config.model_dump())
             inputs = {k: v.to(device) for k, v in inputs.items()}
             with torch.no_grad():
                 outputs = self._model(**inputs)

tests/assets/configs/multiclass.yaml

@@ -29,7 +29,7 @@
       clf_name: [RandomForestClassifier]
       n_estimators: [5, 10]
     - module_name: bert
-      model_config:
+      classification_model_config:
         - model_name: avsolatorio/GIST-small-Embedding-v0
       num_train_epochs: [1]
       batch_size: [8, 16]

tests/assets/configs/multilabel.yaml

@@ -25,7 +25,7 @@
       clf_name: [RandomForestClassifier]
       n_estimators: [5, 10]
     - module_name: bert
-      model_config:
+      classification_model_config:
         - model_name: avsolatorio/GIST-small-Embedding-v0
       num_train_epochs: [1]
       batch_size: [8]

tests/modules/scoring/test_bert.py

@@ -9,7 +9,7 @@ def test_bert_prediction(dataset):
     """Test that the transformer model can fit and make predictions."""
     data_handler = DataHandler(dataset)
 
-    scorer = BertScorer(model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
+    scorer = BertScorer(classification_model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
 
     scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))
 
@@ -46,7 +46,7 @@ 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 = BertScorer(classification_model_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=8)
 
     scorer.fit(data_handler.train_utterances(0), data_handler.train_labels(0))