From 40e4de10a78cac64fcd686ae8546c13f3874cded Mon Sep 17 00:00:00 2001 From: Mohith Gajjela <109003762+Mohith26@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:17:31 -0500 Subject: [PATCH 1/3] feat: support pluggable stemmer and inline stopwords for BM25 Allow passing a custom stemmer (any object with a stem_word(word) -> str method, per the new Stemmer protocol) to Bm25, overriding the default SnowballStemmer. When a custom stemmer is provided, the supported-languages check is skipped, enabling languages without a Snowball algorithm such as Polish, Czech, Ukrainian, Slovak, Bulgarian, or Vietnamese. Also allow passing stopwords inline, overriding the per-language stopwords file shipped with the model. Both parameters are forwarded to parallel workers. Default behavior is unchanged when neither parameter is given. Fixes #654 --- fastembed/sparse/bm25.py | 42 ++++++++++++++++++++++++++++++++++------ 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index c31acd12c..785953806 100644 --- a/fastembed/sparse/bm25.py +++ b/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, runtime_checkable import mmh3 import numpy as np @@ -22,6 +22,17 @@ from fastembed.common.model_description import SparseModelDescription, ModelSource +@runtime_checkable +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", @@ -84,6 +95,12 @@ class Bm25(SparseTextEmbeddingBase): 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. + stemmer (Stemmer, optional): A custom stemmer to use instead of the default SnowballStemmer. + Any object with a `stem_word(word: str) -> str` method. When provided, the `language` + is not required to be in the list of supported languages, which enables languages + without a Snowball algorithm (e.g. Polish). Defaults to None. + stopwords (Iterable[str], optional): Inline stopwords to use instead of the stopwords file + shipped with the model for the chosen language. Defaults to None. Raises: ValueError: If the model_name is not in the format / e.g. BAAI/bge-base-en. """ @@ -99,15 +116,20 @@ def __init__( token_max_length: int = 40, disable_stemmer: bool = False, specific_model_path: str | None = None, + stemmer: Stemmer | None = None, + stopwords: Iterable[str] | None = None, **kwargs: Any, ): super().__init__(model_name, cache_dir, **kwargs) - if language not in supported_languages: + if stemmer is None and language not in supported_languages: raise ValueError(f"{language} language is not supported") else: self.language = language + self._custom_stemmer = stemmer + self._custom_stopwords = set(stopwords) if stopwords is not None else None + self.k = k self.b = b self.avg_len = avg_len @@ -128,11 +150,17 @@ def __init__( self.disable_stemmer = disable_stemmer if disable_stemmer: - self.stopwords: set[str] = set() - self.stemmer = None + self.stopwords: set[str] = ( + self._custom_stopwords if self._custom_stopwords is not None else set() + ) + self.stemmer: Stemmer | None = None else: - self.stopwords = set(self._load_stopwords(self._model_dir, self.language)) - self.stemmer = SnowballStemmer(language) + self.stopwords = ( + self._custom_stopwords + if self._custom_stopwords is not None + else set(self._load_stopwords(self._model_dir, self.language)) + ) + self.stemmer = stemmer if stemmer is not None else SnowballStemmer(language) self.tokenizer = SimpleTokenizer @@ -191,6 +219,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, } From 9489011af9e8aa7d35e495ceda8b17d0c1133bad Mon Sep 17 00:00:00 2001 From: Mohith Gajjela <109003762+Mohith26@users.noreply.github.com> Date: Tue, 28 Jul 2026 13:17:31 -0500 Subject: [PATCH 2/3] test: cover pluggable stemmer and inline stopwords for BM25 - custom stemmer callable is applied to tokens - default Snowball path is unchanged - unsupported language (Polish) works with a custom stemmer and still raises without one - inline stopwords override the file-based stopwords --- tests/test_sparse_embeddings.py | 58 +++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index e7c1c936c..e462f73f7 100644 --- a/tests/test_sparse_embeddings.py +++ b/tests/test_sparse_embeddings.py @@ -311,6 +311,64 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None: assert result == expected, f"Expected {expected}, but got {result}" +class _TruncatingStemmer: + """A trivial fake stemmer that truncates words to their first three characters.""" + + def stem_word(self, word: str) -> str: + return word[:3] + + +def test_custom_stemmer_is_used() -> None: + model = Bm25("Qdrant/bm25", language="english", stemmer=_TruncatingStemmer()) + + result = model._stem(["running", "jumped", "documents"]) + + assert result == ["run", "jum", "doc"], f"Custom stemmer was not applied, got {result}" + + +def test_default_stemmer_path_unchanged() -> None: + model = Bm25("Qdrant/bm25", language="english") + + result = model._stem(["running", "jumped", "documents"]) + + assert result == ["run", "jump", "document"], f"Default Snowball path changed, got {result}" + + +def test_custom_stemmer_enables_unsupported_language() -> None: + class _FakePolishStemmer: + def stem_word(self, word: str) -> str: + for suffix in ("ami", "ach"): + if word.endswith(suffix): + return word[: -len(suffix)] + return word + + # without a custom stemmer, unsupported languages are still rejected + with pytest.raises(ValueError): + Bm25("Qdrant/bm25", language="polish") + + model = Bm25( + "Qdrant/bm25", + language="polish", + stemmer=_FakePolishStemmer(), + stopwords=["i", "w"], + ) + + assert model.stopwords == {"i", "w"} + + result = model._stem(["kotami", "domach", "w"]) + + assert result == ["kot", "dom"], f"Expected ['kot', 'dom'], but got {result}" + + +def test_inline_stopwords_override_file_stopwords() -> None: + model = Bm25("Qdrant/bm25", language="english", stopwords=["fox"]) + + result = model._stem(["the", "fox"]) + + # "the" is no longer a stopword, "fox" is + assert result == ["the"], f"Inline stopwords were not applied, got {result}" + + def test_if_splade_query_embed_is_inference_free() -> None: is_ci = os.getenv("CI") model = SparseTextEmbedding( From be436680692defdad35b9bc41d3ffb62576a75a8 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Sat, 26 Sep 2026 23:58:31 +0700 Subject: [PATCH 3/3] fix: refine BM25 custom stemmer and stopword handling --- fastembed/sparse/bm25.py | 69 ++++++++++++--------- tests/test_sparse_embeddings.py | 103 +++++++++++++++++++------------- tests/type_stub.py | 7 +++ 3 files changed, 112 insertions(+), 67 deletions(-) diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index 785953806..6afcb673a 100644 --- a/fastembed/sparse/bm25.py +++ b/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, Protocol, Type, runtime_checkable +from typing import Any, Iterable, Protocol, Type import mmh3 import numpy as np @@ -22,7 +22,6 @@ from fastembed.common.model_description import SparseModelDescription, ModelSource -@runtime_checkable class Stemmer(Protocol): """Protocol for custom stemmers pluggable into Bm25. @@ -93,16 +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. - Any object with a `stem_word(word: str) -> str` method. When provided, the `language` - is not required to be in the list of supported languages, which enables languages - without a Snowball algorithm (e.g. Polish). Defaults to None. - stopwords (Iterable[str], optional): Inline stopwords to use instead of the stopwords file - shipped with the model for the chosen language. Defaults to None. + 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 / 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__( @@ -117,18 +123,21 @@ def __init__( disable_stemmer: bool = False, specific_model_path: str | None = None, stemmer: Stemmer | None = None, - stopwords: Iterable[str] | None = None, + stopwords: set[str] | None = None, **kwargs: Any, ): super().__init__(model_name, cache_dir, **kwargs) - if stemmer is None and 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 = set(stopwords) if stopwords is not None else None + self._custom_stopwords = ( + {word.lower() for word in stopwords} if stopwords is not None else None + ) self.k = k self.b = b @@ -149,18 +158,19 @@ def __init__( self.punctuation = set(get_all_punctuation()) self.disable_stemmer = disable_stemmer - if disable_stemmer: - self.stopwords: set[str] = ( - self._custom_stopwords if self._custom_stopwords is not None else set() - ) - self.stemmer: Stemmer | None = None + if self._custom_stopwords is not None: + self.stopwords = self._custom_stopwords + elif disable_stemmer: + self.stopwords = set() else: - self.stopwords = ( - self._custom_stopwords - if self._custom_stopwords is not None - else set(self._load_stopwords(self._model_dir, self.language)) - ) - self.stemmer = stemmer if stemmer is not None else SnowballStemmer(language) + 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 @@ -175,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 [] @@ -279,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) diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index e462f73f7..9fdbb84d7 100644 --- a/tests/test_sparse_embeddings.py +++ b/tests/test_sparse_embeddings.py @@ -311,62 +311,85 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None: assert result == expected, f"Expected {expected}, but got {result}" -class _TruncatingStemmer: - """A trivial fake stemmer that truncates words to their first three characters.""" - +class _PolishStemmer: def stem_word(self, word: str) -> str: - return word[:3] - - -def test_custom_stemmer_is_used() -> None: - model = Bm25("Qdrant/bm25", language="english", stemmer=_TruncatingStemmer()) - - result = model._stem(["running", "jumped", "documents"]) - - assert result == ["run", "jum", "doc"], f"Custom stemmer was not applied, got {result}" - + for suffix in ("ami", "ach"): + if word.endswith(suffix): + return word[: -len(suffix)] + return word -def test_default_stemmer_path_unchanged() -> None: - model = Bm25("Qdrant/bm25", language="english") - - result = model._stem(["running", "jumped", "documents"]) - - assert result == ["run", "jump", "document"], f"Default Snowball path changed, got {result}" +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))) -def test_custom_stemmer_enables_unsupported_language() -> None: - class _FakePolishStemmer: - def stem_word(self, word: str) -> str: - for suffix in ("ami", "ach"): - if word.endswith(suffix): - return word[: -len(suffix)] - return word + 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, + } - # without a custom stemmer, unsupported languages are still rejected - with pytest.raises(ValueError): - Bm25("Qdrant/bm25", language="polish") +@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", - language="polish", - stemmer=_FakePolishStemmer(), - stopwords=["i", "w"], + stopwords=stopwords, + disable_stemmer=disable_stemmer, ) - assert model.stopwords == {"i", "w"} + assert model._stem(["THE", "and", "FOX"]) == expected - result = model._stem(["kotami", "domach", "w"]) - assert result == ["kot", "dom"], f"Expected ['kot', 'dom'], but got {result}" +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_inline_stopwords_override_file_stopwords() -> None: - model = Bm25("Qdrant/bm25", language="english", stopwords=["fox"]) +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) - result = model._stem(["the", "fox"]) - # "the" is no longer a stopword, "fox" is - assert result == ["the"], f"Inline stopwords were not applied, got {result}" +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: diff --git a/tests/type_stub.py b/tests/type_stub.py index b22f80d90..105521b09 100644 --- a/tests/type_stub.py +++ b/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()