From caaa28d75064cd43c4a836edb060716804905310 Mon Sep 17 00:00:00 2001 From: Christopher Keibel <55911084+CKeibel@users.noreply.github.com> Date: Tue, 19 Mar 2024 08:40:01 +0100 Subject: [PATCH] Fix SentenceTransformer encode documentation return type default (numpy 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 <37621491+tomaarsen@users.noreply.github.com> --------- Co-authored-by: christopherkeibel Co-authored-by: Tom Aarsen <37621491+tomaarsen@users.noreply.github.com> --- sentence_transformers/SentenceTransformer.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sentence_transformers/SentenceTransformer.py b/sentence_transformers/SentenceTransformer.py index e1c273727..3d7869ace 100644 --- 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: