BM25 & Inverted Index

STOPWORDS = {"the", "is", "at", "an", "of", "and", "or", "but", "in", "on", "to", "for", ...}

# Example inverted index
{
  "information": {
    "docIDs": [1, 3],
    "frequencies": [1, 2],
    "positions": [[5], [2, 15]]  # positions in each doc
  },
  "retrieval": {
    "docIDs": [1],
    "frequencies": [1],
    "positions": [[7]]
  },
  "system": {
    "docIDs": [1, 3],
    "frequencies": [1, 1],
    "positions": [[8], [3]]
  }
}

Doc 1: "cat sat mat"
Doc 2: "dog sat log"
Doc 3: "cat dog"

(1, "cat", 0), (1, "sat", 1), (1, "mat", 2),
(2, "dog", 0), (2, "sat", 1), (2, "log", 2),
(3, "cat", 0), (3, "dog", 1)

cat:    [1→[0], 3→[0]]
dog:    [2→[0], 3→[1]]
log:    [2→[2]]
mat:    [1→[2]]
sat:    [1→[1], 2→[1]]

Posting["cat"] = [1, 3]
Posting["dog"] = [2, 3]
Result = Intersection([1, 3], [2, 3]) = [3]
→ Only Doc 3 contains both "cat" and "dog"

Result = Union([1, 3], [2, 3]) = [1, 2, 3]

# Calculation:
Original:     [5,        28,       47,       102,      201]
Gaps:         [5,     23(28-5), 19(47-28), 55(102-47), 99(201-102)]
Savings: Gaps are smaller → fewer bytes needed per value

TF(t, d) = count of term t in document d
IDF(t) = log(N / df_t)
         where N = total documents, df_t = documents containing t

Score(q, d) = Σ TF(t, d) × IDF(t)

# Document scores:
Doc 1: information appears 3x, retrieval appears 1x
  Score = 3×1.0 + 1×1.3 = 4.3

Doc 2: information appears 1x, retrieval appears 5x
  Score = 1×1.0 + 5×1.3 = 7.5

Doc 3: information appears 10x, retrieval appears 0x
  Score = 10×1.0 + 0×1.3 = 10.0

Ranking: Doc 3 > Doc 2 > Doc 1

Score(q, d) = Σ_{i=1}^{n} IDF(q_i) × [TF(q_i, d) × (k1 + 1)] / [TF(q_i, d) + k1 × (1 - b + b × |d| / avgdl)]

where:
  q_i = query term i
  d = document
  TF(q_i, d) = count of q_i in d
  |d| = length of d in tokens
  avgdl = average document length
  k1 = saturation parameter (typically 1.2 - 2.0)
  b = length normalization parameter (typically 0.75)
  IDF(q_i) = log((N - df_q_i + 0.5) / (df_q_i + 0.5) + 1)
         where N = total documents, df = documents containing term

IDF(term) = log((N - df + 0.5) / (df + 0.5) + 1)

Example: N=1000 documents
  - "the" appears in 800 docs: IDF = log((1000-800+0.5)/(800+0.5)) ≈ 0.65 (common)
  - "information" in 100 docs: IDF = log((1000-100+0.5)/(100+0.5)) ≈ 2.29 (rare)
  - "xyzzy" in 1 doc: IDF = log((1000-1+0.5)/(1+0.5)) ≈ 6.2 (very rare)

Saturation curve: TF × (k1 + 1) / (TF + k1)

k1 = 0 (no saturation): score independent of TF (bad)
k1 = 1.2 (default): smooth saturation
k1 = 2.0 (high saturation): rewards frequency more
k1 = ∞ (linear): raw frequency

Normalization denominator: TF + k1 × (1 - b + b × |d| / avgdl)

b = 0: no length normalization (pure TF saturation)
b = 0.75 (default): moderate normalization
b = 1.0: full length normalization
  - Short docs: (1 - b + b × |d|/avgdl) < 1 → divisor smaller → score boosts
  - Long docs: divisor larger → score penalizes

Example with avgdl = 100:
  Short doc (50 tokens): divisor = 1.2 × (1 - 0.75 + 0.75 × 0.5) = 0.69
  Long doc (200 tokens): divisor = 1.2 × (1 - 0.75 + 0.75 × 2.0) = 2.1

Query term statistics:
  "machine": appears in 50 docs
  "learning": appears in 60 docs

IDF("machine") = log((1000-50+0.5)/(50+0.5)+1) ≈ 2.95
IDF("learning") = log((1000-60+0.5)/(60+0.5)+1) ≈ 2.82

Document: 180 tokens
  TF("machine") = 2
  TF("learning") = 3

BM25 calculation (k1=1.2, b=0.75):
  length_norm = 1 - 0.75 + 0.75 × (180/150) = 0.25 + 0.9 = 1.15

  score_machine = 2.95 × [2 × 2.2] / [2 + 1.2 × 1.15]
                = 2.95 × 4.4 / 3.38
                ≈ 3.83

  score_learning = 2.82 × [3 × 2.2] / [3 + 1.2 × 1.15]
                 = 2.82 × 6.6 / 4.38
                 ≈ 4.25

Total BM25 = 3.83 + 4.25 ≈ 8.08

for each document in union of posting lists:
    score = 0
    for each query term:
        if term in doc: score += BM25(term, doc)
    accumulate(doc, score)
sort and return top K

score_accumulator = {}
for each query term:
    for each doc in posting_list[term]:
        score_accumulator[doc] += BM25(term, doc)
sort and return top K

while documents remain:
    block = read_documents_until_memory_full()

    # Build in-memory index for this block
    for each document in block:
        for each term in document:
            add(term, docID, position)

    # Sort and write to disk
    sort_terms_by_id()
    write_to_disk("block_" + block_num)

# Merge all blocks
merge_sorted_blocks()

Shard 0: terms a-d
Shard 1: terms e-h
Shard 2: terms i-l
...

Query "machine learning":
  "machine" → hash to Shard X, fetch posting list
  "learning" → hash to Shard Y, fetch posting list
  Parallel merge on different machines

Score = Σ IDF(term) × [TF_weighted × (k1+1)] / [TF_weighted + k1×(1-b+b×|d|/avgdl)]

where TF_weighted = Σ (boost_field × TF_field)

Example:
  Title: boost = 3.0
  Body: boost = 1.0
  Anchor text: boost = 2.0

  If "machine" appears 1x in title, 2x in body:
    TF_weighted = 3.0×1 + 1.0×2 = 5.0

# Grid search on development set
for k1 in [0.5, 1.0, 1.2, 1.5, 2.0]:
    for b in [0.0, 0.25, 0.5, 0.75, 1.0]:
        evaluate_on_dev_set(k1, b)
        # Measure: NDCG, MAP, MRR

import math
import re
from collections import defaultdict
from typing import List, Dict, Tuple

class SimpleInvertedIndex:
    """Minimal inverted index with gap encoding support."""

    def __init__(self):
        self.posting_lists = defaultdict(list)  # term -> [(docID, freq)]
        self.doc_lengths = {}  # docID -> length
        self.num_docs = 0

    def add_document(self, doc_id: int, terms: List[str]):
        """Index a document's terms."""
        self.num_docs = max(self.num_docs, doc_id + 1)
        term_freqs = defaultdict(int)

        for term in terms:
            term_freqs[term] += 1

        self.doc_lengths[doc_id] = len(terms)

        for term, freq in term_freqs.items():
            self.posting_lists[term].append((doc_id, freq))

    def get_postings(self, term: str) -> List[Tuple[int, int]]:
        """Get posting list for a term."""
        return sorted(self.posting_lists.get(term, []))

    def document_frequency(self, term: str) -> int:
        """Count documents containing term."""
        return len(self.posting_lists.get(term, []))

class BM25Scorer:
    """Compute BM25 scores for documents."""

    def __init__(self, index: SimpleInvertedIndex, k1: float = 1.2, b: float = 0.75):
        self.index = index
        self.k1 = k1
        self.b = b
        self.avg_doc_length = sum(index.doc_lengths.values()) / max(1, len(index.doc_lengths))

    def idf(self, term: str) -> float:
        """Compute IDF for a term."""
        df = self.index.document_frequency(term)
        if df == 0:
            return 0
        N = self.index.num_docs
        return math.log((N - df + 0.5) / (df + 0.5) + 1)

    def score_document(self, doc_id: int, query_terms: List[str]) -> float:
        """Compute BM25 score for document and query."""
        score = 0
        doc_length = self.index.doc_lengths.get(doc_id, 0)

        for term in query_terms:
            postings = self.index.get_postings(term)
            tf = 0
            for did, freq in postings:
                if did == doc_id:
                    tf = freq
                    break

            if tf > 0:
                idf = self.idf(term)
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * doc_length / self.avg_doc_length)
                score += idf * numerator / denominator

        return score

    def search(self, query_terms: List[str], top_k: int = 10) -> List[Tuple[int, float]]:
        """Search for documents, return top-K."""
        candidate_docs = set()
        for term in query_terms:
            postings = self.index.get_postings(term)
            for doc_id, _ in postings:
                candidate_docs.add(doc_id)

        scores = []
        for doc_id in candidate_docs:
            score = self.score_document(doc_id, query_terms)
            if score > 0:
                scores.append((doc_id, score))

        scores.sort(key=lambda x: -x[1])
        return scores[:top_k]

class SimpleTokenizer:
    """Basic tokenization and preprocessing."""

    def __init__(self):
        self.stopwords = {'the', 'is', 'at', 'an', 'of', 'and', 'or', 'but', 'in', 'on', 'to', 'for', 'a', 'be', 'by'}

    def tokenize(self, text: str) -> List[str]:
        """Tokenize and normalize text."""
        text = text.lower()
        tokens = re.findall(r'\b\w+\b', text)
        return [t for t in tokens if t not in self.stopwords]

# Example usage with a small corpus
if __name__ == "__main__":
    corpus = [
        "Machine learning algorithms discover patterns in data",
        "Information retrieval systems rank relevant documents",
        "Neural networks process information through layers",
        "Search engines use ranking algorithms to find documents",
        "Data science applies machine learning to real-world problems"
    ]

    # Build index
    tokenizer = SimpleTokenizer()
    index = SimpleInvertedIndex()

    for doc_id, text in enumerate(corpus):
        terms = tokenizer.tokenize(text)
        index.add_document(doc_id, terms)

    # Search
    scorer = BM25Scorer(index)
    query = "machine learning algorithms"
    query_terms = tokenizer.tokenize(query)

    results = scorer.search(query_terms, top_k=3)

    print(f"Query: {query}")
    print(f"Results:")
    for doc_id, score in results:
        print(f"  Doc {doc_id} (score: {score:.3f}): {corpus[doc_id]}")

    # Output:
    # Query: machine learning algorithms
    # Results:
    #   Doc 0 (score: 4.523): Machine learning algorithms discover patterns in data
    #   Doc 4 (score: 2.341): Data science applies machine learning to real-world problems
    #   Doc 3 (score: 1.789): Search engines use ranking algorithms to find documents

# Evaluate on development set
def evaluate_bm25(corpus, relevance_judgments, k=10):
    """
    relevance_judgments: dict[query] = set[relevant_docIDs]
    """
    tokenizer = SimpleTokenizer()
    index = SimpleInvertedIndex()

    for doc_id, text in enumerate(corpus):
        index.add_document(doc_id, tokenizer.tokenize(text))

    scorer = BM25Scorer(index)

    ndcg_scores = []
    for query, relevant_docs in relevance_judgments.items():
        query_terms = tokenizer.tokenize(query)
        results = scorer.search(query_terms, top_k=k)
        result_docs = [doc_id for doc_id, _ in results]

        # Compute nDCG
        dcg = sum(1/math.log2(i+2) for i, doc_id in enumerate(result_docs) if doc_id in relevant_docs)
        idcg = sum(1/math.log2(i+2) for i in range(min(len(relevant_docs), k)))
        ndcg = dcg / idcg if idcg > 0 else 0
        ndcg_scores.append(ndcg)

    return sum(ndcg_scores) / len(ndcg_scores)