new: sum rows directly for single texts in static embeddings

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

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