Skip to content

Fix semantic_search crashing on a list of numpy embeddings - #4064

Open
MohammadHijjawi97 wants to merge 1 commit into
huggingface:mainfrom
MohammadHijjawi97:fix/semantic-search-list-of-numpy-embeddings
Open

MohammadHijjawi97 wants to merge 1 commit into
huggingface:mainfrom
MohammadHijjawi97:fix/semantic-search-list-of-numpy-embeddings

Conversation

@MohammadHijjawi97

Copy link
Copy Markdown

util.semantic_search() already accepts a 2D numpy array, but a list of 1D numpy vectors crashed in torch.stack. That is the shape SentenceTransformer.encode returns for a single string, so this public pattern failed:

query = model.encode("a query")
corpus = [model.encode(doc) for doc in documents]
semantic_search(query, corpus)

The list path now uses _convert_to_tensor, which already stacks lists of numpy arrays. A 1D corpus embedding is unsqueezed the same way a 1D query already was.

Tests: pytest tests/util/test_retrieval.py -k semantic_search
The new test fails on unmodified main and passes with this change.

Co-authored-by: Cursor <cursoragent@cursor.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant