[fix] Avoid error if prompts & output_value=None (#3327)

* Avoid error if prompts & output_value=None

* Refactor, avoid function & batch_size override

Author: tomaarsen

sentence_transformers/SentenceTransformer.py

@@ -700,9 +700,15 @@ def encode(
                         embeddings.append(token_emb[0 : last_mask_id + 1])
                 elif output_value is None:  # Return all outputs
                     embeddings = []
-                    for sent_idx in range(len(out_features["sentence_embedding"])):
-                        row = {name: out_features[name][sent_idx] for name in out_features}
-                        embeddings.append(row)
+                    for idx in range(len(out_features["sentence_embedding"])):
+                        batch_item = {}
+                        for name, value in out_features.items():
+                            try:
+                                batch_item[name] = value[idx]
+                            except TypeError:
+                                # Handle non-indexable values (like prompt_length)
+                                batch_item[name] = value
+                        embeddings.append(batch_item)
                 else:  # Sentence embeddings
                     embeddings = out_features[output_value]
                     embeddings = embeddings.detach()

tests/test_sentence_transformer.py

@@ -376,6 +376,27 @@ def test_save_load_prompts() -> None:
         assert fresh_model.default_prompt_name == "query"
 
 
+def test_prompt_output_value_None(stsb_bert_tiny_model_reused) -> None:
+    model = stsb_bert_tiny_model_reused
+    outputs = model.encode(
+        ["Text one", "Text two"],
+        prompt="query: ",
+        output_value=None,
+    )
+    assert len(outputs) == 2
+    assert isinstance(outputs, list)
+    expected_keys = {
+        "input_ids",
+        "token_type_ids",
+        "attention_mask",
+        "sentence_embedding",
+        "token_embeddings",
+        "prompt_length",
+    }
+    assert set(outputs[0].keys()) == expected_keys
+    assert set(outputs[1].keys()) == expected_keys
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA must be available to test float16 support.")
 def test_load_with_torch_dtype() -> None:
     model = SentenceTransformer("sentence-transformers-testing/stsb-bert-tiny-safetensors")