new: support pluggable stemmer and inline stopwords for Bm25

fastembed/sparse/bm25.py

@@ -2,7 +2,7 @@
 from collections import defaultdict
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Iterable, Type
+from typing import Any, Iterable, Protocol, Type
 
 import mmh3
 import numpy as np
@@ -22,6 +22,16 @@
 from fastembed.common.model_description import SparseModelDescription, ModelSource
 
 
+class Stemmer(Protocol):
+    """Protocol for custom stemmers pluggable into Bm25.
+
+    Any object exposing a `stem_word` method, e.g. `py_rust_stemmers.SnowballStemmer`,
+    satisfies this protocol.
+    """
+
+    def stem_word(self, word: str) -> str: ...
+
+
 supported_languages = [
     "arabic",
     "danish",
@@ -82,10 +92,23 @@ class Bm25(SparseTextEmbeddingBase):
         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.
-        language (str): Specifies the language for the stemmer.
-        disable_stemmer (bool): Disable the stemmer.
+        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.
+        disable_stemmer (bool): Disable stemming and bundled stopwords. Explicit inline
+            stopwords still apply. Cannot be combined with a custom stemmer. Lowercasing,
+            tokenization, punctuation removal and token_max_length filtering still apply.
+        stemmer (Stemmer, optional): A custom stemmer to use instead of the default SnowballStemmer.
+            Its stem_word method receives lowercase tokens after stopword and length filtering,
+            and returns a stem or an empty string to discard the token. The same stemmer is used
+            for documents and queries. For parallel embedding it must be deepcopyable and
+            pickleable, with its class defined in an importable module. Defaults to None.
+        stopwords (set[str], optional): Inline stopwords to use instead of the stopwords file
+            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 the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        ValueError: If a default stemmer is unavailable for the language, or stemmer is supplied with
+            disable_stemmer=True.
     """
 
     def __init__(
@@ -99,14 +122,22 @@ def __init__(
         token_max_length: int = 40,
         disable_stemmer: bool = False,
         specific_model_path: str | None = None,
+        stemmer: Stemmer | None = None,
+        stopwords: set[str] | None = None,
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, **kwargs)
 
-        if language not in supported_languages:
+        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:
             raise ValueError(f"{language} language is not supported")
-        else:
-            self.language = language
+        self.language = language
+
+        self._custom_stemmer = stemmer
+        self._custom_stopwords = (
+            {word.lower() for word in stopwords} if stopwords is not None else None
+        )
 
         self.k = k
         self.b = b
@@ -127,11 +158,18 @@ def __init__(
         self.punctuation = set(get_all_punctuation())
         self.disable_stemmer = disable_stemmer
 
-        if disable_stemmer:
-            self.stopwords: set[str] = set()
-            self.stemmer = None
+        if self._custom_stopwords is not None:
+            self.stopwords = self._custom_stopwords
+        elif disable_stemmer:
+            self.stopwords = set()
         else:
             self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
+
+        if stemmer is not None:
+            self.stemmer = stemmer
+        elif disable_stemmer:
+            self.stemmer = None
+        else:
             self.stemmer = SnowballStemmer(language)
 
         self.tokenizer = SimpleTokenizer
@@ -147,6 +185,9 @@ def _list_supported_models(cls) -> list[SparseModelDescription]:
 
     @classmethod
     def _load_stopwords(cls, model_dir: Path, language: str) -> list[str]:
+        if language not in supported_languages:
+            return []
+
         stopwords_path = model_dir / f"{language}.txt"
         if not stopwords_path.exists():
             return []
@@ -191,6 +232,8 @@ def _embed_documents(
                 "language": self.language,
                 "token_max_length": self.token_max_length,
                 "disable_stemmer": self.disable_stemmer,
+                "stemmer": self._custom_stemmer,
+                "stopwords": self._custom_stopwords,
                 "local_files_only": local_files_only,
                 "specific_model_path": specific_model_path,
             }
@@ -249,7 +292,9 @@ def _stem(self, tokens: list[str]) -> list[str]:
             if len(token) > self.token_max_length:
                 continue
 
-            stemmed_token = self.stemmer.stem_word(lower_token) if self.stemmer else lower_token
+            stemmed_token = (
+                self.stemmer.stem_word(lower_token) if self.stemmer is not None else lower_token
+            )
 
             if stemmed_token:
                 stemmed_tokens.append(stemmed_token)

tests/test_sparse_embeddings.py

@@ -311,6 +311,87 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
     assert result == expected, f"Expected {expected}, but got {result}"
 
 
+class _PolishStemmer:
+    def stem_word(self, word: str) -> str:
+        for suffix in ("ami", "ach"):
+            if word.endswith(suffix):
+                return word[: -len(suffix)]
+        return word
+
+
+def test_custom_stemmer_applies_to_documents_and_queries() -> None:
+    model = SparseTextEmbedding(
+        "Qdrant/bm25",
+        language="polish",
+        stemmer=_PolishStemmer(),
+        stopwords={"W"},
+        k=1.0,
+        b=0.0,
+    )
+    text = "Kotami kotach w domach!"
+    document = next(iter(model.embed(text)))
+    query = next(iter(model.query_embed(text)))
+
+    assert dict(zip(document.indices, document.values)) == pytest.approx(
+        {Bm25.compute_token_id("kot"): 4 / 3, Bm25.compute_token_id("dom"): 1.0}
+    )
+    assert dict(zip(query.indices, query.values)) == {
+        Bm25.compute_token_id("kot"): 1.0,
+        Bm25.compute_token_id("dom"): 1.0,
+    }
+
+
+@pytest.mark.parametrize(
+    "disable_stemmer,stopwords,expected",
+    [
+        (False, None, ["fox"]),
+        (True, None, ["the", "and", "fox"]),
+        (False, set(), ["the", "and", "fox"]),
+        (True, set(), ["the", "and", "fox"]),
+        (False, {"The", "FOX"}, ["and"]),
+        (True, {"The", "FOX"}, ["and"]),
+    ],
+)
+def test_stopword_precedence(
+    stopwords: set[str] | None, disable_stemmer: bool, expected: list[str]
+) -> None:
+    model = Bm25(
+        "Qdrant/bm25",
+        stopwords=stopwords,
+        disable_stemmer=disable_stemmer,
+    )
+
+    assert model._stem(["THE", "and", "FOX"]) == expected
+
+
+def test_unsupported_language_requires_custom_or_disabled_stemmer() -> None:
+    with pytest.raises(ValueError, match="polish language is not supported"):
+        Bm25("Qdrant/bm25", language="polish", stopwords=set())
+
+
+def test_custom_and_disabled_stemmer_are_mutually_exclusive() -> None:
+    with pytest.raises(ValueError, match="stemmer cannot be supplied with disable_stemmer=True"):
+        Bm25("Qdrant/bm25", stemmer=_PolishStemmer(), disable_stemmer=True)
+
+
+def test_parallel_matches_serial() -> None:
+    model = SparseTextEmbedding(
+        "Qdrant/bm25",
+        language="polish",
+        stemmer=_PolishStemmer(),
+        stopwords={"w", "kotami"},
+    )
+    documents = ["Kotami kotach w domach", "W domach", "", "kotami"]
+    serial = list(model.embed(documents))
+    # More documents than batch_size ensures this actually starts worker processes.
+    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())
+
+
 def test_if_splade_query_embed_is_inference_free() -> None:
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(

tests/type_stub.py

@@ -3,6 +3,11 @@
 from fastembed.rerank.cross_encoder import TextCrossEncoder
 
 
+class CustomStemmer:
+    def stem_word(self, word: str) -> str:
+        return word
+
+
 text_embedder = TextEmbedding(cache_dir="models")
 late_interaction_embedder = LateInteractionTextEmbedding(model_name="", cache_dir="models")
 reranker = TextCrossEncoder(model_name="", cache_dir="models")
@@ -16,6 +21,8 @@
     token_max_length=1,
     disable_stemmer=False,
     specific_model_path="models",
+    stemmer=CustomStemmer(),
+    stopwords={"the", "a"},
 )
 
 text_embedder.list_supported_models()