Fix semantic_search_usearch() for single query

[karmi]

This patch fixes a bug where the sentence_transformers.quantization.semantic_search_usearch() method would fail with TypeError: 'numpy.float32' object is not iterable where only a single query is used.

Example code:

import numpy as np
import sentence_transformers

query_embeddings = np.array([[1, 2, 3]], dtype=np.int8)
corpus_embeddings = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int8)

hits, duration = sentence_transformers.quantization.semantic_search_usearch(
  query_embeddings=query_embeddings,
  corpus_embeddings=corpus_embeddings,
  corpus_precision="binary",
  rescore=False,
  top_k=min(len(corpus_embeddings), 10))

This raises TypeError: 'numpy.float32' object is not iterable without this patch, because the code expects the scores and indices to be a 2D array:

sentence-transformers/sentence_transformers/quantization.py

Lines 300 to 309 in 0253363

	outputs = (
	[
	[
	{"corpus_id": int(neighbor), "score": float(score)}
	for score, neighbor in zip(scores[query_id], indices[query_id])
	]
	for query_id in range(len(query_embeddings))
	],
	delta_t,
	)

However, the usearch implementation in the Python client returns 1D arrays for a single query, and a 2D array for multiple queries:

def search(self, vectors: np.ndarray, count: int) -> Matches:
    if vectors.ndim == 1 or (vectors.ndim == 2 and vectors.shape[0] == 1):
        return self.search_one(vectors, count)
    else:
        return self.search_many(vectors, count)

The problem also manifests in the example code here:

sentence-transformers/examples/applications/embedding-quantization/semantic_search_usearch.py

Lines 16 to 19 in 0253363

	queries = [
	"How do I become a good programmer?",
	"How do I become a good data scientist?",
	]

When just a single query is passed, it fails with the error above.

I wasn't sure if I should add tests for this patch — I can definitely do so, when needed.

[tomaarsen]

Hello!

... or (vectors.ndim == 2 and vectors.shape[0] == 1)

Oh I see! That's unexpected! I appreciate the details to reproduce and regarding the cause of the issue.
No tests will be necessary for this - I've opted to keep the tests minimal for the semantic_search_... quantization helper functions, as they're designed primarily to be quick helpers to get some experimentation going, to later be replaced by a more robust solution such as this (i.e. one that doesn't create an index on the fly 😄)

I've reproduced the issue and can confirm that your solution fixes it. Thanks a bunch! I'll merge this once the tests are green.

• Tom Aarsen

[karmi]

Thank you!