BM25 Index
mediumagents_retrievalBuild a BM25 index — the keyword retriever behind the lexical half of a RAG pipeline, the part that answers "which documents mention this?".
BM25Index(documents, k1=1.5, b=0.75) indexes a collection once; a document is
identified by its position in documents. search(query, top_k) then returns
at most top_k (doc_id, score) pairs, highest score first, ties broken by the
smaller doc_id, and only documents scoring above 0 — a document that
shares no term with the query is not a result.
The score of document D for a query is a sum over the query's tokens
(duplicates included — "cat cat" counts twice):
tf(t, D) * (k1 + 1)
score(D) = SUM idf(t) * -----------------------------------------
t in q tf(t, D) + k1 * (1 - b + b * |D| / avgdl)
N - df(t) + 0.5
idf(t) = ln( 1 + ------------------- )
df(t) + 0.5
tf(t, D) is how often t occurs in D, df(t) how many documents contain
t at all, N the number of documents, |D| the token count of D and
avgdl the mean document length. That 1 + in the idf is the Lucene variant:
it keeps the idf positive even for a term present in every document, which is
what makes "score above 0" mean exactly "contains a query term".
tokenize is given — use it, so your terms and the tests' agree. The
collection is non-empty and at least one document has a token, but a single
document may well be empty. The split between the two halves of the pipeline is
part of the task: __init__ tokenizes the collection, search may only
tokenize the query — so the per-document scoring lives in _score_tokens,
which both entry points share.
- no rank_bm25, bm25s, sklearn, gensim, nltk, whoosh
Your solution
Edit BM25Index and run the real pytest suite in your browser — no install required. Your code is saved locally.
Loading the Python runtime (first run only)…
Hints
Hint 1
__init__ is the offline half: per document keep its token count and a Counter of its terms, plus one dict term -> [doc ids containing it] (the inverted index) and the idf of every term you saw.
Hint 2
search is the online half: tokenize the query once, take the union of the query terms' posting lists as the candidates, score only those, then sort by (-score, doc_id) and cut to top_k.