Add SetFitModel.to (#229) (#236)

* Add SetFitModel.to

* Add docstring for 'to' to clarify that we don't copy

Note that this does mean that we differ from e.g. torch, which creates a moved copy rather than moving the e.g. tensor itself to a different device.

* Add tests for 'to'.

If CUDA is available, then try to move a model from the CPU to the GPU and back to the CPU.

* Reformat test_modeling

Co-authored-by: Tom Aarsen <Cubiegamedev@gmail.com>

Co-authored-by: Jegor Kitškerkin <jegor.kitskerkin@gmail.com>

src/setfit/modeling.py

@@ -330,6 +330,22 @@ def predict_proba(self, x_test: List[str]) -> Union[torch.Tensor, np.ndarray]:
         embeddings = self.model_body.encode(x_test, normalize_embeddings=self.normalize_embeddings)
         return self.model_head.predict_proba(embeddings)
 
+    def to(self, device: Union[str, torch.device]) -> "SetFitModel":
+        """Move this SetFitModel to `device`, and then return `self`. This method does not copy.
+
+        Args:
+            device (Union[str, torch.device]): The identifier of the device to move the model to.
+
+        Returns:
+            SetFitModel: Returns the original model, but now on the desired device.
+        """
+        self.model_body = self.model_body.to(device)
+
+        if isinstance(self.model_head, torch.nn.Module):
+            self.model_head = self.model_head.to(device)
+
+        return self
+
     def __call__(self, inputs):
         return self.predict(inputs)
 

tests/test_modeling.py

@@ -1,6 +1,7 @@
 from unittest import TestCase
 
 import numpy as np
+import torch
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
 from sklearn.linear_model import LogisticRegression
@@ -209,3 +210,30 @@ def test_setfit_from_pretrained_local_model_with_head(tmp_path):
     model = SetFitModel.from_pretrained(str(tmp_path.absolute()))
 
     assert isinstance(model, SetFitModel)
+
+
+def test_to_logistic_head():
+    model = SetFitModel.from_pretrained("sentence-transformers/paraphrase-albert-small-v2")
+    devices = (
+        [torch.device("cpu"), torch.device("cuda", 0), torch.device("cpu")]
+        if torch.cuda.is_available()
+        else [torch.device("cpu")]
+    )
+    for device in devices:
+        model.to(device)
+        assert model.model_body.device == device
+
+
+def test_to_torch_head():
+    model = SetFitModel.from_pretrained(
+        "sentence-transformers/paraphrase-albert-small-v2", use_differentiable_head=True
+    )
+    devices = (
+        [torch.device("cpu"), torch.device("cuda", 0), torch.device("cpu")]
+        if torch.cuda.is_available()
+        else [torch.device("cpu")]
+    )
+    for device in devices:
+        model.to(device)
+        assert model.model_body.device == device
+        assert model.model_head.device == device