Keyword Ranking for In-Memory Text Documents

From Keywords to Relevance: The TF-IDF Intuition

At its simplest, keyword search seems easy: if a document contains a word from the search query, it's a match. But how do we rank these matches? A document mentioning "performance" once is probably less relevant than one discussing it ten times. And a match on "pgvector" is likely more significant than a match on "the".

This is where statistical relevance scoring comes in. The foundational concept is TF-IDF, which stands for Term Frequency-Inverse Document Frequency. It's a method to weigh the importance of a word within a document relative to a collection of documents (a "corpus").

It breaks down into two parts:

1. Term Frequency (TF): This measures how often a term appears in a specific document. The intuition is simple: the more a document talks about a topic, the more relevant it is to that topic.
2. Inverse Document Frequency (IDF): This measures how rare a term is across all documents in your corpus. Common words like "a", "the", or "is" will appear in almost every document, so they have a low IDF score. Rare, specific terms like "BullMQ" or "ACID" will appear in few documents, giving them a high IDF score. The IDF score boosts the importance of these unique, informative terms.

By multiplying these two values (TF * IDF), we get a score that reflects how important a word is to a specific document within the context of the whole collection.

The following video from The Coding Train provides a fantastic, hands-on explanation of why we need TF-IDF and how the two components work together.

The overall process of turning a document into a set of TF-IDF scores can be visualized as a straightforward pipeline.

To calculate a document's total relevance to a search query, you simply sum the TF-IDF scores of all query terms that appear in that document. Documents with higher total scores are ranked higher.

Here is a simple example of the calculations in a spreadsheet format, which can help solidify the concepts.

A Better Way: Okapi BM25

While TF-IDF is a great starting point, it has a couple of known weaknesses:

• Term Frequency Saturation: It assumes that a word appearing 20 times is twice as important as a word appearing 10 times. In reality, after a certain point, more mentions don't necessarily make a document much more relevant. The relevance "saturates".
• Document Length: It doesn't explicitly account for document length. A longer document has a higher chance of containing query terms, which can give it an unfair advantage over shorter, more concise documents.

To address this, the information retrieval community developed Okapi BM25 (Best Match 25). It's an evolution of TF-IDF and is considered a state-of-the-art ranking function, widely used in search engines.

BM25 refines the TF-IDF formula by adding two simple, tunable parameters:

• k1: Controls term frequency saturation. A higher k1 means that term frequency continues to contribute to the score for longer, while a lower k1 causes the score to plateau more quickly. A typical value is between 1.2 and 2.0.
• b: Controls how much to penalize longer documents. A value of 0 means no penalty, while a value of 1 applies a full penalty. A typical value is 0.75.

This short video gives a great, no-nonsense overview of the intuition behind BM25.

Implementing BM25 in TypeScript

Now, let's get practical and implement a BM25 ranker in TypeScript. We'll use the excellent article "Machine Learning: Full-Text Search in Javascript" by Burak Kanber as our guide. It provides a full, clear implementation that we can adapt.

The process involves a few distinct steps:

1. Tokenization: Breaking text down into individual words (tokens).
2. Indexing: Processing each document to build up the necessary statistics (like term counts and document lengths).
3. Searching: Calculating the BM25 score for each document against a query and returning a sorted list.

Let's read through the key parts of the implementation from the article.

A Self-Contained TypeScript Implementation

Now, let's consolidate these ideas into a TypeScript class that you can use directly in your Problem-Finder project. This implementation simplifies the tokenization from the article by omitting stemming and stop-word removal to avoid external dependencies, but it retains the full power of the BM25 algorithm.

// A simplified tokenizer that lowercases and splits by non-alphanumeric characters
function tokenize(text: string): string[] {
  return text
    .toLowerCase()
    .split(/[^a-z0-9]+/)
    .filter(token => token.length > 0);
}

interface Document {
  id: string;
  body: string;
}

interface IndexedDocument {
  id: string;
  body: string;
  tokens: string[];
  termCount: number;
  // A map of term -> { count: number, freq: number }
  terms: Record<string, { count: number; freq: number }>;
  _score?: number; // Used during search
}

export class BM25Ranker {
  private k1: number;
  private b: number;
  private documents: Record<string, IndexedDocument> = {};
  private terms: Record<string, { n: number; idf: number }> = {};
  private totalDocuments = 0;
  private totalDocumentTermLength = 0;
  private averageDocumentLength = 0;

  constructor({ k1 = 1.2, b = 0.75 }: { k1?: number; b?: number } = {}) {
    this.k1 = k1;
    this.b = b;
  }

  public addDocument(doc: Document): void {
    const { id, body } = doc;
    const tokens = tokenize(body);
    
    const indexedDoc: IndexedDocument = {
      id,
      body,
      tokens,
      termCount: tokens.length,
      terms: {},
    };

    // Update corpus stats
    this.totalDocuments++;
    this.totalDocumentTermLength += indexedDoc.termCount;
    this.averageDocumentLength = this.totalDocumentTermLength / this.totalDocuments;

    // Calculate term frequencies for this document
    const termCounts: Record<string, number> = {};
    for (const token of tokens) {
      termCounts[token] = (termCounts[token] || 0) + 1;
    }

    for (const term in termCounts) {
      const count = termCounts[term];
      indexedDoc.terms[term] = {
        count: count,
        freq: count / indexedDoc.termCount,
      };

      // Initialize or update corpus-wide term stats
      if (!this.terms[term]) {
        this.terms[term] = { n: 0, idf: 0 };
      }
      this.terms[term].n++;
    }

    this.documents[id] = indexedDoc;
  }

  public updateIdf(): void {
    for (const term in this.terms) {
      const { n } = this.terms[term];
      // The BM25 IDF formula
      const idf = Math.log(
        (this.totalDocuments - n + 0.5) / (n + 0.5) + 1
      );
      this.terms[term].idf = Math.max(idf, 0); // Ensure non-negative
    }
  }

  public search(query: string): IndexedDocument[] {
    const queryTerms = tokenize(query);
    const results: IndexedDocument[] = [];

    for (const id in this.documents) {
      const doc = this.documents[id];
      doc._score = 0;

      for (const queryTerm of queryTerms) {
        if (!this.terms[queryTerm] || !doc.terms[queryTerm]) {
          continue; // Term not in corpus or not in this doc
        }

        const idf = this.terms[queryTerm].idf;
        const tf = doc.terms[queryTerm].count;
        const docLength = doc.termCount;

        // The core BM25 formula component for a single term
        const numerator = tf * (this.k1 + 1);
        const denominator = tf + this.k1 * (1 - this.b + (this.b * docLength) / this.averageDocumentLength);
        
        doc._score += idf * (numerator / denominator);
      }

      if (doc._score > 0) {
        results.push(doc);
      }
    }

    // Sort by score in descending order
    return results.sort((a, b) => (b._score || 0) - (a._score || 0));
  }
}

// --- Example Usage for Problem-Finder ---
const redditPosts: Document[] = [
    { id: 'p1', body: 'My Node.js backend is slow when processing large files.' },
    { id: 'p2', body: 'How to optimize PostgreSQL query performance with indexing?' },
    { id: 'p3', body: 'Struggling with slow performance in my Express.js server.' },
    { id: 'p4', body: 'What are best practices for database connection pooling?' }
];

const ranker = new BM25Ranker();

// 1. Index all documents
redditPosts.forEach(post => ranker.addDocument(post));

// 2. Calculate IDF scores for the whole corpus
ranker.updateIdf();

// 3. Perform a search
const query = 'slow node performance';
const searchResults = ranker.search(query);

console.log(searchResults.map(r => ({ id: r.id, score: r._score })));
// Expected output: p1 and p3 will be ranked highest.

This class provides a complete, in-memory keyword search engine. You can instantiate it, add your Reddit posts to it, and then execute searches, getting back a list of posts ranked by relevance.

Conclusion

In this lesson, you've built the second, equally important half of your hybrid search system. While our previous lesson focused on what text means, this lesson focused on what it says.

Here are your key takeaways:

• Keyword ranking moves beyond simple string matching by scoring documents based on term statistics.
• TF-IDF is the foundational concept, valuing terms that are frequent in a document but rare in the overall corpus.
• Okapi BM25 is the modern standard. It improves on TF-IDF by intelligently handling term frequency saturation and document length, using the tunable k1 and b parameters.
• Implementing BM25 in memory involves a two-stage process: an indexing step to pre-calculate statistics, and a search step to score documents against a query.

You now have two powerful ranking algorithms at your disposal: a semantic ranker (cosine similarity) and a lexical ranker (BM25). Each has its own strengths. In our next lesson, we will learn how to combine their outputs into a single, superior ranking using a technique called Reciprocal Rank Fusion (RRF). This will complete the core logic of your hybrid search engine.