Fix SentenceTransformer encode documentation return type default (num…

…py vectors) (#2546) * Fix SentenceTransformer encode documentation return type default * fix encode documentation * fix to numpy array * Add more detailed description to return values Co-authored-by: Tom Aarsen <[email protected]> --------- Co-authored-by: christopherkeibel <[email protected]> Co-authored-by: Tom Aarsen <[email protected]>
UKPLab · Mar 19, 2024 · caaa28d · caaa28d
1 parent 87f4180
commit caaa28d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/sentence_transformers/SentenceTransformer.py b/sentence_transformers/SentenceTransformer.py
@@ -281,8 +281,9 @@ def encode(
         :param normalize_embeddings: Whether to normalize returned vectors to have length 1. In that case,
             the faster dot-product (util.dot_score) instead of cosine similarity can be used.
 
-        :return: By default, a list of tensors is returned. If convert_to_tensor, a stacked tensor is returned.
-            If convert_to_numpy, a numpy matrix is returned.
+        :return: By default, a 2d numpy array with shape [num_inputs, output_dimension] is returned. If only one string
+            input is provided, then the output is a 1d array with shape [output_dimension]. If `convert_to_tensor`, a
+            torch Tensor is returned instead.
         """
         self.eval()
         if show_progress_bar is None: