batch inference should return same shape as individual inference

fastembed/late_interaction/colbert.py

@@ -46,7 +46,8 @@ def _post_process_onnx_output(
         self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any
     ) -> Iterable[NumpyArray]:
         if not is_doc:
-            return output.model_output
+            for embedding in output.model_output:
+                yield embedding
 
         if output.input_ids is None or output.attention_mask is None:
             raise ValueError(
@@ -62,7 +63,9 @@ def _post_process_onnx_output(
         norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
         norm_clamped = np.maximum(norm, 1e-12)
         output.model_output /= norm_clamped
-        return output.model_output
+
+        for embedding, attention_mask in zip(output.model_output, output.attention_mask):
+            yield embedding[attention_mask == 1]
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any

tests/test_late_interaction_embeddings.py

@@ -170,6 +170,23 @@ def test_batch_embedding(model_name: str):
         delete_model_cache(model.model._model_dir)
 
 
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_batch_inference_size_same_as_single_inference(model_name: str):
+    is_ci = os.getenv("CI")
+
+    model = LateInteractionTextEmbedding(model_name=model_name)
+    docs_to_embed = [
+        "short document",
+        "A bit longer document, which should not affect the size"
+    ]
+    result = list(model.embed(docs_to_embed, batch_size=1))
+    result_2 = list(model.embed(docs_to_embed, batch_size=2))
+    assert len(result[0]) == len(result_2[0])
+
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
 @pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
 def test_single_embedding(model_name: str):
     is_ci = os.getenv("CI")
@@ -219,17 +236,16 @@ def test_parallel_processing(token_dim: int, model_name: str):
 
     docs = ["hello world", "flag embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
-    embeddings = np.stack(embeddings, axis=0)
 
     embeddings_2 = list(model.embed(docs, batch_size=10, parallel=None))
-    embeddings_2 = np.stack(embeddings_2, axis=0)
 
     embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
-    embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape[0] == len(docs) and embeddings.shape[-1] == token_dim
-    assert np.allclose(embeddings, embeddings_2, atol=1e-3)
-    assert np.allclose(embeddings, embeddings_3, atol=1e-3)
+    assert len(embeddings) == len(docs) and embeddings[0].shape[-1] == token_dim
+
+    for i in range(len(embeddings)):
+        assert np.allclose(embeddings[i], embeddings_2[i], atol=1e-3)
+        assert np.allclose(embeddings[i], embeddings_3[i], atol=1e-3)
 
     if is_ci:
         delete_model_cache(model.model._model_dir)