upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

batch inference should return same shape as individual inference (#547)

Browse Source
This commit is contained in:
Andrey Vasnetsov
2025-08-06 22:44:28 +02:00
committed by GitHub
parent acec31277b
commit faf3d9fe18
2 changed files with 27 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
fastembed/late_interaction/colbert.py
View File

@@ -46,7 +46,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
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 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
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

28
tests/test_late_interaction_embeddings.py
View File

@@ -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)