sparse-embedding¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/sparse-embedding
项目说明¶
Educational Sparse Vector Search Engine¶
An educational implementation of a sparse vector search engine using inverted index and BM25 algorithm. This project demonstrates the fundamental concepts of information retrieval with extensive logging and visualization features for learning purposes.
Features¶
- Full BM25 Implementation: Complete implementation of the BM25 ranking algorithm
- Advanced Tokenization: Comprehensive tokenizer handling numbers, codes, technical terms, and mixed case
- Inverted Index: Efficient inverted index data structure for term lookup
- HTTP API: RESTful API built with FastAPI
- Interactive Web UI: Browser-based interface for indexing and searching
- Extensive Logging: Detailed educational logging throughout indexing and search operations
- Index Visualization: APIs to inspect the internal structure of the index
- In-Memory Storage: Simple in-memory storage for educational purposes
Tokenization Capabilities¶
The TextProcessor now provides comprehensive tokenization for real-world text:
- Numbers:
404,3.14,2.0.1 - Codes:
XK9-2B4-7Q1,API_KEY_123 - Technical Terms:
C++,.NET,Node.js - Mixed Case:
JavaScript,PyTorch,iPhone - Email:
user@example.com - Hex Codes:
#FF5733,0x1234 - Acronyms:
API,HTTP,NASA - Alphanumeric:
Python3,ES6,HTML5
Architecture¶
Core Components¶
- TextProcessor: Advanced tokenizer handling words, numbers, codes, technical terms, and mixed case
- InvertedIndex: Maintains the inverted index structure with term frequencies and document frequencies
- BM25: Implements the BM25 ranking algorithm for relevance scoring
- SparseSearchEngine: Main engine coordinating all components
- HTTP Server: FastAPI-based server exposing the search engine functionality
BM25 Algorithm¶
BM25 (Best Matching 25) is a probabilistic ranking function that scores documents based on query terms. The algorithm uses:
- Term Frequency (TF): How often a term appears in a document
- Inverse Document Frequency (IDF): How rare or common a term is across all documents
- Document Length Normalization: Adjusts scores based on document length
Key parameters:
- k1 (default 1.5): Controls term frequency saturation
- b (default 0.75): Controls length normalization
Installation¶
- Install dependencies:
cli.py（下节的命令行工具）只依赖 Python 标准库，无需安装任何第三方包即可离线运行；server.py / demo.py 才需要上面的 FastAPI 等依赖。
命令行工具 cli.py（实验 3-5，推荐入口）¶
cli.py 提供一个完全离线的命令行入口：在一个内置的 10 篇小型语料上运行 BM25 稀疏检索、复现书中“逐词 IDF/TF/BM25 贡献”的日志，并在带标注的评测集上计算 recall/precision/MRR。所有参数都有中文 --help。
python cli.py --help # 查看全部参数（中文）
python cli.py # 默认演示：查询 "model distillation"
python cli.py -q "model distillation" --explain # 逐词展示 TF/IDF/BM25 贡献（对应书中日志）
python cli.py --eval # 在标注集上计算 recall@k / precision@k / MRR
python cli.py -q "cat" # 观察 BM25 读不懂同义词的短板（kitten/feline 漏召回）
python cli.py --corpus my.json -q "查询" -o out.json # 自定义语料 + 结果落盘
python cli.py --k1 2.0 -b 0.5 -q "..." # 调 BM25 参数 k1 / b
python cli.py --method splade -q "..." # 学习型稀疏检索 SPLADE（需预先下载模型）
主要参数：
| 参数 | 说明 |
|---|---|
-q, --query |
查询字符串（默认 model distillation） |
-c, --corpus |
语料文件（.json 文档数组或 .jsonl 每行一篇）；缺省用内置示例语料 |
-m, --method |
bm25（默认，离线）或 splade（学习型稀疏，需下载模型） |
-k, --top-k |
返回前 k 条（默认 5） |
-o, --output |
把结果 / 评测指标写入 JSON 文件 |
--eval |
在标注集上评测 recall@k / precision@k / MRR |
--labels |
自定义评测标注 {query: [相关doc_id,...]} |
--explain |
对命中文档逐词展示 TF / IDF / BM25 贡献 |
--k1 / -b |
BM25 词频饱和参数 k1、文档长度归一化参数 b |
-v, --verbose |
打开引擎 DEBUG 日志（分词、倒排索引构建、打分全过程） |
检索质量评测（--eval）¶
内置标注集覆盖精确关键词、错误码、专有名称与“只有同义表达”的查询。python cli.py --eval 的真实输出（k=5）：
查询 'model distillation' recall@5=1.00 precision@5=1.00 RR=1.00
查询 'HTTP 404 error' recall@5=1.00 precision@5=0.50 RR=1.00
查询 'XK9-2B4-7Q1' recall@5=1.00 precision@5=1.00 RR=1.00
查询 'BM25 ranking function' recall@5=1.00 precision@5=1.00 RR=1.00
查询 'cat' recall@5=0.00 precision@5=0.00 RR=0.00 <- 漏召回(同义词短板)
宏平均 recall@5=0.800 precision@5=0.700 MRR=0.800 漏召回率(1-recall@5)=0.200
结果直观印证了书中的结论：BM25 在精确关键词、错误码、专有名称上表现极佳（recall=1.0），但读不懂同义词——查询 cat 无法命中只写了 kitten / feline 的文档（recall=0）。这一长一短正是引入混合检索（见实验 3-6 retrieval-pipeline）的动机。
学习型稀疏检索（--method splade）¶
--method splade 对应书中提到的学习型稀疏检索（SPLADE）：用掩码语言模型为每个词项打权重，并能为原文未出现但语义相关的词项补权重（术语扩展）。它需要下载预训练模型 naver/splade-cocondenser-ensembledistil（依赖 torch、transformers）。离线环境下无法下载权重时，命令会快速给出清晰提示（并说明 BM25 路径无需任何模型即可离线运行），不会卡在网络下载上。联网环境可先执行 huggingface-cli download naver/splade-cocondenser-ensembledistil 缓存模型后再运行。
Usage¶
Starting the Server¶
The server will start on http://localhost:4241
Web Interface¶
Open your browser and navigate to http://localhost:4241 to access the interactive web UI.
API Endpoints¶
Index a Document¶
POST /index
{
"text": "Your document text here",
"metadata": {"title": "Document Title", "category": "Category"}
}
Search Documents¶
Get Statistics¶
Get Index Structure¶
Retrieve Document by ID¶
Clear Index¶
Running the Demo¶
The demo script will: 1. Clear any existing index 2. Index sample documents about programming and computer science 3. Display index statistics 4. Show the internal index structure 5. Perform various search queries 6. Demonstrate document retrieval
Educational Features¶
Extensive Logging¶
The system provides detailed logging at every step: - Document tokenization process - Term frequency calculations - IDF score computations - BM25 scoring for each term - Query processing steps - Candidate document identification
Index Visualization¶
The /index/structure endpoint returns:
- Inverted index mapping (terms to documents)
- Document statistics (length, unique terms, top terms)
- BM25 parameters
- Global term frequency distribution
Debug Information in Search Results¶
Each search result includes debug information: - Matched query terms - Document length - Term frequencies for query terms - Individual term contributions to the final score
Example Output¶
Indexing Log¶
2024-01-15 10:30:45 - Indexing document with ID 0
2024-01-15 10:30:45 - Document text: Python is a high-level programming...
2024-01-15 10:30:45 - Tokenizing text of length 142
2024-01-15 10:30:45 - Found 23 raw tokens
2024-01-15 10:30:45 - After removing stop words: 15 tokens
2024-01-15 10:30:45 - Document 0: 15 tokens, 12 unique terms
2024-01-15 10:30:45 - Document 0 indexed successfully
Search Log¶
2024-01-15 10:31:20 - Searching for: 'machine learning algorithms'
2024-01-15 10:31:20 - Query terms after processing: ['machine', 'learning', 'algorithms']
2024-01-15 10:31:20 - Term 'machine' appears in 2 documents
2024-01-15 10:31:20 - Term 'learning' appears in 3 documents
2024-01-15 10:31:20 - Term 'algorithms' appears in 1 document
2024-01-15 10:31:20 - Found 4 candidate documents
2024-01-15 10:31:20 - IDF for 'machine': N=10, df=2, idf=1.7918
2024-01-15 10:31:20 - Term 'machine' in doc 1: tf=2, dl=35, score=3.2451
2024-01-15 10:31:20 - Document 1 total score: 7.8923
2024-01-15 10:31:20 - Returning top 3 results
API Documentation¶
Interactive API documentation is available at http://localhost:4241/docs when the server is running.
Project Structure¶
sparse-embedding/
├── bm25_engine.py # Core search engine implementation
├── cli.py # Offline argparse CLI: BM25/SPLADE search + recall/precision eval
├── server.py # FastAPI HTTP server
├── demo.py # Demonstration script
├── requirements.txt # Python dependencies
└── README.md # This file
Learning Resources¶
This implementation demonstrates: - How inverted indices work in search engines - The mathematics behind BM25 ranking - Term frequency and document frequency concepts - The importance of text preprocessing - How sparse vectors represent documents - RESTful API design for search systems
Limitations¶
This is an educational implementation with some limitations: - In-memory storage (not persistent) - Basic tokenization (could be improved with lemmatization) - English-only stop words - No support for phrase queries - No query expansion or synonyms - Single-threaded processing
Further Improvements¶
Potential enhancements for learning: - Add persistence with a database - Implement more advanced text processing (stemming, lemmatization) - Add support for phrase queries - Implement query expansion - Add multi-language support - Implement index compression techniques - Add support for field-specific searching - Implement more ranking algorithms (TF-IDF, Okapi BM25+)
源代码¶
bm25_engine.py¶
"""
BM25 Sparse Vector Search Engine
An educational implementation of BM25 algorithm with inverted index
"""
import math
import re
import logging
from collections import defaultdict, Counter
from typing import List, Dict, Set, Tuple, Optional
import json
# Configure logging for educational purposes
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
class TextProcessor:
"""Text preprocessing for indexing and searching"""
def __init__(self):
# Common English stop words
self.stop_words = {
'the', 'is', 'at', 'which', 'on', 'a', 'an', 'as', 'are', 'was',
'been', 'be', 'have', 'has', 'had', 'do', 'does', 'did', 'will',
'would', 'could', 'should', 'may', 'might', 'must', 'can', 'this',
'that', 'these', 'those', 'i', 'you', 'he', 'she', 'it', 'we',
'they', 'what', 'who', 'when', 'where', 'why', 'how', 'all', 'each',
'every', 'both', 'few', 'more', 'most', 'other', 'some', 'such',
'only', 'own', 'same', 'so', 'than', 'too', 'very', 'just'
}
logger.info(f"TextProcessor initialized with {len(self.stop_words)} stop words")
def tokenize(self, text: str, remove_stop_words: bool = True) -> List[str]:
"""Tokenize text into words, numbers, and codes.
Handles:
- Words (preserving case for acronyms)
- Numbers (404, 500, 3.14)
- Codes (XK9-2B4-7Q1, API_KEY, user@example.com)
- Technical terms (C++, .NET, Node.js)
"""
logger.debug(f"Tokenizing text of length {len(text)}")
# Comprehensive tokenization patterns
patterns = [
# Email addresses (keep whole)
r'\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b',
# URLs (simplified)
r'https?://[^\s]+',
# API keys, codes with hyphens/underscores (e.g., XK9-2B4-7Q1, API_KEY_123)
r'\b[A-Z0-9]+(?:[_-][A-Z0-9]+)+\b',
# Technical terms with special chars (C++, C#, .NET)
r'\b[A-Z]\+\+|\b[A-Z]#|\.[A-Z]+[a-zA-Z]*',
# Version numbers (3.14, 2.0.1)
r'\b\d+(?:\.\d+)+\b',
# Hex codes (#FF5733, 0x1234)
r'#[0-9A-Fa-f]{3,8}\b|0x[0-9A-Fa-f]+\b',
# Numbers (including decimals)
r'\b\d+(?:\.\d+)?\b',
# Acronyms and uppercase words (USA, NASA, API)
r'\b[A-Z]{2,}\b',
# Mixed case words (JavaScript, PyTorch)
r'\b[A-Z][a-z]+[A-Z][a-zA-Z]*\b',
# Alphanumeric combinations (Python3, ES6, 3DS)
r'\b[A-Za-z]+\d+\b|\b\d+[A-Za-z]+\b',
# Regular words (including apostrophes)
r"\b[a-zA-Z]+(?:'[a-z]+)?\b",
]
# Combine all patterns
combined_pattern = '|'.join(f'({p})' for p in patterns)
# Extract all tokens
raw_tokens = re.findall(combined_pattern, text, re.IGNORECASE)
# Flatten the results (findall with groups returns tuples)
tokens = []
for match in raw_tokens:
token = next(t for t in match if t) # Get the non-empty match
# Preserve case for:
# - All uppercase words (API, USA)
# - Mixed case (JavaScript, PyTorch)
# - Codes with special chars
# - Numbers
# - Alphanumeric combinations (Python3, ES6)
if (token.isupper() and len(token) > 1) or \
any(c.isupper() for c in token[1:]) or \
any(c in '-_@.#+' for c in token) or \
any(c.isdigit() for c in token) or \
token.startswith('.'):
tokens.append(token)
else:
# Convert regular words to lowercase
tokens.append(token.lower())
logger.debug(f"Found {len(tokens)} raw tokens")
if remove_stop_words:
# Only remove stop words from lowercase word tokens
filtered_tokens = []
for token in tokens:
# Keep if: not a lowercase word, or not in stop words
if not token.islower() or token not in self.stop_words:
filtered_tokens.append(token)
tokens = filtered_tokens
logger.debug(f"After removing stop words: {len(tokens)} tokens")
return tokens
class InvertedIndex:
"""Inverted index data structure for efficient term lookup"""
def __init__(self):
# Main inverted index: term -> set of document IDs
self.index: Dict[str, Set[int]] = defaultdict(set)
# Document frequency: term -> number of documents containing term
self.document_frequency: Dict[str, int] = defaultdict(int)
# Term frequency in documents: doc_id -> term -> frequency
self.term_frequency: Dict[int, Counter] = {}
# Document lengths (number of terms)
self.doc_lengths: Dict[int, int] = {}
# Original documents for retrieval
self.documents: Dict[int, str] = {}
# Document metadata
self.doc_metadata: Dict[int, Dict] = {}
# Statistics
self.total_documents = 0
self.total_terms = 0
self.unique_terms = 0
logger.info("InvertedIndex initialized")
def add_document(self, doc_id: int, text: str, metadata: Optional[Dict] = None):
"""Add a document to the index"""
logger.info(f"Adding document {doc_id} to index")
logger.debug(f"Document text: {text[:100]}..." if len(text) > 100 else f"Document text: {text}")
# Store original document
self.documents[doc_id] = text
if metadata:
self.doc_metadata[doc_id] = metadata
logger.debug(f"Document metadata: {metadata}")
# Process text
processor = TextProcessor()
tokens = processor.tokenize(text)
# Count term frequencies
term_freq = Counter(tokens)
self.term_frequency[doc_id] = term_freq
self.doc_lengths[doc_id] = len(tokens)
logger.debug(f"Document {doc_id}: {len(tokens)} tokens, {len(term_freq)} unique terms")
# Update inverted index
for term in term_freq:
self.index[term].add(doc_id)
# Update document frequency
for term in term_freq:
if doc_id not in self.index[term]:
self.document_frequency[term] += 1
self.total_documents += 1
self._update_statistics()
logger.info(f"Document {doc_id} indexed successfully")
def _update_statistics(self):
"""Update index statistics"""
self.unique_terms = len(self.index)
self.total_terms = sum(self.doc_lengths.values())
logger.debug(f"Index statistics: {self.total_documents} documents, "
f"{self.unique_terms} unique terms, {self.total_terms} total terms")
def get_posting_list(self, term: str) -> Set[int]:
"""Get document IDs containing the term"""
return self.index.get(term, set())
def get_statistics(self) -> Dict:
"""Get comprehensive index statistics"""
stats = {
'total_documents': self.total_documents,
'unique_terms': self.unique_terms,
'total_terms': self.total_terms,
'average_document_length': self.total_terms / self.total_documents if self.total_documents > 0 else 0,
'terms_by_frequency': self._get_term_frequency_distribution()
}
return stats
def _get_term_frequency_distribution(self, top_n: int = 10) -> List[Tuple[str, int]]:
"""Get top N most frequent terms across all documents"""
global_term_freq = Counter()
for doc_term_freq in self.term_frequency.values():
global_term_freq.update(doc_term_freq)
return global_term_freq.most_common(top_n)
def get_index_structure(self) -> Dict:
"""Get a visualization-friendly representation of the index"""
structure = {
'inverted_index': {},
'document_info': {},
'statistics': self.get_statistics()
}
# Include top terms in the structure
for term, doc_ids in list(self.index.items())[:20]: # Limit to 20 terms for readability
structure['inverted_index'][term] = {
'document_ids': list(doc_ids),
'document_frequency': len(doc_ids)
}
# Include document information
for doc_id in self.documents:
structure['document_info'][doc_id] = {
'length': self.doc_lengths[doc_id],
'unique_terms': len(self.term_frequency[doc_id]),
'top_terms': self.term_frequency[doc_id].most_common(5)
}
return structure
class BM25:
"""BM25 ranking algorithm implementation"""
def __init__(self, index: InvertedIndex, k1: float = 1.5, b: float = 0.75):
"""
Initialize BM25 with tuning parameters
k1: controls term frequency saturation (typically 1.2 to 2.0)
b: controls length normalization (0.0 to 1.0)
"""
self.index = index
self.k1 = k1
self.b = b
# Calculate average document length
self.avgdl = 0
if index.total_documents > 0:
self.avgdl = sum(index.doc_lengths.values()) / index.total_documents
logger.info(f"BM25 initialized with k1={k1}, b={b}, avgdl={self.avgdl:.2f}")
def calculate_idf(self, term: str) -> float:
"""Calculate Inverse Document Frequency for a term"""
N = self.index.total_documents
df = len(self.index.get_posting_list(term))
if df == 0:
return 0
# BM25 IDF formula
idf = math.log((N - df + 0.5) / (df + 0.5) + 1)
logger.debug(f"IDF for '{term}': N={N}, df={df}, idf={idf:.4f}")
return idf
def calculate_term_score(self, term: str, doc_id: int) -> float:
"""Calculate BM25 score for a single term in a document"""
# Get term frequency in document
tf = self.index.term_frequency.get(doc_id, Counter()).get(term, 0)
if tf == 0:
return 0
# Get document length
dl = self.index.doc_lengths.get(doc_id, 0)
# Calculate IDF
idf = self.calculate_idf(term)
# BM25 term score formula
numerator = tf * (self.k1 + 1)
denominator = tf + self.k1 * (1 - self.b + self.b * (dl / self.avgdl))
score = idf * (numerator / denominator)
logger.debug(f"Term '{term}' in doc {doc_id}: tf={tf}, dl={dl}, score={score:.4f}")
return score
def score_document(self, query_terms: List[str], doc_id: int) -> float:
"""Calculate total BM25 score for a document given query terms"""
total_score = 0
term_scores = {}
for term in query_terms:
term_score = self.calculate_term_score(term, doc_id)
term_scores[term] = term_score
total_score += term_score
logger.debug(f"Document {doc_id} total score: {total_score:.4f}")
logger.debug(f"Term contributions: {term_scores}")
return total_score
def search(self, query: str, top_k: int = 10) -> List[Tuple[int, float, Dict]]:
"""
Search for documents matching the query
Returns list of (doc_id, score, debug_info) tuples
"""
logger.info(f"Searching for: '{query}'")
# Process query
processor = TextProcessor()
query_terms = processor.tokenize(query)
logger.info(f"Query terms after processing: {query_terms}")
# Find candidate documents (documents containing at least one query term)
candidate_docs = set()
term_doc_mapping = {}
for term in query_terms:
# Try exact match first
docs = self.index.get_posting_list(term)
# If no exact match and term is not a number/code, try lowercase
if not docs and not term[0].isdigit() and '-' not in term:
docs = self.index.get_posting_list(term.lower())
candidate_docs.update(docs)
term_doc_mapping[term] = docs
logger.debug(f"Term '{term}' appears in {len(docs)} documents")
logger.info(f"Found {len(candidate_docs)} candidate documents")
# Score each candidate document
doc_scores = []
for doc_id in candidate_docs:
score = self.score_document(query_terms, doc_id)
# Collect debug information
debug_info = {
'matched_terms': [term for term in query_terms
if doc_id in self.index.get_posting_list(term)],
'doc_length': self.index.doc_lengths[doc_id],
'term_frequencies': {term: self.index.term_frequency[doc_id].get(term, 0)
for term in query_terms}
}
doc_scores.append((doc_id, score, debug_info))
# Sort by score (descending)
doc_scores.sort(key=lambda x: x[1], reverse=True)
# Return top k results
results = doc_scores[:top_k]
logger.info(f"Returning top {len(results)} results")
for rank, (doc_id, score, _) in enumerate(results, 1):
logger.info(f"Rank {rank}: Document {doc_id} (score: {score:.4f})")
return results
class SparseSearchEngine:
"""Main search engine combining all components"""
def __init__(self):
self.index = InvertedIndex()
self.bm25 = None
self.next_doc_id = 0
# Map external doc_id to internal doc_id
self.external_to_internal = {}
# Map internal doc_id to external doc_id
self.internal_to_external = {}
logger.info("SparseSearchEngine initialized")
def index_document(self, text: str, metadata: Optional[Dict] = None, external_doc_id: Optional[str] = None) -> str:
"""Index a new document and return its ID"""
# Generate internal ID
internal_doc_id = self.next_doc_id
self.next_doc_id += 1
# Use external_doc_id if provided, otherwise use internal ID as string
if external_doc_id:
doc_id_str = external_doc_id
else:
doc_id_str = str(internal_doc_id)
# Store mappings
self.external_to_internal[doc_id_str] = internal_doc_id
self.internal_to_external[internal_doc_id] = doc_id_str
logger.info(f"Indexing document with external ID '{doc_id_str}' (internal ID {internal_doc_id})")
self.index.add_document(internal_doc_id, text, metadata)
# Reinitialize BM25 with updated index
self.bm25 = BM25(self.index)
return doc_id_str
def index_batch(self, documents: List[Dict]) -> List[str]:
"""Index multiple documents at once"""
logger.info(f"Batch indexing {len(documents)} documents")
doc_ids = []
for doc in documents:
text = doc.get('text', '')
metadata = doc.get('metadata', None)
external_doc_id = doc.get('doc_id', None)
doc_id = self.index_document(text, metadata, external_doc_id)
doc_ids.append(doc_id)
logger.info(f"Batch indexing complete. Indexed {len(doc_ids)} documents")
return doc_ids
def search(self, query: str, top_k: int = 10) -> List[Dict]:
"""Search for documents matching the query"""
if self.bm25 is None:
logger.warning("No documents indexed yet")
return []
logger.info(f"Executing search query: '{query}'")
results = self.bm25.search(query, top_k)
# Format results
formatted_results = []
for internal_doc_id, score, debug_info in results:
# Get external doc_id
external_doc_id = self.internal_to_external.get(internal_doc_id, str(internal_doc_id))
result = {
'doc_id': external_doc_id, # Use external doc_id
'score': score,
'text': self.index.documents[internal_doc_id],
'metadata': self.index.doc_metadata.get(internal_doc_id, {}),
'debug': debug_info
}
formatted_results.append(result)
return formatted_results
def get_document(self, doc_id) -> Optional[Dict]:
"""Retrieve a document by ID (can be internal or external)"""
# Check if it's an external doc_id
if isinstance(doc_id, str) and doc_id in self.external_to_internal:
internal_id = self.external_to_internal[doc_id]
elif isinstance(doc_id, int) and doc_id in self.index.documents:
internal_id = doc_id
doc_id = self.internal_to_external.get(internal_id, str(internal_id))
else:
return None
return {
'doc_id': doc_id, # Return external doc_id
'text': self.index.documents[internal_id],
'metadata': self.index.doc_metadata.get(internal_id, {}),
'statistics': {
'length': self.index.doc_lengths[internal_id],
'unique_terms': len(self.index.term_frequency[internal_id]),
'top_terms': self.index.term_frequency[internal_id].most_common(10)
}
}
def get_index_info(self) -> Dict:
"""Get comprehensive information about the index"""
return {
'statistics': self.index.get_statistics(),
'structure': self.index.get_index_structure(),
'bm25_params': {
'k1': self.bm25.k1 if self.bm25 else None,
'b': self.bm25.b if self.bm25 else None,
'avgdl': self.bm25.avgdl if self.bm25 else None
}
}
def clear_index(self):
"""Clear all indexed documents"""
logger.info("Clearing index")
self.index = InvertedIndex()
self.bm25 = None
self.next_doc_id = 0
self.external_to_internal = {}
self.internal_to_external = {}
logger.info("Index cleared")
cli.py¶
#!/usr/bin/env python3
"""
稀疏检索命令行工具（实验 3-5）
在一个小型示例语料上运行 BM25 稀疏检索，支持：
- 自定义语料 / 查询 / top-k / 输出文件
- --explain 复现书中"逐词 IDF / TF / BM25 贡献"的日志
- --eval 在带标注的小型评测集上计算 recall@k / precision@k / MRR
- --method splade 学习型稀疏检索（需要下载模型，离线环境会给出提示）
不带任何参数运行时，等价于书中实验 3-5 的默认演示（查询"model distillation"）。
"""
import argparse
import json
import logging
import sys
from typing import Dict, List, Optional, Set, Tuple
from bm25_engine import SparseSearchEngine
# ---------------------------------------------------------------------------
# 内置示例语料与标注（英文，与引擎的分词器能力一致，可完全离线复现）
# 语料刻意混合了：普通词、专有代码、技术缩写、以及只有"同义表达"的文档，
# 用来同时展示 BM25 在精确关键词匹配上的强项与在同义词上的短板。
# ---------------------------------------------------------------------------
DEFAULT_CORPUS: List[Dict] = [
{"doc_id": "doc_1", "title": "Python Language",
"text": "Python is a high-level programming language known for readability and a simple syntax."},
{"doc_id": "doc_2", "title": "JavaScript Runtime",
"text": "JavaScript runs in the browser and on servers via Node.js for full-stack web development."},
{"doc_id": "doc_3", "title": "Model Distillation",
"text": "Model distillation compresses a large teacher model into a smaller student model while preserving accuracy."},
{"doc_id": "doc_4", "title": "Knowledge Distillation",
"text": "Knowledge distillation transfers knowledge from a big neural network to a compact model for efficient inference."},
{"doc_id": "doc_5", "title": "BM25 Ranking",
"text": "BM25 is a probabilistic ranking function using term frequency and inverse document frequency."},
{"doc_id": "doc_6", "title": "HTTP Errors",
"text": "The HTTP 404 error code means the requested resource was not found on the web server."},
{"doc_id": "doc_7", "title": "A Playful Kitten",
"text": "A cute kitten chased a ball of yarn across the living room floor all afternoon."},
{"doc_id": "doc_8", "title": "Silent Hunter",
"text": "The feline predator stalked its prey silently through the tall grass at dusk."},
{"doc_id": "doc_9", "title": "Hardware Fault",
"text": "Error code XK9-2B4-7Q1 indicates a hardware fault in the storage controller board."},
{"doc_id": "doc_10", "title": "Transformers",
"text": "Transformer models use self-attention to process input sequences in parallel efficiently."},
]
# query -> 相关文档 doc_id 集合（人工标注的 ground truth）
DEFAULT_LABELS: Dict[str, List[str]] = {
"model distillation": ["doc_3", "doc_4"],
"HTTP 404 error": ["doc_6"],
"XK9-2B4-7Q1": ["doc_9"],
"BM25 ranking function": ["doc_5"],
# 相关文档用 kitten / feline 表达"猫"，故意不含字面 "cat"，
# 用于演示稀疏检索读不懂同义词的短板（BM25 会漏召回）。
"cat": ["doc_7", "doc_8"],
}
DEFAULT_QUERY = "model distillation"
def _quiet_logging(verbose: bool) -> None:
"""默认压低引擎日志；--verbose / --explain 时放开到 DEBUG 以展示计算过程。"""
level = logging.DEBUG if verbose else logging.WARNING
logging.getLogger().setLevel(level)
logging.getLogger("bm25_engine").setLevel(level)
def load_corpus(path: Optional[str]) -> List[Dict]:
"""加载语料。支持 .json（文档数组）与 .jsonl（每行一个文档）。"""
if not path:
return DEFAULT_CORPUS
docs: List[Dict] = []
with open(path, "r", encoding="utf-8") as f:
if path.endswith(".jsonl"):
for line in f:
line = line.strip()
if line:
docs.append(json.loads(line))
else:
data = json.load(f)
docs = data["documents"] if isinstance(data, dict) else data
if not docs:
raise ValueError(f"语料文件为空：{path}")
return docs
def load_labels(path: Optional[str]) -> Dict[str, List[str]]:
"""加载评测标注：{query: [relevant_doc_id, ...]}。"""
if not path:
return DEFAULT_LABELS
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def build_engine(corpus: List[Dict], k1: float, b: float) -> SparseSearchEngine:
"""把语料灌进引擎；用给定的 k1/b 重建 BM25。"""
engine = SparseSearchEngine()
engine.index_batch([
{"text": d["text"],
"doc_id": d.get("doc_id"),
"metadata": {"title": d.get("title", "")}}
for d in corpus
])
# index_batch 内部每篇都会重建 BM25，这里再显式用目标参数固定一次
from bm25_engine import BM25
engine.bm25 = BM25(engine.index, k1=k1, b=b)
return engine
def explain_result(engine: SparseSearchEngine, query: str, doc_id: str) -> List[Tuple[str, int, int, float, float]]:
"""复现书中日志：对命中文档，逐个查询词给出 TF / 文档长度 / IDF / BM25 贡献。"""
from bm25_engine import TextProcessor
internal = engine.external_to_internal[doc_id]
terms = TextProcessor().tokenize(query)
rows = []
for term in terms:
tf = engine.index.term_frequency[internal].get(term, 0)
if tf == 0:
continue
dl = engine.index.doc_lengths[internal]
idf = engine.bm25.calculate_idf(term)
contrib = engine.bm25.calculate_term_score(term, internal)
rows.append((term, tf, dl, idf, contrib))
return rows
def run_search(engine: SparseSearchEngine, query: str, top_k: int,
explain: bool) -> List[Dict]:
"""执行单条查询并打印结果，返回结构化结果供 --output 落盘。"""
results = engine.search(query, top_k=top_k)
print(f"\n查询: '{query}' (BM25, top-{top_k})")
print("-" * 60)
if not results:
print(" 没有命中任何文档（所有查询词都不在倒排索引中）。")
return []
out = []
for rank, r in enumerate(results, 1):
title = r["metadata"].get("title", "")
print(f" #{rank} {r['doc_id']} score={r['score']:.4f} {title}")
print(f" 命中词: {r['debug']['matched_terms']}")
print(f" 预览: {r['text'][:80]}...")
if explain:
rows = explain_result(engine, query, r["doc_id"])
for term, tf, dl, idf, contrib in rows:
print(f" └ '{term}': TF={tf}, 文档长度={dl}词, "
f"IDF={idf:.4f}, BM25贡献={contrib:.4f}")
out.append({
"rank": rank,
"doc_id": r["doc_id"],
"score": r["score"],
"title": title,
"matched_terms": r["debug"]["matched_terms"],
})
return out
def _metrics_for_query(retrieved: List[str], relevant: Set[str], k: int) -> Dict:
"""单条查询的 recall@k / precision@k / 命中排名（用于 MRR）。"""
topk = retrieved[:k]
hits = [d for d in topk if d in relevant]
recall = len(set(hits)) / len(relevant) if relevant else 0.0
precision = len(hits) / len(topk) if topk else 0.0
rr = 0.0
for i, d in enumerate(retrieved, 1):
if d in relevant:
rr = 1.0 / i
break
return {"recall": recall, "precision": precision, "rr": rr,
"hits": hits, "retrieved": topk}
def run_eval(engine: SparseSearchEngine, labels: Dict[str, List[str]],
k: int) -> Dict:
"""在标注集上做检索评测，打印每条查询指标 + 宏平均。"""
print(f"\n{'='*60}")
print(f"检索质量评测 (recall@{k} / precision@{k} / MRR)")
print(f"{'='*60}")
per_query = {}
sum_recall = sum_prec = sum_rr = 0.0
for query, rel_list in labels.items():
relevant = set(rel_list)
results = engine.search(query, top_k=max(k, 10))
retrieved = [r["doc_id"] for r in results]
m = _metrics_for_query(retrieved, relevant, k)
per_query[query] = m
sum_recall += m["recall"]
sum_prec += m["precision"]
sum_rr += m["rr"]
flag = "" if m["recall"] > 0 else " <- 漏召回(同义词短板)" if query == "cat" else " <- 漏召回"
print(f"\n查询 '{query}' 相关文档={sorted(relevant)}")
print(f" 召回排序: {retrieved[:k]}")
print(f" recall@{k}={m['recall']:.2f} precision@{k}={m['precision']:.2f} RR={m['rr']:.2f}{flag}")
n = len(labels)
macro = {
"recall@k": sum_recall / n,
"precision@k": sum_prec / n,
"mrr": sum_rr / n,
"miss_rate@k": 1.0 - sum_recall / n,
}
print(f"\n{'-'*60}")
print(f"宏平均 recall@{k}={macro['recall@k']:.3f} "
f"precision@{k}={macro['precision@k']:.3f} "
f"MRR={macro['mrr']:.3f} 漏召回率(1-recall@{k})={macro['miss_rate@k']:.3f}")
return {"k": k, "per_query": {q: {kk: vv for kk, vv in m.items() if kk != "retrieved"}
for q, m in per_query.items()},
"macro": macro}
def run_splade(query: str, corpus: List[Dict], top_k: int) -> Optional[List[Dict]]:
"""学习型稀疏检索（SPLADE）。需要 transformers + torch + 预训练模型。
离线环境无法下载模型时，会打印清晰提示并返回 None（不影响参数解析验证）。
"""
model_name = "naver/splade-cocondenser-ensembledistil"
try:
import torch # noqa: F401
from transformers import AutoModelForMaskedLM, AutoTokenizer
except Exception as e:
print("\n[SPLADE] 需要依赖 transformers 与 torch，当前环境缺失：", e)
print(" 安装：pip install torch transformers")
print(" （BM25 路径无需任何模型，可完全离线运行）")
return None
try:
# 只用本地缓存加载，避免离线环境卡在无休止的网络下载上。
print(f"\n[SPLADE] 尝试从本地缓存加载模型 {model_name} ...")
tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
model = AutoModelForMaskedLM.from_pretrained(model_name, local_files_only=True)
model.eval()
except Exception:
print(f"\n[SPLADE] 本地缓存中没有模型 {model_name}，且离线环境无法下载权重。")
print(" 请先在联网环境执行一次以下命令把模型缓存到本地，再重跑本命令：")
print(f" huggingface-cli download {model_name}")
print(" （BM25 路径不依赖任何模型，可完全离线复现书中实验 3-5）")
return None
import torch
def encode(text: str) -> Dict[str, float]:
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256)
with torch.no_grad():
logits = model(**inputs).logits # [1, seq, vocab]
# SPLADE: log(1+ReLU(logits)) 后在序列维做 max-pool，得到词表维稀疏权重
weights = torch.max(
torch.log1p(torch.relu(logits)) * inputs["attention_mask"].unsqueeze(-1),
dim=1,
).values.squeeze(0)
nz = torch.nonzero(weights).squeeze(-1)
return {int(i): float(weights[i]) for i in nz}
q_vec = encode(query)
scored = []
for d in corpus:
d_vec = encode(d["text"])
score = sum(w * d_vec.get(t, 0.0) for t, w in q_vec.items())
scored.append((d.get("doc_id"), score, d.get("title", "")))
scored.sort(key=lambda x: x[1], reverse=True)
print(f"\n查询: '{query}' (SPLADE, top-{top_k})")
print("-" * 60)
out = []
for rank, (doc_id, score, title) in enumerate(scored[:top_k], 1):
print(f" #{rank} {doc_id} score={score:.4f} {title}")
out.append({"rank": rank, "doc_id": doc_id, "score": score, "title": title})
return out
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="cli.py",
description="稀疏检索命令行工具（实验 3-5）：在小型语料上运行 BM25 / SPLADE 稀疏检索并评测检索质量。",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""示例：
python cli.py # 默认演示（查询 "model distillation"）
python cli.py -q "HTTP 404 error" --explain # 展示逐词 TF/IDF/BM25 贡献
python cli.py --eval # 在标注集上算 recall/precision/MRR
python cli.py -q "cat" # 观察 BM25 的同义词短板
python cli.py --corpus my.json -q "..." -o out.json
python cli.py --method splade -q "model distillation" # 学习型稀疏检索（需模型）
""",
)
parser.add_argument("-q", "--query", default=DEFAULT_QUERY,
help=f"查询字符串（默认: '{DEFAULT_QUERY}'）")
parser.add_argument("-c", "--corpus", default=None,
help="语料文件路径（.json 文档数组 或 .jsonl 每行一篇）；缺省用内置示例语料")
parser.add_argument("-m", "--method", choices=["bm25", "splade"], default="bm25",
help="检索方法：bm25(默认,离线) 或 splade(学习型稀疏,需下载模型)")
parser.add_argument("-k", "--top-k", type=int, default=5,
help="返回前 k 条结果（默认: 5）")
parser.add_argument("-o", "--output", default=None,
help="把结果/评测指标以 JSON 写入该文件")
parser.add_argument("--eval", action="store_true",
help="在标注集上评测 recall@k / precision@k / MRR，而非只跑单条查询")
parser.add_argument("--labels", default=None,
help="评测标注文件 {query: [相关doc_id,...]}；缺省用内置标注")
parser.add_argument("--explain", action="store_true",
help="对每条命中文档展示逐词 TF/IDF/BM25 贡献（复现书中日志）")
parser.add_argument("--k1", type=float, default=1.5,
help="BM25 词频饱和参数 k1（默认: 1.5）")
parser.add_argument("-b", "--b", type=float, default=0.75,
help="BM25 文档长度归一化参数 b（默认: 0.75）")
parser.add_argument("-v", "--verbose", action="store_true",
help="打开引擎 DEBUG 日志（展示分词、倒排索引构建、打分全过程）")
return parser
def main(argv: Optional[List[str]] = None) -> int:
args = build_parser().parse_args(argv)
_quiet_logging(args.verbose)
corpus = load_corpus(args.corpus)
print(f"已加载语料：{len(corpus)} 篇文档"
+ ("（内置示例）" if not args.corpus else f"（来自 {args.corpus}）"))
payload: Dict = {"method": args.method, "query": args.query, "top_k": args.top_k}
if args.method == "splade":
results = run_splade(args.query, corpus, args.top_k)
if results is None:
return 0 # 已给出模型缺失提示，视为正常退出
payload["results"] = results
else:
engine = build_engine(corpus, k1=args.k1, b=args.b)
print(f"BM25 参数：k1={args.k1}, b={args.b}, avgdl={engine.bm25.avgdl:.2f}")
if args.eval:
labels = load_labels(args.labels)
payload["eval"] = run_eval(engine, labels, args.top_k)
else:
payload["results"] = run_search(engine, args.query, args.top_k, args.explain)
if args.output:
with open(args.output, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
print(f"\n已写入结果：{args.output}")
return 0
if __name__ == "__main__":
sys.exit(main())
config.py¶
"""Configuration for the sparse embedding service."""
import os
# Server configuration
SPARSE_PORT = int(os.getenv("SPARSE_PORT", "4241")) # Port 4241 to avoid conflicts
SPARSE_HOST = os.getenv("SPARSE_HOST", "0.0.0.0")
# Logging
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
demo.py¶
"""
Demo script for the Educational Sparse Vector Search Engine
Shows how to use the engine with sample documents and queries
"""
import requests
import json
import time
import logging
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# Server URL
BASE_URL = "http://localhost:4241"
def wait_for_server(max_attempts=10):
"""Wait for server to be ready"""
logger.info("Waiting for server to be ready...")
for i in range(max_attempts):
try:
response = requests.get(f"{BASE_URL}/stats")
if response.status_code == 200:
logger.info("Server is ready!")
return True
except:
pass
time.sleep(1)
return False
def clear_index():
"""Clear the index before demo"""
logger.info("Clearing existing index...")
response = requests.delete(f"{BASE_URL}/index")
if response.status_code == 200:
logger.info("Index cleared successfully")
return response.json()
def index_sample_documents():
"""Index a collection of sample documents"""
logger.info("\n" + "="*50)
logger.info("INDEXING SAMPLE DOCUMENTS")
logger.info("="*50)
sample_documents = [
{
"text": "Python is a high-level programming language known for its simplicity and readability. It supports multiple programming paradigms including procedural, object-oriented, and functional programming.",
"metadata": {"title": "Python Programming", "category": "programming"}
},
{
"text": "Machine learning is a subset of artificial intelligence that enables computers to learn from data without being explicitly programmed. It uses algorithms to identify patterns and make decisions.",
"metadata": {"title": "Introduction to Machine Learning", "category": "AI"}
},
{
"text": "Natural language processing (NLP) is a field of AI that focuses on the interaction between computers and human language. It involves tasks like text classification, sentiment analysis, and machine translation.",
"metadata": {"title": "NLP Basics", "category": "AI"}
},
{
"text": "Data structures are fundamental concepts in computer science that organize and store data efficiently. Common data structures include arrays, linked lists, trees, graphs, and hash tables.",
"metadata": {"title": "Data Structures Overview", "category": "computer science"}
},
{
"text": "JavaScript is a dynamic programming language commonly used for web development. It runs in browsers and on servers with Node.js, making it versatile for full-stack development.",
"metadata": {"title": "JavaScript Essentials", "category": "programming"}
},
{
"text": "Deep learning is a subset of machine learning that uses neural networks with multiple layers. It has achieved breakthrough results in computer vision, speech recognition, and natural language processing.",
"metadata": {"title": "Deep Learning Introduction", "category": "AI"}
},
{
"text": "Algorithms are step-by-step procedures for solving computational problems. Algorithm analysis involves studying their time and space complexity using Big O notation.",
"metadata": {"title": "Algorithm Analysis", "category": "computer science"}
},
{
"text": "Web development involves creating websites and web applications using technologies like HTML, CSS, JavaScript, and various frameworks. Modern web development often uses React, Vue, or Angular for frontend development.",
"metadata": {"title": "Modern Web Development", "category": "web"}
},
{
"text": "Databases are systems for storing and managing data. Relational databases use SQL and tables, while NoSQL databases offer flexible schemas for unstructured data. Popular choices include PostgreSQL, MongoDB, and Redis.",
"metadata": {"title": "Database Systems", "category": "databases"}
},
{
"text": "Cloud computing provides on-demand computing resources over the internet. Major providers like AWS, Google Cloud, and Azure offer services for storage, computation, and machine learning in the cloud.",
"metadata": {"title": "Cloud Computing Basics", "category": "cloud"}
}
]
# Index documents
doc_ids = []
for i, doc in enumerate(sample_documents, 1):
logger.info(f"\nIndexing document {i}/{len(sample_documents)}: {doc['metadata']['title']}")
response = requests.post(
f"{BASE_URL}/index",
json={"text": doc["text"], "metadata": doc["metadata"]}
)
if response.status_code == 200:
result = response.json()
doc_ids.append(result["doc_id"])
logger.info(f"✓ Indexed with ID: {result['doc_id']}")
else:
logger.error(f"✗ Failed to index document")
logger.info(f"\nSuccessfully indexed {len(doc_ids)} documents")
return doc_ids
def show_statistics():
"""Display index statistics"""
logger.info("\n" + "="*50)
logger.info("INDEX STATISTICS")
logger.info("="*50)
response = requests.get(f"{BASE_URL}/stats")
if response.status_code == 200:
stats = response.json()
logger.info(f"Total documents: {stats['total_documents']}")
logger.info(f"Unique terms: {stats['unique_terms']}")
logger.info(f"Total terms: {stats['total_terms']}")
logger.info(f"Average document length: {stats['average_document_length']:.2f}")
if 'terms_by_frequency' in stats:
logger.info("\nTop 10 most frequent terms:")
for term, freq in stats['terms_by_frequency']:
logger.info(f" - {term}: {freq} occurrences")
return stats
def perform_searches():
"""Perform various search queries to demonstrate the engine"""
logger.info("\n" + "="*50)
logger.info("PERFORMING SEARCHES")
logger.info("="*50)
search_queries = [
("machine learning algorithms", 3),
("programming language", 5),
("database SQL", 3),
("web development JavaScript", 3),
("artificial intelligence", 5),
("data structures algorithms", 3),
("cloud computing AWS", 3),
("neural networks deep learning", 3)
]
for query, top_k in search_queries:
logger.info(f"\n{'─'*40}")
logger.info(f"Query: '{query}' (top {top_k} results)")
logger.info('─'*40)
response = requests.post(
f"{BASE_URL}/search",
json={"query": query, "top_k": top_k}
)
if response.status_code == 200:
results = response.json()
if not results:
logger.info("No results found")
else:
for rank, result in enumerate(results, 1):
logger.info(f"\n Rank {rank}:")
logger.info(f" Score: {result['score']:.4f}")
logger.info(f" Title: {result['metadata'].get('title', 'N/A')}")
logger.info(f" Category: {result['metadata'].get('category', 'N/A')}")
logger.info(f" Text preview: {result['text'][:100]}...")
logger.info(f" Matched terms: {result['debug']['matched_terms']}")
logger.info(f" Document length: {result['debug']['doc_length']} terms")
else:
logger.error(f"Search failed: {response.status_code}")
time.sleep(0.5) # Small delay between searches
def show_index_structure():
"""Display the internal structure of the index"""
logger.info("\n" + "="*50)
logger.info("INDEX STRUCTURE VISUALIZATION")
logger.info("="*50)
response = requests.get(f"{BASE_URL}/index/structure")
if response.status_code == 200:
data = response.json()
# Show BM25 parameters
logger.info("\nBM25 Parameters:")
params = data['bm25_params']
logger.info(f" k1 (term frequency saturation): {params['k1']}")
logger.info(f" b (length normalization): {params['b']}")
logger.info(f" avgdl (average document length): {params['avgdl']:.2f}")
# The actual structure is nested under 'structure' key
structure = data.get('structure', {})
# Show sample of inverted index
logger.info("\nSample of Inverted Index (first 5 terms):")
inv_index = structure.get('inverted_index', {})
if inv_index:
for i, (term, info) in enumerate(list(inv_index.items())[:5]):
logger.info(f" '{term}':")
logger.info(f" - Document frequency: {info['document_frequency']}")
logger.info(f" - Appears in documents: {info['document_ids']}")
else:
logger.info(" No inverted index data available")
# Show document information
logger.info("\nDocument Information:")
doc_info = structure.get('document_info', {})
if doc_info:
for doc_id, info in list(doc_info.items())[:3]: # Show first 3 documents
logger.info(f" Document {doc_id}:")
logger.info(f" - Length: {info['length']} terms")
logger.info(f" - Unique terms: {info['unique_terms']}")
logger.info(f" - Top terms: {[f'{term}({freq})' for term, freq in info['top_terms'][:5]]}")
else:
logger.info(" No document information available")
return data
def test_specific_document_retrieval():
"""Test retrieving specific documents by ID"""
logger.info("\n" + "="*50)
logger.info("DOCUMENT RETRIEVAL TEST")
logger.info("="*50)
# Retrieve document with ID 0
doc_id = 0
logger.info(f"\nRetrieving document with ID {doc_id}...")
response = requests.get(f"{BASE_URL}/document/{doc_id}")
if response.status_code == 200:
document = response.json()
logger.info(f"Document {doc_id}:")
logger.info(f" Title: {document['metadata'].get('title', 'N/A')}")
logger.info(f" Category: {document['metadata'].get('category', 'N/A')}")
logger.info(f" Text: {document['text'][:150]}...")
else:
logger.error(f"Failed to retrieve document: {response.status_code}")
def main():
"""Run the complete demo"""
logger.info("Starting Educational Sparse Vector Search Engine Demo")
logger.info("Make sure the server is running (python server.py)")
# Wait for server
if not wait_for_server():
logger.error("Server is not responding. Please start the server first.")
return
# Clear existing index
clear_index()
# Index sample documents
doc_ids = index_sample_documents()
# Show statistics
show_statistics()
# Show index structure
show_index_structure()
# Perform searches
perform_searches()
# Test document retrieval
test_specific_document_retrieval()
logger.info("\n" + "="*50)
logger.info("DEMO COMPLETED")
logger.info("="*50)
logger.info("\nVisit http://localhost:8000 in your browser for the interactive UI")
logger.info("API documentation available at http://localhost:8000/docs")
if __name__ == "__main__":
main()
quickstart.py¶
#!/usr/bin/env python3
"""
Quick start script for the Educational Sparse Vector Search Engine
Demonstrates basic usage in a simple, interactive way
"""
import logging
from bm25_engine import SparseSearchEngine
# Configure logging to show educational information
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(message)s'
)
logger = logging.getLogger(__name__)
def main():
print("\n" + "="*60)
print(" Educational Sparse Vector Search Engine - Quick Start")
print("="*60)
print("\nThis demo shows the core functionality of BM25 search.\n")
# Initialize the search engine
print("Initializing search engine...")
engine = SparseSearchEngine()
# Sample documents about different programming topics
documents = [
{
"text": "Python is a versatile programming language widely used for web development, data science, machine learning, and automation. Its simple syntax makes it ideal for beginners.",
"title": "Python Overview"
},
{
"text": "JavaScript powers the interactive web. It runs in browsers and on servers with Node.js. Modern JavaScript includes features like async/await, arrow functions, and destructuring.",
"title": "JavaScript Essentials"
},
{
"text": "Machine learning algorithms enable computers to learn from data. Popular algorithms include linear regression, decision trees, neural networks, and support vector machines.",
"title": "ML Algorithms"
},
{
"text": "Web development involves HTML for structure, CSS for styling, and JavaScript for interactivity. Modern frameworks like React, Vue, and Angular simplify complex applications.",
"title": "Web Development"
},
{
"text": "Data structures organize information efficiently. Arrays provide fast access, linked lists enable dynamic sizing, trees support hierarchical data, and hash tables offer constant-time lookups.",
"title": "Data Structures"
},
{
"text": "Databases store and manage data persistently. SQL databases like PostgreSQL use structured tables, while NoSQL databases like MongoDB store flexible documents.",
"title": "Database Systems"
},
{
"text": "Cloud computing provides scalable infrastructure on demand. AWS, Google Cloud, and Azure offer services for compute, storage, networking, and machine learning.",
"title": "Cloud Computing"
},
{
"text": "Software testing ensures code quality. Unit tests verify individual functions, integration tests check component interactions, and end-to-end tests validate entire workflows.",
"title": "Software Testing"
},
{
"text": "Version control systems track code changes over time. Git is the most popular system, enabling collaboration through branches, commits, and pull requests.",
"title": "Version Control"
},
{
"text": "APIs (Application Programming Interfaces) enable communication between software systems. REST APIs use HTTP methods, while GraphQL provides flexible data querying.",
"title": "APIs and Integration"
}
]
# Index documents
print(f"\nIndexing {len(documents)} documents...")
print("-" * 40)
for i, doc in enumerate(documents):
doc_id = engine.index_document(doc["text"], {"title": doc["title"]})
print(f" [{doc_id}] {doc['title']}")
print(f"\n✓ Indexed {len(documents)} documents successfully!")
# Show index statistics
stats = engine.index.get_statistics()
print(f"\nIndex Statistics:")
print(f" • Total documents: {stats['total_documents']}")
print(f" • Unique terms: {stats['unique_terms']}")
print(f" • Average document length: {stats['average_document_length']:.1f} terms")
# Demonstrate searches
print("\n" + "="*60)
print(" Demonstration Searches")
print("="*60)
queries = [
"machine learning algorithms",
"web development JavaScript",
"database SQL NoSQL",
"cloud computing AWS",
"Python programming"
]
for query in queries:
print(f"\n🔍 Query: '{query}'")
print("-" * 40)
results = engine.search(query, top_k=3)
if results:
for rank, result in enumerate(results, 1):
title = result['metadata'].get('title', 'Unknown')
score = result['score']
matched = result['debug']['matched_terms']
print(f"\n #{rank} {title} (Score: {score:.3f})")
print(f" Matched terms: {', '.join(matched)}")
print(f" Preview: {result['text'][:100]}...")
else:
print(" No results found")
# Interactive search
print("\n" + "="*60)
print(" Interactive Search")
print("="*60)
print("\nNow you can try your own searches!")
print("Type 'quit' to exit, 'stats' for statistics, or enter a search query.\n")
while True:
try:
query = input("Enter search query: ").strip()
if query.lower() == 'quit':
print("\nThank you for using the Educational Sparse Vector Search Engine!")
break
if query.lower() == 'stats':
stats = engine.index.get_statistics()
print(f"\nCurrent Index Statistics:")
print(f" • Documents: {stats['total_documents']}")
print(f" • Unique terms: {stats['unique_terms']}")
print(f" • Total terms: {stats['total_terms']}")
print(f" • Top terms: {', '.join([t[0] for t in stats['terms_by_frequency'][:5]])}")
print()
continue
if not query:
continue
# Perform search
results = engine.search(query, top_k=5)
if results:
print(f"\nFound {len(results)} results for '{query}':\n")
for rank, result in enumerate(results, 1):
title = result['metadata'].get('title', 'Unknown')
score = result['score']
matched = result['debug']['matched_terms']
print(f" #{rank} {title}")
print(f" Score: {score:.4f}")
print(f" Matched: {', '.join(matched) if matched else 'None'}")
print(f" Text: {result['text'][:150]}...")
print()
else:
print(f"\nNo results found for '{query}'")
print("Try different keywords or check your spelling.\n")
except KeyboardInterrupt:
print("\n\nExiting...")
break
except Exception as e:
print(f"Error: {e}")
continue
print("\n" + "="*60)
print("\nTo learn more:")
print(" • Run 'python test_engine.py' to see comprehensive tests")
print(" • Run 'python server.py' to start the HTTP API server")
print(" • Run 'python demo.py' for a full demonstration")
print(" • Check the README.md for detailed documentation")
print("\n" + "="*60)
if __name__ == "__main__":
main()
server.py¶
"""
HTTP API Server for Sparse Vector Search Engine
Educational server with extensive logging and visualization
"""
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, Field
from typing import List, Dict, Optional
import uvicorn
import logging
import json
from datetime import datetime
from bm25_engine import SparseSearchEngine
# Configure logging
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# Initialize FastAPI app
app = FastAPI(
title="Educational Sparse Vector Search Engine",
description="BM25-based search engine with inverted index for educational purposes",
version="1.0.0"
)
# Initialize search engine
search_engine = SparseSearchEngine()
# Pydantic models for request/response
class IndexDocumentRequest(BaseModel):
text: str = Field(..., description="Text content to index")
metadata: Optional[Dict] = Field(None, description="Optional metadata")
doc_id: Optional[str] = Field(None, description="Optional external document ID")
class BatchIndexRequest(BaseModel):
documents: List[Dict] = Field(..., description="List of documents to index")
class SearchRequest(BaseModel):
query: str = Field(..., description="Search query")
top_k: int = Field(10, description="Number of results to return")
class DocumentResponse(BaseModel):
doc_id: str # Changed to str to support external IDs
text: str
metadata: Optional[Dict]
score: Optional[float] = None
debug: Optional[Dict] = None
# Root endpoint with UI
@app.get("/", response_class=HTMLResponse)
async def root():
"""Serve a simple HTML interface for the search engine"""
html_content = """
<!DOCTYPE html>
<html>
<head>
<title>Educational Sparse Vector Search Engine</title>
<style>
body {
font-family: Arial, sans-serif;
max-width: 1200px;
margin: 0 auto;
padding: 20px;
background-color: #f5f5f5;
}
h1 {
color: #333;
border-bottom: 2px solid #007bff;
padding-bottom: 10px;
}
.section {
background: white;
border-radius: 8px;
padding: 20px;
margin: 20px 0;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.form-group {
margin: 15px 0;
}
label {
display: block;
margin-bottom: 5px;
font-weight: bold;
color: #555;
}
input, textarea {
width: 100%;
padding: 8px;
border: 1px solid #ddd;
border-radius: 4px;
box-sizing: border-box;
}
button {
background-color: #007bff;
color: white;
padding: 10px 20px;
border: none;
border-radius: 4px;
cursor: pointer;
font-size: 16px;
}
button:hover {
background-color: #0056b3;
}
pre {
background-color: #f8f9fa;
padding: 15px;
border-radius: 4px;
overflow-x: auto;
border: 1px solid #dee2e6;
}
.results {
margin-top: 20px;
}
.result-item {
background: #f8f9fa;
padding: 15px;
margin: 10px 0;
border-radius: 4px;
border-left: 4px solid #007bff;
}
.score {
font-weight: bold;
color: #007bff;
}
.debug-info {
margin-top: 10px;
padding: 10px;
background: #fff;
border: 1px solid #ddd;
border-radius: 4px;
font-size: 0.9em;
}
.stats {
display: grid;
grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
gap: 15px;
margin-top: 15px;
}
.stat-item {
background: #f8f9fa;
padding: 10px;
border-radius: 4px;
text-align: center;
}
.stat-value {
font-size: 24px;
font-weight: bold;
color: #007bff;
}
.stat-label {
font-size: 14px;
color: #666;
margin-top: 5px;
}
</style>
</head>
<body>
<h1>🔍 Educational Sparse Vector Search Engine</h1>
<div class="section">
<h2>Index Documents</h2>
<div class="form-group">
<label for="indexText">Document Text:</label>
<textarea id="indexText" rows="4" placeholder="Enter document text to index..."></textarea>
</div>
<div class="form-group">
<label for="indexMetadata">Metadata (JSON, optional):</label>
<input id="indexMetadata" placeholder='{"title": "Document Title", "author": "Author Name"}'>
</div>
<button onclick="indexDocument()">Index Document</button>
</div>
<div class="section">
<h2>Search</h2>
<div class="form-group">
<label for="searchQuery">Query:</label>
<input id="searchQuery" placeholder="Enter search query...">
</div>
<div class="form-group">
<label for="topK">Number of Results:</label>
<input id="topK" type="number" value="5" min="1" max="100">
</div>
<button onclick="search()">Search</button>
<div id="searchResults" class="results"></div>
</div>
<div class="section">
<h2>Index Statistics</h2>
<button onclick="loadStatistics()">Load Statistics</button>
<div id="statistics"></div>
</div>
<div class="section">
<h2>Index Structure Visualization</h2>
<button onclick="loadIndexStructure()">Load Index Structure</button>
<div id="indexStructure"></div>
</div>
<script>
async function indexDocument() {
const text = document.getElementById('indexText').value;
const metadataStr = document.getElementById('indexMetadata').value;
let metadata = null;
if (metadataStr) {
try {
metadata = JSON.parse(metadataStr);
} catch (e) {
alert('Invalid JSON in metadata field');
return;
}
}
const response = await fetch('/index', {
method: 'POST',
headers: {'Content-Type': 'application/json'},
body: JSON.stringify({text: text, metadata: metadata, doc_id: null})
});
if (response.ok) {
const result = await response.json();
alert(`Document indexed successfully! ID: ${result.doc_id}`);
document.getElementById('indexText').value = '';
document.getElementById('indexMetadata').value = '';
loadStatistics();
} else {
alert('Error indexing document');
}
}
async function search() {
const query = document.getElementById('searchQuery').value;
const topK = document.getElementById('topK').value;
const response = await fetch('/search', {
method: 'POST',
headers: {'Content-Type': 'application/json'},
body: JSON.stringify({query: query, top_k: parseInt(topK)})
});
if (response.ok) {
const results = await response.json();
displaySearchResults(results);
} else {
alert('Error performing search');
}
}
function displaySearchResults(results) {
const container = document.getElementById('searchResults');
if (results.length === 0) {
container.innerHTML = '<p>No results found</p>';
return;
}
let html = '<h3>Search Results</h3>';
results.forEach((result, index) => {
html += `
<div class="result-item">
<div><strong>Rank ${index + 1}</strong> - Doc ID: ${result.doc_id}</div>
<div class="score">Score: ${result.score.toFixed(4)}</div>
<div style="margin-top: 10px;">${result.text}</div>
${result.metadata ? `<div style="margin-top: 10px;"><strong>Metadata:</strong> ${JSON.stringify(result.metadata)}</div>` : ''}
<div class="debug-info">
<strong>Debug Info:</strong>
<pre>${JSON.stringify(result.debug, null, 2)}</pre>
</div>
</div>
`;
});
container.innerHTML = html;
}
async function loadStatistics() {
const response = await fetch('/stats');
if (response.ok) {
const stats = await response.json();
displayStatistics(stats);
}
}
function displayStatistics(stats) {
const container = document.getElementById('statistics');
let html = '<div class="stats">';
html += `
<div class="stat-item">
<div class="stat-value">${stats.total_documents}</div>
<div class="stat-label">Total Documents</div>
</div>
<div class="stat-item">
<div class="stat-value">${stats.unique_terms}</div>
<div class="stat-label">Unique Terms</div>
</div>
<div class="stat-item">
<div class="stat-value">${stats.total_terms}</div>
<div class="stat-label">Total Terms</div>
</div>
<div class="stat-item">
<div class="stat-value">${stats.average_document_length.toFixed(2)}</div>
<div class="stat-label">Avg Doc Length</div>
</div>
`;
html += '</div>';
if (stats.terms_by_frequency && stats.terms_by_frequency.length > 0) {
html += '<h4>Top Terms by Frequency</h4>';
html += '<ul>';
stats.terms_by_frequency.forEach(([term, freq]) => {
html += `<li>${term}: ${freq}</li>`;
});
html += '</ul>';
}
container.innerHTML = html;
}
async function loadIndexStructure() {
const response = await fetch('/index/structure');
if (response.ok) {
const structure = await response.json();
displayIndexStructure(structure);
}
}
function displayIndexStructure(structure) {
const container = document.getElementById('indexStructure');
let html = '<h4>Inverted Index Sample (Top Terms)</h4>';
html += '<pre>' + JSON.stringify(structure.inverted_index, null, 2) + '</pre>';
html += '<h4>Document Information</h4>';
html += '<pre>' + JSON.stringify(structure.document_info, null, 2) + '</pre>';
html += '<h4>BM25 Parameters</h4>';
html += '<pre>' + JSON.stringify(structure.bm25_params, null, 2) + '</pre>';
container.innerHTML = html;
}
// Load statistics on page load
window.onload = function() {
loadStatistics();
};
</script>
</body>
</html>
"""
return html_content
@app.post("/index", response_model=Dict)
async def index_document(request: IndexDocumentRequest):
"""Index a single document"""
logger.info(f"Received index request for document of length {len(request.text)}")
if request.doc_id:
logger.info(f"External doc_id provided: {request.doc_id}")
try:
# Extract doc_id from metadata if not provided directly
external_doc_id = request.doc_id
if not external_doc_id and request.metadata and 'doc_id' in request.metadata:
external_doc_id = request.metadata['doc_id']
doc_id = search_engine.index_document(request.text, request.metadata, external_doc_id)
logger.info(f"Document indexed successfully with ID {doc_id}")
return {
"success": True,
"doc_id": doc_id,
"message": f"Document indexed successfully with ID {doc_id}"
}
except Exception as e:
logger.error(f"Error indexing document: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/index/batch", response_model=Dict)
async def index_batch(request: BatchIndexRequest):
"""Index multiple documents at once"""
logger.info(f"Received batch index request for {len(request.documents)} documents")
try:
doc_ids = search_engine.index_batch(request.documents)
logger.info(f"Batch indexing successful: {len(doc_ids)} documents indexed")
return {
"success": True,
"doc_ids": doc_ids,
"message": f"Successfully indexed {len(doc_ids)} documents"
}
except Exception as e:
logger.error(f"Error in batch indexing: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/search", response_model=List[DocumentResponse])
async def search(request: SearchRequest):
"""Search for documents"""
logger.info(f"Received search request: '{request.query}' (top_k={request.top_k})")
try:
results = search_engine.search(request.query, request.top_k)
logger.info(f"Search completed, returning {len(results)} results")
return results
except Exception as e:
logger.error(f"Error performing search: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/document/{doc_id}", response_model=DocumentResponse)
async def get_document(doc_id: str):
"""Retrieve a specific document by ID"""
logger.info(f"Retrieving document {doc_id}")
document = search_engine.get_document(doc_id)
if document is None:
logger.warning(f"Document {doc_id} not found")
raise HTTPException(status_code=404, detail=f"Document {doc_id} not found")
logger.info(f"Document {doc_id} retrieved successfully")
return document
@app.get("/stats", response_model=Dict)
async def get_statistics():
"""Get index statistics"""
logger.info("Retrieving index statistics")
stats = search_engine.index.get_statistics()
logger.info(f"Statistics retrieved: {stats['total_documents']} documents, "
f"{stats['unique_terms']} unique terms")
return stats
@app.get("/index/structure", response_model=Dict)
async def get_index_structure():
"""Get detailed index structure for visualization"""
logger.info("Retrieving index structure")
info = search_engine.get_index_info()
logger.info("Index structure retrieved successfully")
return info
@app.delete("/index", response_model=Dict)
async def clear_index():
"""Clear all indexed documents"""
logger.warning("Clearing entire index")
search_engine.clear_index()
logger.info("Index cleared successfully")
return {
"success": True,
"message": "Index cleared successfully"
}
@app.get("/logs", response_model=Dict)
async def get_recent_logs(lines: int = Query(100, description="Number of log lines to retrieve")):
"""Get recent application logs for educational purposes"""
# This is a simplified version - in production you'd read from a log file
return {
"message": "Logs are being written to console. Check terminal for detailed logs.",
"log_level": "DEBUG",
"description": "Educational logging is enabled. All indexing and search operations are logged."
}
if __name__ == "__main__":
import sys
# Allow overriding port from command line
port = 4241 # Default to 4241 to avoid conflicts with common services
if len(sys.argv) > 1:
try:
port = int(sys.argv[1])
except ValueError:
pass
logger.info("Starting Educational Sparse Vector Search Engine Server")
logger.info(f"Server will run on http://localhost:{port}")
logger.info(f"Visit http://localhost:{port} for the web interface")
logger.info(f"API documentation available at http://localhost:{port}/docs")
uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")
test_engine.py¶
"""
Test script for the Educational Sparse Vector Search Engine
Tests core functionality and demonstrates educational aspects
"""
import logging
from bm25_engine import TextProcessor, InvertedIndex, BM25, SparseSearchEngine
# Configure logging
logging.basicConfig(
level=logging.DEBUG,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
def test_text_processor():
"""Test text processing functionality"""
logger.info("\n" + "="*50)
logger.info("TESTING TEXT PROCESSOR")
logger.info("="*50)
processor = TextProcessor()
# Test basic tokenization
text1 = "The quick brown fox jumps over the lazy dog."
tokens1 = processor.tokenize(text1, remove_stop_words=False)
logger.info(f"Input: {text1}")
logger.info(f"Tokens (with stop words): {tokens1}")
tokens2 = processor.tokenize(text1, remove_stop_words=True)
logger.info(f"Tokens (without stop words): {tokens2}")
# Test with technical text
text2 = "Machine learning algorithms process data to identify patterns."
tokens3 = processor.tokenize(text2)
logger.info(f"\nInput: {text2}")
logger.info(f"Tokens: {tokens3}")
# Test with mixed case and punctuation
text3 = "Python, JavaScript, and C++ are POPULAR programming languages!"
tokens4 = processor.tokenize(text3)
logger.info(f"\nInput: {text3}")
logger.info(f"Tokens: {tokens4}")
assert len(tokens2) < len(tokens1), "Stop word removal should reduce token count"
logger.info("\n✓ Text processor tests passed")
def test_inverted_index():
"""Test inverted index functionality"""
logger.info("\n" + "="*50)
logger.info("TESTING INVERTED INDEX")
logger.info("="*50)
index = InvertedIndex()
# Add test documents
doc1 = "Python is a programming language"
doc2 = "JavaScript is also a programming language"
doc3 = "Python and JavaScript are both popular"
logger.info("\nAdding documents to index...")
index.add_document(0, doc1, {"title": "Doc1"})
index.add_document(1, doc2, {"title": "Doc2"})
index.add_document(2, doc3, {"title": "Doc3"})
# Test posting lists
logger.info("\nTesting posting lists:")
python_docs = index.get_posting_list("python")
logger.info(f"Documents containing 'python': {python_docs}")
assert python_docs == {0, 2}, "Python should be in docs 0 and 2"
programming_docs = index.get_posting_list("programming")
logger.info(f"Documents containing 'programming': {programming_docs}")
assert programming_docs == {0, 1}, "Programming should be in docs 0 and 1"
# Test statistics
stats = index.get_statistics()
logger.info(f"\nIndex statistics:")
logger.info(f" Total documents: {stats['total_documents']}")
logger.info(f" Unique terms: {stats['unique_terms']}")
logger.info(f" Average doc length: {stats['average_document_length']:.2f}")
assert stats['total_documents'] == 3, "Should have 3 documents"
logger.info("\n✓ Inverted index tests passed")
def test_bm25_scoring():
"""Test BM25 scoring algorithm"""
logger.info("\n" + "="*50)
logger.info("TESTING BM25 SCORING")
logger.info("="*50)
# Create index with test documents
index = InvertedIndex()
index.add_document(0, "The cat sat on the mat", {"title": "Simple"})
index.add_document(1, "The dog sat on the log", {"title": "Similar"})
index.add_document(2, "Cats and dogs are pets", {"title": "Pets"})
index.add_document(3, "The mat was comfortable", {"title": "Mat"})
# Initialize BM25
bm25 = BM25(index, k1=1.5, b=0.75)
# Test IDF calculation
logger.info("\nTesting IDF calculations:")
idf_cat = bm25.calculate_idf("cat")
idf_the = bm25.calculate_idf("the") # Common word
idf_mat = bm25.calculate_idf("mat")
logger.info(f"IDF('cat'): {idf_cat:.4f}")
logger.info(f"IDF('the'): {idf_the:.4f}")
logger.info(f"IDF('mat'): {idf_mat:.4f}")
# IDF of rare words should be higher than common words
assert idf_cat > idf_the, "Rare words should have higher IDF"
# Test document scoring
logger.info("\nTesting document scoring for query 'cat mat':")
query_terms = ["cat", "mat"]
for doc_id in range(4):
score = bm25.score_document(query_terms, doc_id)
logger.info(f"Document {doc_id} score: {score:.4f}")
# Document 0 should have the highest score as it contains both terms
score_0 = bm25.score_document(query_terms, 0)
score_1 = bm25.score_document(query_terms, 1)
assert score_0 > score_1, "Doc with both terms should score higher"
logger.info("\n✓ BM25 scoring tests passed")
def test_search_engine():
"""Test the complete search engine"""
logger.info("\n" + "="*50)
logger.info("TESTING SEARCH ENGINE")
logger.info("="*50)
engine = SparseSearchEngine()
# Index test documents
logger.info("\nIndexing test documents...")
doc_ids = []
test_docs = [
("Information retrieval is the science of searching for information in documents.",
{"topic": "IR"}),
("Search engines use inverted indices to quickly find relevant documents.",
{"topic": "Search"}),
("BM25 is a probabilistic ranking function used in information retrieval.",
{"topic": "BM25"}),
("The inverted index maps terms to the documents that contain them.",
{"topic": "Index"}),
("Relevance ranking determines the order of search results.",
{"topic": "Ranking"})
]
for text, metadata in test_docs:
doc_id = engine.index_document(text, metadata)
doc_ids.append(doc_id)
logger.info(f"Indexed: {metadata['topic']} (ID: {doc_id})")
# Test search queries
test_queries = [
("information retrieval", 3),
("inverted index", 2),
("search ranking", 3),
("BM25 algorithm", 2)
]
for query, top_k in test_queries:
logger.info(f"\nSearching for: '{query}' (top {top_k})")
results = engine.search(query, top_k)
for rank, result in enumerate(results, 1):
logger.info(f" Rank {rank}: {result['metadata']['topic']} "
f"(score: {result['score']:.4f}, "
f"matched: {result['debug']['matched_terms']})")
# Test document retrieval
logger.info("\nTesting document retrieval:")
doc = engine.get_document(0)
assert doc is not None, "Should retrieve document"
logger.info(f"Retrieved document 0: {doc['metadata']['topic']}")
# Test index clearing
logger.info("\nTesting index clearing:")
initial_stats = engine.index.get_statistics()
engine.clear_index()
final_stats = engine.index.get_statistics()
assert final_stats['total_documents'] == 0, "Index should be empty after clearing"
logger.info("✓ Index cleared successfully")
logger.info("\n✓ Search engine tests passed")
def test_edge_cases():
"""Test edge cases and special scenarios"""
logger.info("\n" + "="*50)
logger.info("TESTING EDGE CASES")
logger.info("="*50)
engine = SparseSearchEngine()
# Test empty query
logger.info("\nTesting empty query:")
results = engine.search("", top_k=5)
assert len(results) == 0, "Empty query should return no results"
logger.info("✓ Empty query handled correctly")
# Test query with only stop words
logger.info("\nTesting query with only stop words:")
engine.index_document("This is a test document about nothing specific.")
results = engine.search("the is a", top_k=5)
logger.info(f"Results for stop words query: {len(results)} documents")
# Test single word document
logger.info("\nTesting single word document:")
doc_id = engine.index_document("Python")
doc = engine.get_document(doc_id)
assert doc is not None, "Should index single word document"
logger.info("✓ Single word document indexed")
# Test duplicate documents
logger.info("\nTesting duplicate documents:")
text = "This is a duplicate document"
id1 = engine.index_document(text)
id2 = engine.index_document(text)
assert id1 != id2, "Duplicate documents should have different IDs"
logger.info(f"✓ Duplicate documents have different IDs: {id1}, {id2}")
# Test very long document
logger.info("\nTesting very long document:")
long_text = " ".join(["word" + str(i) for i in range(1000)])
long_doc_id = engine.index_document(long_text)
long_doc = engine.get_document(long_doc_id)
logger.info(f"✓ Long document indexed (length: {long_doc['statistics']['length']} terms)")
# Test special characters
logger.info("\nTesting special characters:")
special_text = "Email: test@example.com, URL: https://example.com, Price: $99.99"
special_id = engine.index_document(special_text)
results = engine.search("email test example", top_k=1)
logger.info(f"✓ Special characters handled, found {len(results)} results")
logger.info("\n✓ All edge case tests passed")
def test_ranking_quality():
"""Test the quality of search result ranking"""
logger.info("\n" + "="*50)
logger.info("TESTING RANKING QUALITY")
logger.info("="*50)
engine = SparseSearchEngine()
# Create documents with varying relevance
docs = [
"Machine learning is a subset of artificial intelligence",
"Deep learning uses neural networks for machine learning",
"Machine learning algorithms learn from data",
"Artificial intelligence includes machine learning and robotics",
"Data science often uses machine learning techniques",
"Neural networks are inspired by biological brains",
"Supervised learning is a type of machine learning",
"Unsupervised learning discovers patterns in data",
"Reinforcement learning uses rewards and penalties",
"Computer vision is an application of deep learning"
]
logger.info("Indexing documents about machine learning...")
for i, doc in enumerate(docs):
engine.index_document(doc, {"id": i})
# Search for "machine learning"
query = "machine learning"
logger.info(f"\nSearching for: '{query}'")
results = engine.search(query, top_k=5)
logger.info("\nTop 5 results:")
for rank, result in enumerate(results, 1):
logger.info(f" Rank {rank}: Score {result['score']:.4f}")
logger.info(f" Text: {result['text']}")
logger.info(f" Term frequencies: {result['debug']['term_frequencies']}")
# Verify that documents with both terms rank higher
first_result_text = results[0]['text'].lower()
assert 'machine' in first_result_text and 'learning' in first_result_text, \
"Top result should contain both query terms"
logger.info("\n✓ Ranking quality test passed")
def main():
"""Run all tests"""
logger.info("Starting Educational Sparse Vector Search Engine Tests")
logger.info("This will test all components and demonstrate educational logging")
try:
test_text_processor()
test_inverted_index()
test_bm25_scoring()
test_search_engine()
test_edge_cases()
test_ranking_quality()
logger.info("\n" + "="*50)
logger.info("ALL TESTS PASSED SUCCESSFULLY!")
logger.info("="*50)
logger.info("\nThe educational sparse vector search engine is working correctly.")
logger.info("Run 'python server.py' to start the HTTP server.")
logger.info("Run 'python demo.py' to see a full demonstration.")
except Exception as e:
logger.error(f"Test failed: {str(e)}", exc_info=True)
raise
if __name__ == "__main__":
main()