Refactor model and trainer kwargs in SetFitClassification

Signed-off-by: Christopher Schröder <chschroeder@users.noreply.github.com>

Author: chschroeder

CHANGELOG.md

@@ -27,6 +27,9 @@ On the other hand, this also allowed us to deal with further issues that contain
   - TransformerBasedClassification:
     - Removed unnecessary `token_type_ids` keyword argument in model call.
     - Additional keyword args for config, tokenizer, and model can now be configured.
+  - SetFitClassification:
+    - Additional keyword args for trainer and model are now attached to `SetFitModelArguments` instead of `SetFitClassification`.
+
 - Embeddings:
   - Prevented unnecessary gradient computations for some embedding types and unified code structure.
 - Pytorch:

MIGRATION_GUIDE.md

@@ -11,9 +11,11 @@ This is not an exhaustive list of changes, but here we try to collect changes th
 - [PoolBasedActiveLearner](https://small-text.readthedocs.io/en/latest/api/active_learner.html#activelearner-api):  
   `initialize_data()` has been changed to `initialize()`. The method now takes a list of initial indices or an initialized first (proxy-)model.
 
+- SetFitClassification: `model_kwargs` and `trainer_kwargs` are now attached to `SetFitModelArguments` instead of `SetFitClassification`.
+
 ### Renamed Classes
 
-The following classes amd variables have been renamed for consistency:
+The following classes and variables have been renamed for consistency:
 
 - KimCNNFactory -> KimCNNClassifierFactory
 

small_text/integrations/transformers/classifiers/setfit.py

@@ -19,6 +19,13 @@
     ModelLoadingStrategy
 )
 
+from small_text.integrations.transformers.utils.setfit import (
+    _check_model_kwargs,
+    _check_trainer_kwargs,
+    _check_train_kwargs,
+    _truncate_texts
+)
+
 try:
     import torch
 
@@ -29,34 +36,54 @@
     from small_text.integrations.transformers.utils.classification import (
         _get_arguments_for_from_pretrained_model
     )
-    from small_text.integrations.transformers.utils.setfit import (
-        _check_model_kwargs,
-        _check_trainer_kwargs,
-        _check_train_kwargs,
-        _truncate_texts
-    )
-except ImportError:
-    pass
+except ImportError as e:
+    print(e)
+    print(e)
 
 
 class SetFitModelArguments(object):
-    """
+    """Model arguments for :py:class:`SetFitClassification`.
+
     .. versionadded:: 1.2.0
     """
 
     def __init__(self,
                  sentence_transformer_model: str,
+                 model_kwargs={},
+                 trainer_kwargs={},
                  model_loading_strategy: ModelLoadingStrategy = ModelLoadingStrategy.DEFAULT):
         """
         Parameters
         ----------
         sentence_transformer_model : str
             Name of a sentence transformer model.
+        model_kwargs : dict, default={}
+            Keyword arguments used for the SetFit model. The keyword `use_differentiable_head` is
+            excluded and managed by this class. The other keywords are directly passed to
+            `SetFitModel.from_pretrained()`. Additional kwargs that will be passed into
+            `SetFitModel.from_pretrained()`. Arguments that are managed by small-text
+            (such as the model name given by `model`) are excluded.
+
+            .. seealso::
+
+                `SetFitModel.from_pretrained()
+                <https://huggingface.co/docs/setfit/en/reference/main#setfit.SetFitModel.from_pretrained>`_
+                in the SetFit documentation.
+        trainer_kwargs : dict
+            Keyword arguments used for the SetFit model. The keyword `batch_size` is excluded and
+            is instead controlled by the keyword `mini_batch_size` of this class. The other
+            keywords are directly passed to `SetFitTrainer.__init__()`.
+
+            .. seealso:: `Trainer
+                         <https://huggingface.co/docs/setfit/en/reference/trainer>`_
+                         in the SetFit documentation.
         model_loading_strategy: ModelLoadingStrategy, default=ModelLoadingStrategy.DEFAULT
             Specifies if there should be attempts to download the model or if only local
             files should be used.
         """
         self.sentence_transformer_model = sentence_transformer_model
+        self.model_kwargs = _check_model_kwargs(model_kwargs)
+        self.trainer_kwargs = _check_trainer_kwargs(trainer_kwargs)
         self.model_loading_strategy = model_loading_strategy
 
 
@@ -135,8 +162,7 @@ class SetFitClassification(SetFitClassificationEmbeddingMixin, Classifier):
     """
 
     def __init__(self, setfit_model_args, num_classes, multi_label=False, max_seq_len=512,
-                 use_differentiable_head=False, mini_batch_size=32, model_kwargs=dict(),
-                 trainer_kwargs=dict(), device=None):
+                 use_differentiable_head=False, mini_batch_size=32, device=None):
         """
         sentence_transformer_model : SetFitModelArguments
             Settings for the sentence transformer model to be used.
@@ -149,21 +175,6 @@ def __init__(self, setfit_model_args, num_classes, multi_label=False, max_seq_le
             Uses a differentiable head instead of a logistic regression for the classification head.
             Corresponds to the keyword argument with the same name in
             `SetFitModel.from_pretrained()`.
-        model_kwargs : dict
-            Keyword arguments used for the SetFit model. The keyword `use_differentiable_head` is
-            excluded and managed by this class. The other keywords are directly passed to
-            `SetFitModel.from_pretrained()`.
-
-            .. seealso:: `SetFit: src/setfit/modeling.py
-                         <https://github.com/huggingface/setfit/blob/main/src/setfit/modeling.py>`_
-
-        trainer_kwargs : dict
-            Keyword arguments used for the SetFit model. The keyword `batch_size` is excluded and
-            is instead controlled by the keyword `mini_batch_size` of this class. The other
-            keywords are directly passed to `SetFitTrainer.__init__()`.
-
-            .. seealso:: `SetFit: src/setfit/trainer.py
-                         <https://github.com/huggingface/setfit/blob/main/src/setfit/trainer.py>`_
         device : str or torch.device, default=None
             Torch device on which the computation will be performed.
         """
@@ -173,10 +184,7 @@ def __init__(self, setfit_model_args, num_classes, multi_label=False, max_seq_le
         self.num_classes = num_classes
         self.multi_label = multi_label
 
-        self.model_kwargs = _check_model_kwargs(model_kwargs)
-        self.trainer_kwargs = _check_trainer_kwargs(trainer_kwargs)
-
-        model_kwargs = self.model_kwargs.copy()
+        model_kwargs = self.setfit_model_args.model_kwargs.copy()
         if self.multi_label and 'multi_target_strategy' not in model_kwargs:
             model_kwargs['multi_target_strategy'] = 'one-vs-rest'
 
@@ -264,7 +272,7 @@ def _fit(self, sub_train, sub_valid, setfit_train_kwargs):
             eval_dataset=sub_valid,
             batch_size=self.mini_batch_size,
             seed=seed,
-            **self.trainer_kwargs
+            **self.setfit_model_args.trainer_kwargs
         )
         trainer.train(max_length=self.max_seq_len, **setfit_train_kwargs)
         return self

tests/unit/small_text/integrations/transformers/classifiers/test_setfit.py

@@ -36,6 +36,36 @@ def test_setfit_model_arguments_init(self):
         sentence_transformer_model = 'sentence-transformers/all-MiniLM-L6-v2'
         args = SetFitModelArguments(sentence_transformer_model)
         self.assertEqual(sentence_transformer_model, args.sentence_transformer_model)
+        self.assertIsNotNone(args.model_kwargs)
+        self.assertEqual(0, len(args.model_kwargs))
+        self.assertIsNotNone(args.trainer_kwargs)
+        self.assertEqual(0, len(args.trainer_kwargs))
+        self.assertIsNotNone(args.model_loading_strategy)
+        self.assertEqual(ModelLoadingStrategy.DEFAULT, args.model_loading_strategy)
+        self.assertFalse(args.compile_model)
+
+    def test_setfit_model_arguments_init_with_model_kwargs(self):
+        sentence_transformer_model = 'sentence-transformers/all-MiniLM-L6-v2'
+        model_kwargs = {'cache_dir': '/tmp/cache'}
+        args = SetFitModelArguments(sentence_transformer_model, model_kwargs=model_kwargs)
+        self.assertEqual(sentence_transformer_model, args.sentence_transformer_model)
+        self.assertIsNotNone(args.model_kwargs)
+        self.assertEqual(1, len(args.model_kwargs))
+        self.assertIsNotNone(args.trainer_kwargs)
+        self.assertEqual(0, len(args.trainer_kwargs))
+        self.assertIsNotNone(args.model_loading_strategy)
+        self.assertEqual(ModelLoadingStrategy.DEFAULT, args.model_loading_strategy)
+        self.assertFalse(args.compile_model)
+
+    def test_setfit_model_arguments_init_with_trainer_kwargs(self):
+        sentence_transformer_model = 'sentence-transformers/all-MiniLM-L6-v2'
+        trainer_kwargs = {'batch_size': 32}
+        args = SetFitModelArguments(sentence_transformer_model, trainer_kwargs=trainer_kwargs)
+        self.assertEqual(sentence_transformer_model, args.sentence_transformer_model)
+        self.assertIsNotNone(args.model_kwargs)
+        self.assertEqual(0, len(args.model_kwargs))
+        self.assertIsNotNone(args.trainer_kwargs)
+        self.assertEqual(1, len(args.trainer_kwargs))
         self.assertIsNotNone(args.model_loading_strategy)
         self.assertEqual(ModelLoadingStrategy.DEFAULT, args.model_loading_strategy)
         self.assertFalse(args.compile_model)
@@ -46,6 +76,10 @@ def test_setfit_model_arguments_init_with_model_loading_strategy(self):
         args = SetFitModelArguments(sentence_transformer_model,
                                     model_loading_strategy=model_loading_strategy)
         self.assertEqual(sentence_transformer_model, args.sentence_transformer_model)
+        self.assertIsNotNone(args.model_kwargs)
+        self.assertEqual(0, len(args.model_kwargs))
+        self.assertIsNotNone(args.trainer_kwargs)
+        self.assertEqual(0, len(args.trainer_kwargs))
         self.assertIsNotNone(args.model_loading_strategy)
         self.assertEqual(model_loading_strategy, args.model_loading_strategy)
         self.assertFalse(args.compile_model)
@@ -61,6 +95,10 @@ def test_transformer_model_arguments_init_with_env_override(self):
             args = SetFitModelArguments(sentence_transformer_model)
 
             self.assertEqual(sentence_transformer_model, args.sentence_transformer_model)
+            self.assertIsNotNone(args.model_kwargs)
+            self.assertEqual(0, len(args.model_kwargs))
+            self.assertIsNotNone(args.trainer_kwargs)
+            self.assertEqual(0, len(args.trainer_kwargs))
             self.assertIsNotNone(args.model_loading_strategy)
             self.assertEqual(ModelLoadingStrategy.ALWAYS_LOCAL, args.model_loading_strategy)
             self.assertFalse(args.compile_model)
@@ -70,6 +108,10 @@ def test_setfit_model_arguments_init_with_compile(self):
         args = SetFitModelArguments(sentence_transformer_model,
                                     compile_model=True)
         self.assertEqual(sentence_transformer_model, args.sentence_transformer_model)
+        self.assertIsNotNone(args.model_kwargs)
+        self.assertEqual(0, len(args.model_kwargs))
+        self.assertIsNotNone(args.trainer_kwargs)
+        self.assertEqual(0, len(args.trainer_kwargs))
         self.assertIsNotNone(args.model_loading_strategy)
         self.assertEqual(ModelLoadingStrategy.DEFAULT, args.model_loading_strategy)
         self.assertTrue(args.compile_model)