Vajra BM25

Vajra (Sanskrit: वज्र, "thunderbolt") is a high-performance BM25 search engine. It uses Category Theory abstractions to reframe the BM25 algorithm, providing a well-structured API with vectorized implementations. Benchmarks show Vajra is faster than BM25S (one of the fastest Python BM25 libraries) at larger corpus sizes while achieving better recall on certain datasets.

What Makes Vajra Different

Vajra implements the standard BM25 ranking algorithm using rigorous mathematical abstractions:

• Morphisms: BM25 scoring as a mathematical arrow (Query, Document) → ℝ
• Coalgebras: Search as state unfolding QueryState → List[SearchResult]
• Functors: The List functor captures multiple-results semantics

While Vajra BM25 uses the same underlying mathematics of BM25, it uses different vocabulary to describe the search process, and the abstractions are more amenable to experimentation and improvement. The core BM25 formula is identical to other implementations—category theory provides the organizational structure.

Installation

# Basic installation (zero dependencies)
pip install vajra-bm25

# With optimizations (NumPy + SciPy for vectorized operations)
pip install vajra-bm25[optimized]

# With index persistence (save/load indices)
pip install vajra-bm25[persistence]

# Everything
pip install vajra-bm25[all]

Quick Start

The Python API for using Vajra BM25 is quite straightforward, and there's currently support for using JSONL document corpuses via the DocumentCorpus class.

from vajra_bm25 import VajraSearch, Document, DocumentCorpus

# Create documents
documents = [
    Document(id="1", title="Category Theory", content="Functors preserve structure"),
    Document(id="2", title="Coalgebras", content="Coalgebras model dynamics"),
    Document(id="3", title="Search Algorithms", content="BFS explores level by level"),
]
corpus = DocumentCorpus(documents)

# Create search engine
engine = VajraSearch(corpus)

# Search
results = engine.search("category functors", top_k=5)

for r in results:
    print(f"{r.rank}. {r.document.title} (score: {r.score:.3f})")

Optimized Usage

For larger corpora (1000+ documents), use the optimized version. This optimized version is much faster.

from vajra_bm25 import VajraSearchOptimized, DocumentCorpus

# Load corpus from JSONL
corpus = DocumentCorpus.load_jsonl("corpus.jsonl")

# Create optimized engine
# Automatically uses sparse matrices for >10K documents
engine = VajraSearchOptimized(corpus)

# Search (vectorized, cached)
results = engine.search("neural networks", top_k=10)

Parallel Batch Processing

For high-throughput scenarios, use the parallel batch processing version. This version is much faster and able to return results for multiple queries in parallel. There's obviously the overhead due to parallelism, which may work against the search algorithm, but in cases where we have memory limitations, this may work better than Vajra Search Optimized.

from vajra_bm25 import VajraSearchParallel

engine = VajraSearchParallel(corpus, max_workers=4)

# Process multiple queries in parallel
queries = ["machine learning", "deep learning", "neural networks"]
batch_results = engine.search_batch(queries, top_k=5)

Performance

Benchmarked on synthetic corpora with 10 queries per run, comparing against rank-bm25 and BM25S:

Speed Comparison

Speedup vs rank-bm25

Recall@10 (vs rank-bm25 baseline)

Key observations:

• Sub-millisecond query latency at all corpus sizes
• Up to 307x speedup over rank-bm25 at 100K documents with Vajra Parallel
• Vajra is faster than BM25S at all corpus sizes
• Vajra achieves better recall at 50K docs (80% vs 56%)
• Recall varies by corpus characteristics (vocabulary overlap, document length distribution)

Vajra achieves these speedups through:

• Vectorized NumPy operations
• Pre-computed IDF values
• Sparse matrix representation
• LRU query caching
• Partial sort for top-k
• Thread pool parallelism (VajraSearchParallel)

For detailed benchmark methodology and results, see docs/benchmarks.md.

JSONL Format

Vajra uses JSONL for corpus persistence:

{"id": "doc1", "title": "First Document", "content": "Content here"}
{"id": "doc2", "title": "Second Document", "content": "More content"}

Load and save:

# Save
corpus.save_jsonl("corpus.jsonl")

# Load
corpus = DocumentCorpus.load_jsonl("corpus.jsonl")

BM25 Parameters

from vajra_bm25 import VajraSearch, BM25Parameters

# Custom BM25 parameters
params = BM25Parameters(
    k1=1.5,  # Term frequency saturation (default: 1.5)
    b=0.75   # Length normalization (default: 0.75)
)

engine = VajraSearch(corpus, params=params)

Categorical Abstractions (Advanced)

For users interested in the category theory foundations:

from vajra_bm25 import (
    Morphism, FunctionMorphism, IdentityMorphism,
    Coalgebra, SearchCoalgebra,
    Functor, ListFunctor,
)

# Morphism composition
f = FunctionMorphism(lambda x: x + 1)
g = FunctionMorphism(lambda x: x * 2)
h = f >> g  # h(x) = (x + 1) * 2

# Identity laws
identity = IdentityMorphism()
assert (f >> identity).apply(5) == f.apply(5)  # f . id = f
assert (identity >> f).apply(5) == f.apply(5)  # id . f = f

There's a better, more rigorous treatment of the concepts of Category Theory by Bartosz Milewski here.

API Reference

Core Classes

• Document(id, title, content, metadata=None) - Immutable document
• DocumentCorpus(documents) - Collection of documents
• VajraSearch(corpus, params=None) - Base search engine
• VajraSearchOptimized(corpus, k1=1.5, b=0.75) - Vectorized search
• VajraSearchParallel(corpus, max_workers=4) - Parallel batch search

Search Results

@dataclass
class SearchResult:
    document: Document  # The matched document
    score: float        # BM25 relevance score
    rank: int           # Position in results (1-indexed)

Why Category Theory?

Category theory provides:

1. Unified abstractions - Same Coalgebra.structure_map() interface for graph search and document retrieval
2. Explicit type signatures - BM25: (Query, Document) → ℝ makes inputs/outputs clear
3. Composable pipelines - preprocess >> score >> rank as morphism composition

What it doesn't provide:

• Performance improvements (those come from NumPy/sparse matrices)
• Novel algorithms (BM25 is BM25)
• Runtime machinery (it's just well-organized code)

The honest summary: category theory is a design vocabulary, not a runtime mechanism.

Development

# Clone repository
git clone https://github.com/aiexplorations/vajra_bm25.git
cd vajra_bm25

# Install in development mode
pip install -e ".[dev]"

# Run tests
pytest tests/ -v

# Run with coverage
pytest --cov=vajra_bm25 --cov-report=html

License

MIT License - see LICENSE for details.

Acknowledgments

• BM25 algorithm: Robertson & Zaragoza, "The Probabilistic Relevance Framework"
• Category theory foundations: Rutten, "Universal Coalgebra: A Theory of Systems"
• Built and explored in the State Dynamic Modeling project
• Inspired by the Category Theory lectures by Bartosz Milewski which are here on YouTube.