Fix #3185: Support convert_to_numpy for output_value="token_embeddings"

sentence_transformers/SentenceTransformer.py

@@ -539,10 +539,6 @@ def encode(
         if convert_to_tensor:
             convert_to_numpy = False
 
-        if output_value != "sentence_embedding":
-            convert_to_tensor = False
-            convert_to_numpy = False
-
         input_was_string = False
         if isinstance(sentences, str) or not hasattr(
             sentences, "__len__"
@@ -669,9 +665,9 @@ def encode(
         elif convert_to_numpy:
             if not isinstance(all_embeddings, np.ndarray):
                 if all_embeddings and all_embeddings[0].dtype == torch.bfloat16:
-                    all_embeddings = np.asarray([emb.float().numpy() for emb in all_embeddings])
+                    all_embeddings = np.asarray([emb.float().cpu().numpy() for emb in all_embeddings])
                 else:
-                    all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
+                    all_embeddings = np.asarray([emb.cpu().numpy() for emb in all_embeddings])
         elif isinstance(all_embeddings, np.ndarray):
             all_embeddings = [torch.from_numpy(embedding) for embedding in all_embeddings]
 

tests/test_sentence_transformer.py

@@ -718,6 +718,35 @@ def test_empty_encode(stsb_bert_tiny_model: SentenceTransformer) -> None:
     assert embeddings.shape == (0,)
 
 
+@pytest.mark.parametrize(
+    ["convert_to_tensor", "convert_to_numpy", "expected_type"],
+    [
+        (True, False, torch.Tensor),
+        (False, False, torch.Tensor),
+        (None, False, torch.Tensor),
+        (True, True, torch.Tensor),
+        (False, True, np.ndarray),
+        (None, True, np.ndarray),
+        (True, None, torch.Tensor),
+        (False, None, np.ndarray),
+        (None, None, np.ndarray),
+    ],
+)
+def test_encode_token_embeddings_type(
+    stsb_bert_tiny_model_reused: SentenceTransformer, convert_to_tensor: bool, convert_to_numpy: bool, expected_type
+) -> None:
+    model = stsb_bert_tiny_model_reused
+
+    encode_kwargs = {}
+    if convert_to_tensor is not None:
+        encode_kwargs["convert_to_tensor"] = convert_to_tensor
+    if convert_to_numpy is not None:
+        encode_kwargs["convert_to_numpy"] = convert_to_numpy
+    embeddings = model.encode("Hello, World!", output_value="token_embeddings", **encode_kwargs)
+    assert isinstance(embeddings, expected_type)
+    assert embeddings.shape == (6, 128)
+
+
 @pytest.mark.skipif(not is_peft_available(), reason="PEFT must be available to test adapter methods.")
 def test_multiple_adapters() -> None:
     text = "Hello, World!"