From 45b37739b3d2d0b048a4f41fa5b87c7f2ec42280 Mon Sep 17 00:00:00 2001 From: Andrey Vasnetsov Date: Sun, 3 Aug 2025 19:51:03 -0600 Subject: [PATCH] batch inference should return same shape as individual inference --- fastembed/late_interaction/colbert.py | 7 ++++-- tests/test_late_interaction_embeddings.py | 28 ++++++++++++++++++----- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index 9d1ccf9a2..e3eaaf8a8 100644 --- a/fastembed/late_interaction/colbert.py +++ b/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 diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index 9f433ea64..f5b313304 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/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)