fix: validate BM25 average document length

fastembed/sparse/bm25.py

@@ -1,5 +1,6 @@
 import os
 from collections import defaultdict
+from math import isfinite
 from multiprocessing import get_all_start_methods
 from pathlib import Path
 from typing import Any, Iterable, Protocol, Type
@@ -95,7 +96,8 @@ class Bm25(SparseTextEmbeddingBase):
             I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.
         b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
             Defaults to 0.75.
-        avg_len (float, optional): The average length of the documents in the corpus. Defaults to 256.0.
+        avg_len (float, optional): The average length of the documents in the corpus.
+            Must be finite and greater than 0. Defaults to 256.0.
         language (str): Selects the default Snowball stemmer and bundled stopwords.
             Unsupported languages require a custom stemmer or disable_stemmer=True,
             and have no bundled stopwords.
@@ -111,8 +113,8 @@ class Bm25(SparseTextEmbeddingBase):
             shipped with the model for the chosen language. Words are lowercased and matched
             before stemming. None uses the default stopwords; an empty set disables filtering.
     Raises:
-        ValueError: If a default stemmer is unavailable for the language, or stemmer is supplied with
-            disable_stemmer=True.
+        ValueError: If avg_len is non-positive or non-finite, a default stemmer is unavailable
+            for the language, or stemmer is supplied with disable_stemmer=True.
     """
 
     def __init__(
@@ -130,8 +132,12 @@ def __init__(
         stopwords: set[str] | None = None,
         **kwargs: Any,
     ):
+        """Initialize BM25, rejecting invalid avg_len before loading model files."""
         super().__init__(model_name, cache_dir, **kwargs)
 
+        if not isfinite(avg_len) or avg_len <= 0:
+            raise ValueError("avg_len must be finite and greater than 0")
+
         if disable_stemmer and stemmer is not None:
             raise ValueError("stemmer cannot be supplied with disable_stemmer=True")
         if not disable_stemmer and stemmer is None and language not in supported_languages:

tests/test_bm25_parameters.py

@@ -0,0 +1,84 @@
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from fastembed.sparse.bm25 import Bm25
+from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
+
+
+@pytest.mark.parametrize("model_class", [Bm25, SparseTextEmbedding])
+@pytest.mark.parametrize("avg_len", [0.0, -0.0, -1.0, float("nan"), float("inf"), -float("inf")])
+def test_bm25_rejects_invalid_average_length(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+    model_class: type[Bm25] | type[SparseTextEmbedding],
+    avg_len: float,
+) -> None:
+    """Reject non-positive and non-finite lengths through both public entry points."""
+
+    def unexpected_download(*args, **kwargs):
+        """Fail if invalid parameters reach model download."""
+        pytest.fail("Invalid avg_len must be rejected before downloading model files")
+
+    monkeypatch.setattr(Bm25, "download_model", unexpected_download)
+    with pytest.raises(ValueError, match="avg_len must be finite and greater than 0"):
+        model_class(
+            "Qdrant/bm25",
+            avg_len=avg_len,
+            disable_stemmer=True,
+            specific_model_path=str(tmp_path),
+            cache_dir=str(tmp_path),
+        )
+
+
+@pytest.mark.parametrize("model_class", [Bm25, SparseTextEmbedding])
+@pytest.mark.parametrize(
+    "avg_len,expected_weights",
+    [
+        (0.5, {"fox": 4.4 / 7.7, "dog": 2.2 / 6.7}),
+        (3.0, {"fox": 1.375, "dog": 1.0}),
+        (256.0, {"fox": 4.4 / 2.310546875, "dog": 2.2 / 1.310546875}),
+    ],
+)
+def test_bm25_positive_average_length_preserves_weights(
+    tmp_path: Path,
+    model_class: type[Bm25] | type[SparseTextEmbedding],
+    avg_len: float,
+    expected_weights: dict[str, float],
+) -> None:
+    """Keep BM25 weights and empty-input behavior unchanged for positive lengths."""
+    model = model_class(
+        "Qdrant/bm25",
+        avg_len=avg_len,
+        disable_stemmer=True,
+        specific_model_path=str(tmp_path),
+        cache_dir=str(tmp_path),
+    )
+
+    embedding = next(iter(model.embed("fox fox dog")))
+    assert np.isfinite(embedding.values).all()
+    assert dict(zip(embedding.indices, embedding.values)) == pytest.approx(
+        {Bm25.compute_token_id(token): weight for token, weight in expected_weights.items()}
+    )
+    assert next(iter(model.embed(""))).indices.size == 0
+
+
+def test_bm25_average_length_preserved_in_parallel(tmp_path: Path) -> None:
+    """Pass the configured corpus length to workers without changing embedding values."""
+    model = SparseTextEmbedding(
+        "Qdrant/bm25",
+        avg_len=64.0,
+        disable_stemmer=True,
+        specific_model_path=str(tmp_path),
+        cache_dir=str(tmp_path),
+    )
+    documents = ["fox fox dog", "dog", "", "fox"]
+    serial = list(model.embed(documents, batch_size=2))
+    # More documents than batch_size forces worker processes to be started.
+    parallel = list(model.embed(documents, batch_size=2, parallel=2))
+
+    assert len(serial) == len(parallel) == len(documents)
+    for expected, actual in zip(serial, parallel):
+        assert actual.indices.tolist() == expected.indices.tolist()
+        assert actual.values.tolist() == pytest.approx(expected.values.tolist())