diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index c31acd12c..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, 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 / 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) diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index e7c1c936c..9fdbb84d7 100644 --- a/tests/test_sparse_embeddings.py +++ b/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( 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()