diff --git a/fastembed/text/static_embedding.py b/fastembed/text/static_embedding.py index 241886aa..62699f05 100644 --- a/fastembed/text/static_embedding.py +++ b/fastembed/text/static_embedding.py @@ -215,12 +215,20 @@ def _sum_rows(self, ids: NDArray[np.int64], lengths: NDArray[np.int64]) -> Numpy The sums are a product of a (documents x distinct tokens) count matrix with the rows of those tokens: one matrix multiplication, several times faster than segment sums with np.add.reduceat or a loop over the documents. Documents whose count matrix would exceed - POOLING_CELLS are halved until it fits, which bounds the temporary memory. + POOLING_CELLS are halved until it fits, which bounds the temporary memory. A single + document sums its rows directly. """ assert self.embeddings is not None if not ids.size: return np.zeros((len(lengths), self.embeddings.shape[1]), dtype=np.float32) rows = ids if self.token_mapping is None else self.token_mapping[ids] + if len(lengths) == 1: + # A single text, typically a query: summing its rows skips the fixed cost of building + # a count matrix, which dominates for short texts. + vectors = self.embeddings[rows] + if self.weights is not None: + vectors = vectors * self.weights[ids, None] + return vectors.sum(axis=0, keepdims=True, dtype=np.float32) distinct_rows, row_index = np.unique(rows, return_inverse=True) if len(lengths) > 1 and len(lengths) * len(distinct_rows) > self.POOLING_CELLS: half = len(lengths) // 2 diff --git a/tests/test_static_embeddings.py b/tests/test_static_embeddings.py index 78a8d5c7..fb958c0e 100644 --- a/tests/test_static_embeddings.py +++ b/tests/test_static_embeddings.py @@ -173,3 +173,17 @@ def test_pooling_budget_splits_without_changing_embeddings(model: TextEmbedding) tiny_budget = TextEmbedding(CI_MODEL) tiny_budget.model.POOLING_CELLS = 5 assert np.allclose(np.stack(list(tiny_budget.embed(docs))), expected, atol=1e-6) + + +def test_single_text_matches_batch(model: TextEmbedding) -> None: + docs = [ + "hello world", + "", + "🤗", + "hello 🤗 world", + "a much longer document about vectors " * 30, + ] + batch = np.stack(list(model.embed(docs, batch_size=len(docs)))) + single = np.stack([next(iter(model.embed(doc))) for doc in docs]) + assert np.allclose(single, batch, atol=1e-6) + assert single.dtype == np.float32