semantic_router.encoders.bm25
BM25Encoder Objects
class BM25Encoder(SparseEncoder, FittableMixin, AsymmetricSparseMixin)BM25Encoder, running a vectorized version of ATIRE BM25 algorithm
Concept:
- BM25 uses scoring between queries & corpus to retrieve the most relevant documents ∈ corpus
- most vector databases (VDB) store embedded documents and score them versus received queries for retrieval
- we need to break up the BM25 formula into
encode_queriesandencode_documents, with the latter to be stored in VDB - dot product of
encode_queries(q)andencode_documents([D_0, D_1, ...])is the BM25 score of the documents[D_0, D_1, ...]for the given queryq - we train a BM25 encoder's normalization parameters on a sufficiently large corpus to capture target language distribution
- these trained parameter allow us to balance TF & IDF of query & documents for retrieval (read more on how BM25 fixes issues with TF-IDF)
The reference formula is section 4.1 of the ATIRE paper:
RSV_d = Σ_{t ∈ q} log(N / df_t) · (k1 + 1) · tf_td
/ (k1 · ((1 - b) + b · L_d / L_avg) + tf_td)Deviations from the paper
This implementation is faithful to the formula above on the document side, but deviates from it in three documented ways. Each is deliberate; see the inline comments at the point of use for the details.
-
IDF smoothing. We compute
log((N + 1) / (df_t + 0.5))rather thanlog(N / df_t), to avoid a division by zero for unseen terms. This preserves the strictly-positive IDF that the ATIRE variant exists for. -
Query weight normalization. The paper produces an unbounded relevance score, which is fine for a search engine that only ranks within a single query. A router instead compares scores against a fixed threshold, so scores must be comparable across queries. We therefore L1-normalize the query weights and scale them by
1 / (k1 + 1), which bounds the query·document dot product to [0, 1] -- the same range as the cosine similarity returned by the dense encoders, so the two can be blended in a hybridencode_documents0 and thresholded alike.This is ranking-neutral within a query (it is a single positive constant per query) but it does discard cross-query IDF magnitude: two queries with different total IDF mass can produce the same score.
-
Parameter defaults. We default to
encode_documents1 rather than the paper'sencode_documents2. See theencode_documents3 andencode_documents4 parameter docs below.
ATIRE Paper: https://www.cs.otago.ac.nz/research/student-publications/atire-opensource.pdf Pinecone Implementation: https://github.com/pinecone-io/pinecone-text/blob/8399f9ff28c4652766c35165c0db9b0eff309077/pinecone_text/sparse/bm25_encoder.py
Arguments:
encode_documents3 (encode_documents6): normalizer parameter that limits how much a single query termencode_documents7 can affect score for documentencode_documents8. Defaults to 1.5, the standard Robertson value, not the 0.9 reported in the ATIRE paper. The paper's value was trained on the INEX 2008 Wikipedia collection, whose documents average 878 words (paper table 5), whereas this encoder is fitted on route utterances averaging on the order of ten tokens. Transferring the paper's parameters to utterance-length text is not obviously more correct, and changing the default would silently shift the scores that existing thresholds were tuned against.encode_documents4 (encode_documents6): normalizer parameter that balances the effect of a single document length compared to the average document length. Defaults to 0.75 rather than the paper's 0.4, for the same reason asencode_documents3. Note thatencode_documents4 is considerably more influential on short text: over utterances spanning 0.3x to 2.2x the average length,encode_queries(q)3 spreads the document factor by ~2.2x versus ~1.4x forencode_queries(q)4. Lower it towards the paper's value if length is over-weighted for your route set.encode_queries(q)5 (encode_queries(q)6): number of documents in the trained corpusencode_queries(q)7 (encode_queries(q)8): float representing the average document length in the trained corpusencode_queries(q)9 (encode_documents([D_0, D_1, ...])0numpy.ndarrayencode_documents([D_0, D_1, ...])1): (1, tokenizer.vocab_size) shaped array, denoting how many documents containencode_documents([D_0, D_1, ...])2
fit
def fit(routes: List[Route]) -> "BM25Encoder"Trains the encoder weights on the provided routes.
Arguments:
routes(List[Route]): List of routes to train the encoder on.
encode_queries
def encode_queries(queries: list[str]) -> list[SparseEmbedding]Returns BM25 scores for queries using precomputed corpus scores.
Arguments:
queries(list): List of queries to encode
Returns:
list[SparseEmbedding]: BM25 scores for each query against the corpus
encode_documents
def encode_documents(documents: list[str],
batch_size: int | None = None) -> list[SparseEmbedding]Returns document term frequency normed by itself & average trained corpus length
(This is the right-hand side of the BM25 equation, which gets matmul-ed with the query IDF component)
LaTeX:
where:
f(d_i, D) is frequency of term d_i ∈ D
|D| is the document length
avgdl is average document length in trained corpus
Arguments:
documents(list): List of queries to encode
Returns:
list[SparseEmbedding]: Encoded queries (as either sparse or dict)
model
def model(docs: List[str]) -> list[SparseEmbedding]Encode documents using BM25, with different encoding for queries vs documents to be indexed.
Arguments:
docs: List of documents to encodeis_query: If True, use query encoding, else use document encoding
Returns:
List of sparse embeddings