From 7323574d64fb14ca1a8208f966cccab7414cf3f0 Mon Sep 17 00:00:00 2001 From: Karel Minarik Date: Wed, 3 Apr 2024 17:07:42 +0200 Subject: [PATCH] Fix semantic_search_usearch() for single query This patch fixes a bug where the semantic_search_usearch() method would fail with `TypeError: 'numpy.float32' object is not iterable` where only a single query is used. --- sentence_transformers/quantization.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/sentence_transformers/quantization.py b/sentence_transformers/quantization.py index 44eaf9e6b..d21248b28 100644 --- a/sentence_transformers/quantization.py +++ b/sentence_transformers/quantization.py @@ -277,6 +277,11 @@ def semantic_search_usearch( scores = matches.distances indices = matches.keys + if scores.ndim < 2: + scores = np.atleast_2d(scores) + if indices.ndim < 2: + indices = np.atleast_2d(indices) + # If rescoring is enabled, we need to rescore the results using the rescore_embeddings if rescore_embeddings is not None: top_k_embeddings = np.array([corpus_index.get(query_indices) for query_indices in indices])