mirror of
https://github.com/qdrant/fastembed.git
synced 2026-10-03 19:37:46 -05:00
Mohith GajjelaandGeorge Panchuk
9ab3ffe9bc
* 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 * 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 * fix: refine BM25 custom stemmer and stopword handling --------- Co-authored-by: George Panchuk <george.panchuk@qdrant.tech>
64 lines
2.7 KiB
Python
64 lines
2.7 KiB
Python
from fastembed import TextEmbedding, LateInteractionTextEmbedding, SparseTextEmbedding
|
|
from fastembed.sparse.bm25 import Bm25
|
|
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")
|
|
sparse_embedder = SparseTextEmbedding(model_name="", cache_dir="models")
|
|
bm25_embedder = Bm25(
|
|
model_name="",
|
|
k=1.0,
|
|
b=1.0,
|
|
avg_len=1.0,
|
|
language="",
|
|
token_max_length=1,
|
|
disable_stemmer=False,
|
|
specific_model_path="models",
|
|
stemmer=CustomStemmer(),
|
|
stopwords={"the", "a"},
|
|
)
|
|
|
|
text_embedder.list_supported_models()
|
|
text_embedder.embed(documents=[""], batch_size=1, parallel=1)
|
|
text_embedder.embed(documents="", parallel=None, task_id=1)
|
|
text_embedder.query_embed(query=[""], batch_size=1, parallel=1)
|
|
text_embedder.query_embed(query="", parallel=None)
|
|
text_embedder.passage_embed(texts=[""], batch_size=1, parallel=1)
|
|
text_embedder.passage_embed(texts=[""], parallel=None)
|
|
|
|
late_interaction_embedder.list_supported_models()
|
|
late_interaction_embedder.embed(documents=[""], batch_size=1, parallel=1)
|
|
late_interaction_embedder.embed(documents="", parallel=None)
|
|
late_interaction_embedder.query_embed(query=[""], batch_size=1, parallel=1)
|
|
late_interaction_embedder.query_embed(query="", parallel=None)
|
|
late_interaction_embedder.passage_embed(texts=[""], batch_size=1, parallel=1)
|
|
late_interaction_embedder.passage_embed(texts=[""], parallel=None)
|
|
|
|
reranker.list_supported_models()
|
|
reranker.rerank(query="", documents=[""], batch_size=1, parallel=1)
|
|
reranker.rerank(query="", documents=[""], parallel=None)
|
|
reranker.rerank_pairs(pairs=[("", "")], batch_size=1, parallel=1)
|
|
reranker.rerank_pairs(pairs=[("", "")], parallel=None)
|
|
|
|
sparse_embedder.list_supported_models()
|
|
sparse_embedder.embed(documents=[""], batch_size=1, parallel=1)
|
|
sparse_embedder.embed(documents="", batch_size=1, parallel=None)
|
|
sparse_embedder.query_embed(query=[""], batch_size=1, parallel=1)
|
|
sparse_embedder.query_embed(query="", batch_size=1, parallel=None)
|
|
sparse_embedder.passage_embed(texts=[""], batch_size=1, parallel=1)
|
|
sparse_embedder.passage_embed(texts=[""], batch_size=1, parallel=None)
|
|
|
|
bm25_embedder.list_supported_models()
|
|
bm25_embedder.embed(documents=[""], batch_size=1, parallel=1)
|
|
bm25_embedder.embed(documents="", batch_size=1, parallel=None)
|
|
bm25_embedder.query_embed(query=[""], batch_size=1, parallel=1)
|
|
bm25_embedder.query_embed(query="", batch_size=1, parallel=None)
|
|
bm25_embedder.raw_embed(documents=[""])
|