ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,347 @@
|
||||
# Sparse Vector Search Engine (BM25) / 稀疏向量搜索引擎（BM25）
|
||||
|
||||
> Companion material for *AI Agents in Depth*, Chapter 3 — **Experiment 3-5**: educational BM25 / inverted-index sparse search with offline CLI evaluation.
|
||||
> 配套《深入理解 AI Agent》第 3 章 **实验 3-5**：从零理解 BM25 与倒排索引，含可离线评测的 CLI。
|
||||
|
||||
← [Chapter 3 index / 返回第 3 章目录](../README.md)
|
||||
|
||||
---
|
||||
|
||||
## English
|
||||
|
||||
### Overview
|
||||
|
||||
An educational sparse vector search engine using an inverted index and BM25. It demonstrates core IR concepts with extensive logging and visualization.
|
||||
|
||||
### Features
|
||||
|
||||
- **Full BM25 implementation**
|
||||
- **Advanced tokenization**: numbers, codes, technical terms, mixed case
|
||||
- **Inverted index** for term lookup
|
||||
- **HTTP API** (FastAPI)
|
||||
- **Interactive Web UI** for index/search
|
||||
- **Educational logging** through index and search
|
||||
- **Index visualization** APIs
|
||||
- **In-memory storage** (educational simplicity)
|
||||
|
||||
#### Tokenization capabilities
|
||||
|
||||
- **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**: `#FF5733`, `0x1234`
|
||||
- **Acronyms**: `API`, `HTTP`, `NASA`
|
||||
- **Alphanumeric**: `Python3`, `ES6`, `HTML5`
|
||||
|
||||
### Architecture
|
||||
|
||||
1. **TextProcessor**: tokenizer for words, numbers, codes, technical terms, mixed case
|
||||
2. **InvertedIndex**: term/document frequencies
|
||||
3. **BM25**: ranking
|
||||
4. **SparseSearchEngine**: orchestration
|
||||
5. **HTTP Server**: FastAPI surface
|
||||
|
||||
BM25 uses TF, IDF, and document-length normalization. Key params: `k1` (default 1.5), `b` (default 0.75).
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
# From the repository root: use the shared Chapter 3 environment
|
||||
uv sync --locked --python 3.12 --extra ch3
|
||||
|
||||
# Activate it before changing directories:
|
||||
# macOS/Linux:
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell: .venv\Scripts\Activate.ps1
|
||||
# Windows cmd: .venv\Scripts\activate.bat
|
||||
|
||||
# pip fallback when uv is not installed:
|
||||
# python -m pip install -e ".[ch3]"
|
||||
|
||||
cd chapter3/sparse-embedding
|
||||
|
||||
# Single-project compatibility path, still supported during migration:
|
||||
# python -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
`cli.py` (below) uses only the Python standard library and runs offline with no third-party packages; `server.py` / `demo.py` need the shared `ch3` environment or the compatibility `requirements.txt` path.
|
||||
|
||||
### CLI tool `cli.py` (Experiment 3-5, recommended entry)
|
||||
|
||||
Fully offline CLI: BM25 on a built-in 10-doc corpus, per-term TF/IDF/BM25 contribution logs (as in the book), and labelled recall/precision/MRR. All flags have Chinese `--help`.
|
||||
|
||||
```bash
|
||||
python cli.py --help # all flags (Chinese)
|
||||
python cli.py # default demo query "model distillation"
|
||||
python cli.py -q "model distillation" --explain # per-term TF/IDF/BM25
|
||||
python cli.py --eval # recall@k / precision@k / MRR
|
||||
python cli.py -q "cat" # synonym failure (kitten/feline miss)
|
||||
python cli.py --corpus my.json -q "查询" -o out.json
|
||||
python cli.py --k1 2.0 -b 0.5 -q "..."
|
||||
python cli.py --method splade -q "..." # SPLADE (needs downloaded model)
|
||||
```
|
||||
|
||||
| Flag | Description |
|
||||
| --- | --- |
|
||||
| `-q, --query` | Query string (default `model distillation`) |
|
||||
| `-c, --corpus` | Corpus (`.json` array or `.jsonl`); default built-in sample |
|
||||
| `-m, --method` | `bm25` (default, offline) or `splade` (learned sparse; needs model) |
|
||||
| `-k, --top-k` | Top-k (default 5) |
|
||||
| `-o, --output` | Write results/metrics JSON |
|
||||
| `--eval` | Evaluate recall@k / precision@k / MRR on labels |
|
||||
| `--labels` | Custom labels `{query: [doc_id,...]}` |
|
||||
| `--explain` | Per-term TF / IDF / BM25 on hits |
|
||||
| `--k1` / `-b` | BM25 k1 and b |
|
||||
| `-v, --verbose` | Engine DEBUG logs |
|
||||
|
||||
#### Retrieval quality (`--eval`)
|
||||
|
||||
Built-in labels cover exact keywords, error codes, proper names, and synonym-only queries. Real `python cli.py --eval` output (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 excels on exact keywords, codes, and names (recall=1.0) but misses synonyms—query `cat` does not hit docs that only say `kitten` / `feline`. That gap motivates hybrid search (Experiment 3-6 `retrieval-pipeline`).
|
||||
|
||||
#### Learned sparse (`--method splade`)
|
||||
|
||||
SPLADE weights terms with a masked LM and can expand semantically related terms. Needs pretrained `naver/splade-cocondenser-ensembledistil` (`torch`, `transformers`). Offline without weights, the command fails fast with a clear message (BM25 path needs no model). Online: `huggingface-cli download naver/splade-cocondenser-ensembledistil` then run.
|
||||
|
||||
### Server usage
|
||||
|
||||
```bash
|
||||
python server.py
|
||||
```
|
||||
|
||||
Server: `http://localhost:4241`. Web UI: open that URL. API docs: `http://localhost:4241/docs`.
|
||||
|
||||
#### API endpoints
|
||||
|
||||
```bash
|
||||
POST /index
|
||||
{
|
||||
"text": "Your document text here",
|
||||
"metadata": {"title": "Document Title", "category": "Category"}
|
||||
}
|
||||
|
||||
POST /search
|
||||
{
|
||||
"query": "your search query",
|
||||
"top_k": 10
|
||||
}
|
||||
|
||||
GET /stats
|
||||
GET /index/structure
|
||||
GET /document/{doc_id}
|
||||
DELETE /index
|
||||
```
|
||||
|
||||
#### Demo
|
||||
|
||||
```bash
|
||||
python demo.py
|
||||
```
|
||||
|
||||
Demo: clear index → sample CS docs → stats → index structure → sample queries → document get.
|
||||
|
||||
### Educational features
|
||||
|
||||
Logging covers tokenization, TF, IDF, per-term BM25, query processing, candidates. `/index/structure` returns inverted map, doc stats, BM25 params, global TF distribution. Search results include matched terms, doc length, TFs, per-term score contributions.
|
||||
|
||||
### Project structure
|
||||
|
||||
```
|
||||
sparse-embedding/
|
||||
├── bm25_engine.py # Core engine
|
||||
├── cli.py # Offline CLI: BM25/SPLADE + metrics
|
||||
├── server.py # FastAPI server
|
||||
├── demo.py # Demo script
|
||||
├── requirements.txt
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### Limitations
|
||||
|
||||
In-memory only; basic tokenization (no lemmatization); English stopwords; no phrase queries / synonyms / multi-thread.
|
||||
|
||||
---
|
||||
|
||||
## 中文
|
||||
|
||||
### 概述
|
||||
|
||||
基于倒排索引与 BM25 的教学型稀疏向量搜索引擎，用详细日志与可视化帮助理解信息检索基本概念。
|
||||
|
||||
### 功能特性
|
||||
|
||||
- **完整 BM25 实现**
|
||||
- **高级分词**：数字、编码、技术术语、大小写混合
|
||||
- **倒排索引**
|
||||
- **HTTP API**（FastAPI）
|
||||
- **交互式 Web UI**
|
||||
- **教学日志**
|
||||
- **索引结构可视化 API**
|
||||
- **内存存储**（教学简化）
|
||||
|
||||
#### 分词能力
|
||||
|
||||
- **数字**：`404`、`3.14`、`2.0.1`
|
||||
- **编码**：`XK9-2B4-7Q1`、`API_KEY_123`
|
||||
- **技术术语**：`C++`、`.NET`、`Node.js`
|
||||
- **大小写混合**：`JavaScript`、`PyTorch`、`iPhone`
|
||||
- **邮箱**：`user@example.com`
|
||||
- **十六进制**：`#FF5733`、`0x1234`
|
||||
- **缩写**：`API`、`HTTP`、`NASA`
|
||||
- **字母数字**：`Python3`、`ES6`、`HTML5`
|
||||
|
||||
### 架构
|
||||
|
||||
1. **TextProcessor**：分词
|
||||
2. **InvertedIndex**：词频 / 文档频率
|
||||
3. **BM25**：相关性打分
|
||||
4. **SparseSearchEngine**：总控
|
||||
5. **HTTP Server**：FastAPI
|
||||
|
||||
BM25 使用 TF、IDF 与文档长度归一化。关键参数：`k1`（默认 1.5）、`b`（默认 0.75）。
|
||||
|
||||
### 安装
|
||||
|
||||
```bash
|
||||
# 在仓库根目录使用统一的第 3 章环境
|
||||
uv sync --locked --python 3.12 --extra ch3
|
||||
|
||||
# 切换目录前先激活环境：
|
||||
# macOS/Linux：
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell：.venv\Scripts\Activate.ps1
|
||||
# Windows cmd：.venv\Scripts\activate.bat
|
||||
|
||||
# 未安装 uv 时可用 pip 兜底：
|
||||
# python -m pip install -e ".[ch3]"
|
||||
|
||||
cd chapter3/sparse-embedding
|
||||
|
||||
# 迁移期间仍支持单项目兼容路径：
|
||||
# python -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
`cli.py` 只依赖 Python 标准库，无需第三方包即可离线运行；`server.py` / `demo.py` 需要统一 `ch3` 环境或兼容 `requirements.txt` 路径提供的 FastAPI 等依赖。
|
||||
|
||||
### 命令行工具 cli.py（实验 3-5，推荐入口）
|
||||
|
||||
完全离线：在内置 10 篇小语料上跑 BM25、复现书中「逐词 IDF/TF/BM25 贡献」日志，并在标注集上算 recall/precision/MRR。参数均有中文 `--help`。
|
||||
|
||||
```bash
|
||||
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" # 观察同义词短板（kitten/feline 漏召回）
|
||||
python cli.py --corpus my.json -q "查询" -o out.json
|
||||
python cli.py --k1 2.0 -b 0.5 -q "..."
|
||||
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` | 在标注集上评测 |
|
||||
| `--labels` | 自定义标注 `{query: [doc_id,...]}` |
|
||||
| `--explain` | 逐词 TF / IDF / BM25 |
|
||||
| `--k1` / `-b` | BM25 参数 |
|
||||
| `-v, --verbose` | DEBUG 日志 |
|
||||
|
||||
#### 检索质量评测（`--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 在精确关键词、错误码、专有名称上极佳，但读不懂同义词——查询 `cat` 无法命中只写 `kitten` / `feline` 的文档。这正是引入混合检索（实验 3-6 `retrieval-pipeline`）的动机。
|
||||
|
||||
#### 学习型稀疏检索（`--method splade`）
|
||||
|
||||
用掩码语言模型为词项打权，并可为语义相关词项补权。需下载 `naver/splade-cocondenser-ensembledistil`（依赖 `torch`、`transformers`）。离线无权重时会快速给出清晰提示。联网可先 `huggingface-cli download naver/splade-cocondenser-ensembledistil`。
|
||||
|
||||
### 服务端用法
|
||||
|
||||
```bash
|
||||
python server.py
|
||||
```
|
||||
|
||||
服务地址：`http://localhost:4241`。Web UI 打开该地址。API 文档：`http://localhost:4241/docs`。
|
||||
|
||||
#### API 端点
|
||||
|
||||
```bash
|
||||
POST /index
|
||||
{
|
||||
"text": "Your document text here",
|
||||
"metadata": {"title": "Document Title", "category": "Category"}
|
||||
}
|
||||
|
||||
POST /search
|
||||
{
|
||||
"query": "your search query",
|
||||
"top_k": 10
|
||||
}
|
||||
|
||||
GET /stats
|
||||
GET /index/structure
|
||||
GET /document/{doc_id}
|
||||
DELETE /index
|
||||
```
|
||||
|
||||
#### 运行演示
|
||||
|
||||
```bash
|
||||
python demo.py
|
||||
```
|
||||
|
||||
演示：清空索引 → 示例文档 → 统计 → 索引结构 → 查询 → 按 ID 取文档。
|
||||
|
||||
### 教学特性
|
||||
|
||||
日志覆盖分词、TF、IDF、逐词 BM25、查询处理、候选文档。`/index/structure` 返回倒排映射、文档统计、BM25 参数、全局词频分布。检索结果含匹配词、文档长度、词频与分项得分。
|
||||
|
||||
### 项目结构
|
||||
|
||||
```
|
||||
sparse-embedding/
|
||||
├── bm25_engine.py # 核心检索引擎
|
||||
├── cli.py # 离线 CLI：BM25/SPLADE + 指标
|
||||
├── server.py # FastAPI 服务
|
||||
├── demo.py # 演示脚本
|
||||
├── requirements.txt
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### 局限
|
||||
|
||||
仅内存存储；分词较基础（无词形还原）；英文停用词；不支持短语查询 / 同义扩展 / 多线程。
|
||||
|
||||
---
|
||||
|
||||
## Notes / 说明
|
||||
|
||||
- Related next step / 相关后续：[`../retrieval-pipeline/`](../retrieval-pipeline/) hybrid dense+sparse pipeline (Exp. 3-6).
|
||||
- 相关后续：[`../retrieval-pipeline/`](../retrieval-pipeline/) 混合稠密+稀疏流水线（实验 3-6）。
|
||||
@@ -0,0 +1,191 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Acceptance campaign for Chapter 3 Experiments 3-5.
|
||||
|
||||
This intentionally exercises the repository's from-scratch inverted index and
|
||||
BM25 implementation. It checks one score against an independent, explicit
|
||||
calculation and then measures the exact-keyword/synonym contrast on a labelled
|
||||
corpus. No third-party retrieval implementation is used as an oracle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
PROJECT_DIR = Path(__file__).resolve().parent
|
||||
CHAPTER_DIR = PROJECT_DIR.parent
|
||||
sys.path.insert(0, str(CHAPTER_DIR))
|
||||
|
||||
from experiment_utils import write_campaign_evidence # noqa: E402
|
||||
from bm25_engine import BM25, InvertedIndex, TextProcessor # noqa: E402
|
||||
from cli import DEFAULT_CORPUS, DEFAULT_LABELS, build_engine # noqa: E402
|
||||
|
||||
|
||||
K1 = 1.5
|
||||
B = 0.75
|
||||
TOP_K = 5
|
||||
|
||||
|
||||
def hand_calculation() -> dict[str, Any]:
|
||||
"""Compare engine output with a hand-calculable RSJ BM25 example."""
|
||||
|
||||
texts = [
|
||||
"rare rare common",
|
||||
"common common",
|
||||
"common filler filler filler",
|
||||
"filler",
|
||||
]
|
||||
index = InvertedIndex()
|
||||
for doc_id, text in enumerate(texts):
|
||||
index.add_document(doc_id, text, {"source": "hand-check"})
|
||||
bm25 = BM25(index, k1=K1, b=B)
|
||||
|
||||
term = "rare"
|
||||
doc_id = 0
|
||||
n_docs = len(texts)
|
||||
df = 1
|
||||
tf = 2
|
||||
dl = 3
|
||||
avgdl = sum(len(TextProcessor().tokenize(text)) for text in texts) / n_docs
|
||||
idf = math.log((n_docs - df + 0.5) / (df + 0.5))
|
||||
numerator = tf * (K1 + 1)
|
||||
denominator = tf + K1 * (1 - B + B * (dl / avgdl))
|
||||
expected = idf * numerator / denominator
|
||||
raw_engine_idf = bm25.calculate_raw_idf(term)
|
||||
actual = bm25.calculate_term_score(term, doc_id)
|
||||
tolerance = 1e-12
|
||||
|
||||
return {
|
||||
"corpus": texts,
|
||||
"term": term,
|
||||
"doc_id": doc_id,
|
||||
"parameters": {"N": n_docs, "df": df, "tf": tf, "dl": dl, "avgdl": avgdl, "k1": K1, "b": B},
|
||||
"formula": "ln((N-df+0.5)/(df+0.5)) * tf*(k1+1) / (tf+k1*(1-b+b*dl/avgdl))",
|
||||
"intermediate": {
|
||||
"independent_raw_idf": idf,
|
||||
"engine_raw_idf": raw_engine_idf,
|
||||
"scoring_idf": bm25.calculate_idf(term),
|
||||
"numerator": numerator,
|
||||
"denominator": denominator,
|
||||
},
|
||||
"expected_score": expected,
|
||||
"engine_score": actual,
|
||||
"absolute_error": abs(expected - actual),
|
||||
"tolerance": tolerance,
|
||||
"posting_list": sorted(index.get_posting_list(term)),
|
||||
"recorded_document_frequency": index.document_frequency[term],
|
||||
"passed": (
|
||||
abs(idf - raw_engine_idf) <= tolerance
|
||||
and abs(expected - actual) <= tolerance
|
||||
and index.document_frequency[term] == df
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def labelled_benchmark() -> dict[str, Any]:
|
||||
started = time.perf_counter()
|
||||
engine = build_engine(DEFAULT_CORPUS, k1=K1, b=B)
|
||||
build_ms = (time.perf_counter() - started) * 1000
|
||||
rows: list[dict[str, Any]] = []
|
||||
|
||||
for query, relevant_list in DEFAULT_LABELS.items():
|
||||
query_start = time.perf_counter()
|
||||
results = engine.search(query, top_k=TOP_K)
|
||||
latency_ms = (time.perf_counter() - query_start) * 1000
|
||||
retrieved = [result["doc_id"] for result in results]
|
||||
relevant = set(relevant_list)
|
||||
hits = [doc_id for doc_id in retrieved if doc_id in relevant]
|
||||
recall = len(set(hits)) / len(relevant)
|
||||
reciprocal_rank = next(
|
||||
(1.0 / rank for rank, doc_id in enumerate(retrieved, 1) if doc_id in relevant),
|
||||
0.0,
|
||||
)
|
||||
category = "synonym-only" if query == "cat" else "exact-keyword"
|
||||
rows.append(
|
||||
{
|
||||
"query": query,
|
||||
"category": category,
|
||||
"relevant": sorted(relevant),
|
||||
"retrieved": retrieved,
|
||||
"hits": hits,
|
||||
"recall_at_5": recall,
|
||||
"reciprocal_rank": reciprocal_rank,
|
||||
"latency_ms": round(latency_ms, 3),
|
||||
"results": [
|
||||
{
|
||||
"rank": rank,
|
||||
"doc_id": result["doc_id"],
|
||||
"score": result["score"],
|
||||
"matched_terms": result["debug"]["matched_terms"],
|
||||
"term_frequencies": result["debug"]["term_frequencies"],
|
||||
}
|
||||
for rank, result in enumerate(results, 1)
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
exact = [row for row in rows if row["category"] == "exact-keyword"]
|
||||
synonym = [row for row in rows if row["category"] == "synonym-only"]
|
||||
return {
|
||||
"corpus": DEFAULT_CORPUS,
|
||||
"labels": DEFAULT_LABELS,
|
||||
"parameters": {"k1": K1, "b": B, "top_k": TOP_K},
|
||||
"index_statistics": engine.index.get_statistics(),
|
||||
"build_latency_ms": round(build_ms, 3),
|
||||
"queries": rows,
|
||||
"metrics": {
|
||||
"exact_keyword_recall_at_5": sum(row["recall_at_5"] for row in exact) / len(exact),
|
||||
"exact_keyword_mrr": sum(row["reciprocal_rank"] for row in exact) / len(exact),
|
||||
"synonym_only_recall_at_5": sum(row["recall_at_5"] for row in synonym) / len(synonym),
|
||||
"synonym_only_mrr": sum(row["reciprocal_rank"] for row in synonym) / len(synonym),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
hand = hand_calculation()
|
||||
benchmark = labelled_benchmark()
|
||||
metrics = benchmark["metrics"]
|
||||
acceptance = {
|
||||
"uses_from_scratch_engine": True,
|
||||
"hand_score_matches": hand["passed"],
|
||||
"inverted_index_df_matches": hand["recorded_document_frequency"] == hand["parameters"]["df"],
|
||||
"all_exact_keyword_queries_recalled": metrics["exact_keyword_recall_at_5"] == 1.0,
|
||||
"synonym_only_failure_observed": metrics["synonym_only_recall_at_5"] == 0.0,
|
||||
"transparent_tf_idf_scores_retained": all(
|
||||
"term_frequencies" in result
|
||||
for row in benchmark["queries"]
|
||||
for result in row["results"]
|
||||
),
|
||||
}
|
||||
passed = all(acceptance.values())
|
||||
evidence = {
|
||||
"status": "passed" if passed else "failed",
|
||||
"method": {
|
||||
"implementation": "chapter3/sparse-embedding/bm25_engine.py",
|
||||
"algorithm": "from-scratch inverted index + Robertson/Sparck Jones BM25",
|
||||
"third_party_retrieval_library": None,
|
||||
},
|
||||
"hand_calculation": hand,
|
||||
"benchmark": benchmark,
|
||||
"summary": metrics,
|
||||
"acceptance": acceptance,
|
||||
}
|
||||
manifest = write_campaign_evidence(
|
||||
PROJECT_DIR,
|
||||
"3-5",
|
||||
evidence,
|
||||
input_paths=[__file__, PROJECT_DIR / "bm25_engine.py", PROJECT_DIR / "cli.py"],
|
||||
)
|
||||
print(f"hand score error: {hand['absolute_error']:.3g}")
|
||||
print(f"exact recall@5: {metrics['exact_keyword_recall_at_5']:.3f}")
|
||||
print(f"synonym recall@5: {metrics['synonym_only_recall_at_5']:.3f}")
|
||||
print(f"evidence: {manifest['run_dir']}")
|
||||
return 0 if passed else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,529 @@
|
||||
"""
|
||||
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
|
||||
|
||||
# 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',
|
||||
# Words with apostrophes
|
||||
r"\b[a-zA-Z]+'[a-zA-Z]+\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
|
||||
r"\b[a-zA-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}")
|
||||
|
||||
is_update = doc_id in self.documents
|
||||
if is_update:
|
||||
old_terms = set(self.term_frequency.get(doc_id, Counter()).keys())
|
||||
for term in old_terms:
|
||||
if term in self.index and doc_id in self.index[term]:
|
||||
self.index[term].remove(doc_id)
|
||||
if not self.index[term]:
|
||||
del self.index[term]
|
||||
if term in self.document_frequency:
|
||||
self.document_frequency[term] -= 1
|
||||
if self.document_frequency[term] <= 0:
|
||||
del self.document_frequency[term]
|
||||
|
||||
# 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")
|
||||
|
||||
for term in term_freq:
|
||||
if term not in self.index or doc_id not in self.index[term]:
|
||||
self.document_frequency[term] = self.document_frequency.get(term, 0) + 1
|
||||
self.index[term].add(doc_id)
|
||||
|
||||
if not is_update:
|
||||
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_raw_idf(self, term: str) -> float:
|
||||
"""Calculate the Robertson/Sparck Jones IDF printed in Chapter 3."""
|
||||
N = self.index.total_documents
|
||||
df = len(self.index.get_posting_list(term))
|
||||
|
||||
if df == 0:
|
||||
return 0
|
||||
|
||||
N = max(N, df)
|
||||
val = (N - df + 0.5) / (df + 0.5)
|
||||
if val <= 0:
|
||||
return 0.0
|
||||
return math.log(val)
|
||||
|
||||
def calculate_idf(self, term: str) -> float:
|
||||
"""Return RSJ IDF with a small floor for corpus-ubiquitous terms.
|
||||
|
||||
Raw RSJ IDF is negative when a term occurs in more than half of a tiny
|
||||
corpus. Letting that value flow into ranking perversely rewards a
|
||||
document for matching fewer query terms. Production BM25 variants
|
||||
conventionally floor or smooth that edge case; the raw value remains
|
||||
available through :meth:`calculate_raw_idf` for transparent teaching
|
||||
and hand calculation.
|
||||
"""
|
||||
df = len(self.index.get_posting_list(term))
|
||||
if df == 0:
|
||||
return 0
|
||||
raw_idf = self.calculate_raw_idf(term)
|
||||
idf = max(raw_idf, 1e-6)
|
||||
|
||||
logger.debug(
|
||||
f"IDF for '{term}': N={self.index.total_documents}, df={df}, "
|
||||
f"raw_idf={raw_idf:.4f}, scoring_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
|
||||
if self.avgdl == 0:
|
||||
return 0.0
|
||||
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).
|
||||
# Resolve each term to the variant that actually matched the index:
|
||||
# when the lowercase fallback finds the docs, scoring must use the
|
||||
# lowercase term too, otherwise tf lookups return 0 and the fallback
|
||||
# candidates all score 0.0.
|
||||
candidate_docs = set()
|
||||
term_doc_mapping = {}
|
||||
resolved_terms = []
|
||||
|
||||
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 term and not term[0].isdigit() and '-' not in term:
|
||||
lowered = term.lower()
|
||||
docs = self.index.get_posting_list(lowered)
|
||||
if docs:
|
||||
term = lowered
|
||||
|
||||
resolved_terms.append(term)
|
||||
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(resolved_terms, doc_id)
|
||||
|
||||
# Collect debug information
|
||||
debug_info = {
|
||||
'matched_terms': [term for term in resolved_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 resolved_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")
|
||||
@@ -0,0 +1,341 @@
|
||||
#!/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())
|
||||
@@ -0,0 +1,10 @@
|
||||
"""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")
|
||||
@@ -0,0 +1,286 @@
|
||||
"""
|
||||
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", timeout=30)
|
||||
if response.status_code == 200:
|
||||
logger.info("Server is ready!")
|
||||
return True
|
||||
except Exception:
|
||||
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", timeout=30)
|
||||
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"]}, timeout=30
|
||||
)
|
||||
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", timeout=30)
|
||||
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}, timeout=30
|
||||
)
|
||||
|
||||
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", timeout=30)
|
||||
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}", timeout=30)
|
||||
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()
|
||||
@@ -0,0 +1,184 @@
|
||||
#!/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()
|
||||
@@ -0,0 +1,5 @@
|
||||
fastapi==0.104.1
|
||||
uvicorn[standard]==0.24.0
|
||||
pydantic==2.5.0
|
||||
requests==2.31.0
|
||||
python-multipart==0.0.6
|
||||
@@ -0,0 +1,503 @@
|
||||
"""
|
||||
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")
|
||||
@@ -0,0 +1,17 @@
|
||||
import pytest
|
||||
from bm25_engine import InvertedIndex, BM25, TextProcessor
|
||||
|
||||
|
||||
def test_bm25_search_empty_term_in_query_terms_no_index_error():
|
||||
index = InvertedIndex()
|
||||
index.add_document(1, "hello world")
|
||||
bm25 = BM25(index)
|
||||
|
||||
orig_tokenize = TextProcessor.tokenize
|
||||
try:
|
||||
TextProcessor.tokenize = lambda self, text, remove_stop_words=True: ["hello", ""]
|
||||
results = bm25.search("hello")
|
||||
assert len(results) == 1
|
||||
assert results[0][0] == 1
|
||||
finally:
|
||||
TextProcessor.tokenize = orig_tokenize
|
||||
@@ -0,0 +1,321 @@
|
||||
"""
|
||||
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()
|
||||
@@ -0,0 +1,10 @@
|
||||
import pytest
|
||||
from bm25_engine import TextProcessor
|
||||
|
||||
|
||||
def test_tokenize_preserves_apostrophe_contractions():
|
||||
processor = TextProcessor()
|
||||
tokens = processor.tokenize("don't user's it's text")
|
||||
assert "don't" in tokens
|
||||
assert "user's" in tokens
|
||||
assert "it's" in tokens
|
||||
@@ -0,0 +1,44 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-5",
|
||||
"run_id": "20260729T183232Z-3_5-f9f70b37",
|
||||
"created_at": "2026-07-29T18:32:32.574550+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/validation/runs/20260729T183232Z-3_5-f9f70b37",
|
||||
"artifacts": {
|
||||
"evidence.json": "69bc798be2f18e814b53665d29692cb4182e70561f4dcf222f25bf36d21eae0e",
|
||||
"receipts.json": "37517e5f3dc66819f61f5a7bb8ace1921282415f10551d2defa5c3eb0985b570",
|
||||
"manifest.json": "632bd554d004aa4139907c02006d34f6c0c192a3b8e940689b5d154eadbc67d0"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/benchmark.py",
|
||||
"sha256": "f3601a917de0b1a169691fbb5fb4323b7310220af9eb2221414ed74220a8da89",
|
||||
"bytes": 7038
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/bm25_engine.py",
|
||||
"sha256": "efcaa2eef462b05d2e1dc9d9f2c5dc0d548a529ca17315ec52e0b0e85fe636c0",
|
||||
"bytes": 20138
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/cli.py",
|
||||
"sha256": "6b76c01fd3889a029b5c37a4f589e3a438bc093c406e144e5e4209441c556ec2",
|
||||
"bytes": 16004
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"exact_keyword_recall_at_5": 1.0,
|
||||
"exact_keyword_mrr": 1.0,
|
||||
"synonym_only_recall_at_5": 0.0,
|
||||
"synonym_only_mrr": 0.0
|
||||
},
|
||||
"acceptance": {
|
||||
"uses_from_scratch_engine": true,
|
||||
"hand_score_matches": true,
|
||||
"inverted_index_df_matches": true,
|
||||
"all_exact_keyword_queries_recalled": true,
|
||||
"synonym_only_failure_observed": true,
|
||||
"transparent_tf_idf_scores_retained": true
|
||||
}
|
||||
}
|
||||
+379
@@ -0,0 +1,379 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-5",
|
||||
"run_id": "20260729T183232Z-3_5-f9f70b37",
|
||||
"provenance": {
|
||||
"captured_at": "2026-07-29T18:32:32.564511+00:00",
|
||||
"git_revision": "4a7f37cf278bd15948c409f14533017c4c7fbc29",
|
||||
"python": "3.11.4 (main, Jul 5 2023, 08:40:20) [Clang 14.0.6 ]",
|
||||
"platform": "macOS-26.3-arm64-arm-64bit",
|
||||
"credential_presence": {
|
||||
"ARK_API_KEY": true,
|
||||
"MOONSHOT_API_KEY": true,
|
||||
"OPENAI_API_KEY": true,
|
||||
"GEMINI_API_KEY": true,
|
||||
"SILICONFLOW_API_KEY": true
|
||||
}
|
||||
},
|
||||
"status": "passed",
|
||||
"method": {
|
||||
"implementation": "chapter3/sparse-embedding/bm25_engine.py",
|
||||
"algorithm": "from-scratch inverted index + Robertson/Sparck Jones BM25",
|
||||
"third_party_retrieval_library": null
|
||||
},
|
||||
"hand_calculation": {
|
||||
"corpus": [
|
||||
"rare rare common",
|
||||
"common common",
|
||||
"common filler filler filler",
|
||||
"filler"
|
||||
],
|
||||
"term": "rare",
|
||||
"doc_id": 0,
|
||||
"parameters": {
|
||||
"N": 4,
|
||||
"df": 1,
|
||||
"tf": 2,
|
||||
"dl": 3,
|
||||
"avgdl": 2.5,
|
||||
"k1": 1.5,
|
||||
"b": 0.75
|
||||
},
|
||||
"formula": "ln((N-df+0.5)/(df+0.5)) * tf*(k1+1) / (tf+k1*(1-b+b*dl/avgdl))",
|
||||
"intermediate": {
|
||||
"independent_raw_idf": 0.8472978603872037,
|
||||
"engine_raw_idf": 0.8472978603872037,
|
||||
"scoring_idf": 0.8472978603872037,
|
||||
"numerator": 5.0,
|
||||
"denominator": 3.7249999999999996
|
||||
},
|
||||
"expected_score": 1.1373125642781259,
|
||||
"engine_score": 1.1373125642781259,
|
||||
"absolute_error": 0.0,
|
||||
"tolerance": 1e-12,
|
||||
"posting_list": [
|
||||
0
|
||||
],
|
||||
"recorded_document_frequency": 1,
|
||||
"passed": true
|
||||
},
|
||||
"benchmark": {
|
||||
"corpus": [
|
||||
{
|
||||
"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."
|
||||
}
|
||||
],
|
||||
"labels": {
|
||||
"model distillation": [
|
||||
"doc_3",
|
||||
"doc_4"
|
||||
],
|
||||
"HTTP 404 error": [
|
||||
"doc_6"
|
||||
],
|
||||
"XK9-2B4-7Q1": [
|
||||
"doc_9"
|
||||
],
|
||||
"BM25 ranking function": [
|
||||
"doc_5"
|
||||
],
|
||||
"cat": [
|
||||
"doc_7",
|
||||
"doc_8"
|
||||
]
|
||||
},
|
||||
"parameters": {
|
||||
"k1": 1.5,
|
||||
"b": 0.75,
|
||||
"top_k": 5
|
||||
},
|
||||
"index_statistics": {
|
||||
"total_documents": 10,
|
||||
"unique_terms": 98,
|
||||
"total_terms": 114,
|
||||
"average_document_length": 11.4,
|
||||
"terms_by_frequency": [
|
||||
[
|
||||
"model",
|
||||
4
|
||||
],
|
||||
[
|
||||
"for",
|
||||
3
|
||||
],
|
||||
[
|
||||
"and",
|
||||
3
|
||||
],
|
||||
[
|
||||
"in",
|
||||
3
|
||||
],
|
||||
[
|
||||
"web",
|
||||
2
|
||||
],
|
||||
[
|
||||
"distillation",
|
||||
2
|
||||
],
|
||||
[
|
||||
"knowledge",
|
||||
2
|
||||
],
|
||||
[
|
||||
"to",
|
||||
2
|
||||
],
|
||||
[
|
||||
"frequency",
|
||||
2
|
||||
],
|
||||
[
|
||||
"error",
|
||||
2
|
||||
]
|
||||
]
|
||||
},
|
||||
"build_latency_ms": 1.668,
|
||||
"queries": [
|
||||
{
|
||||
"query": "model distillation",
|
||||
"category": "exact-keyword",
|
||||
"relevant": [
|
||||
"doc_3",
|
||||
"doc_4"
|
||||
],
|
||||
"retrieved": [
|
||||
"doc_3",
|
||||
"doc_4"
|
||||
],
|
||||
"hits": [
|
||||
"doc_3",
|
||||
"doc_4"
|
||||
],
|
||||
"recall_at_5": 1.0,
|
||||
"reciprocal_rank": 1.0,
|
||||
"latency_ms": 0.276,
|
||||
"results": [
|
||||
{
|
||||
"rank": 1,
|
||||
"doc_id": "doc_3",
|
||||
"score": 3.121561765506991,
|
||||
"matched_terms": [
|
||||
"model",
|
||||
"distillation"
|
||||
],
|
||||
"term_frequencies": {
|
||||
"model": 3,
|
||||
"distillation": 1
|
||||
}
|
||||
},
|
||||
{
|
||||
"rank": 2,
|
||||
"doc_id": "doc_4",
|
||||
"score": 2.219735866426749,
|
||||
"matched_terms": [
|
||||
"model",
|
||||
"distillation"
|
||||
],
|
||||
"term_frequencies": {
|
||||
"model": 1,
|
||||
"distillation": 1
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"query": "HTTP 404 error",
|
||||
"category": "exact-keyword",
|
||||
"relevant": [
|
||||
"doc_6"
|
||||
],
|
||||
"retrieved": [
|
||||
"doc_6",
|
||||
"doc_9"
|
||||
],
|
||||
"hits": [
|
||||
"doc_6"
|
||||
],
|
||||
"recall_at_5": 1.0,
|
||||
"reciprocal_rank": 1.0,
|
||||
"latency_ms": 0.269,
|
||||
"results": [
|
||||
{
|
||||
"rank": 1,
|
||||
"doc_id": "doc_6",
|
||||
"score": 4.994285959345282,
|
||||
"matched_terms": [
|
||||
"HTTP",
|
||||
"404",
|
||||
"error"
|
||||
],
|
||||
"term_frequencies": {
|
||||
"HTTP": 1,
|
||||
"404": 1,
|
||||
"error": 1
|
||||
}
|
||||
},
|
||||
{
|
||||
"rank": 2,
|
||||
"doc_id": "doc_9",
|
||||
"score": 1.2953611811041896,
|
||||
"matched_terms": [
|
||||
"error"
|
||||
],
|
||||
"term_frequencies": {
|
||||
"HTTP": 0,
|
||||
"404": 0,
|
||||
"error": 1
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"query": "XK9-2B4-7Q1",
|
||||
"category": "exact-keyword",
|
||||
"relevant": [
|
||||
"doc_9"
|
||||
],
|
||||
"retrieved": [
|
||||
"doc_9"
|
||||
],
|
||||
"hits": [
|
||||
"doc_9"
|
||||
],
|
||||
"recall_at_5": 1.0,
|
||||
"reciprocal_rank": 1.0,
|
||||
"latency_ms": 0.129,
|
||||
"results": [
|
||||
{
|
||||
"rank": 1,
|
||||
"doc_id": "doc_9",
|
||||
"score": 1.9537998395246958,
|
||||
"matched_terms": [
|
||||
"XK9-2B4-7Q1"
|
||||
],
|
||||
"term_frequencies": {
|
||||
"XK9-2B4-7Q1": 1
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"query": "BM25 ranking function",
|
||||
"category": "exact-keyword",
|
||||
"relevant": [
|
||||
"doc_5"
|
||||
],
|
||||
"retrieved": [
|
||||
"doc_5"
|
||||
],
|
||||
"hits": [
|
||||
"doc_5"
|
||||
],
|
||||
"recall_at_5": 1.0,
|
||||
"reciprocal_rank": 1.0,
|
||||
"latency_ms": 0.261,
|
||||
"results": [
|
||||
{
|
||||
"rank": 1,
|
||||
"doc_id": "doc_5",
|
||||
"score": 5.626316650182078,
|
||||
"matched_terms": [
|
||||
"BM25",
|
||||
"ranking",
|
||||
"function"
|
||||
],
|
||||
"term_frequencies": {
|
||||
"BM25": 1,
|
||||
"ranking": 1,
|
||||
"function": 1
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"query": "cat",
|
||||
"category": "synonym-only",
|
||||
"relevant": [
|
||||
"doc_7",
|
||||
"doc_8"
|
||||
],
|
||||
"retrieved": [],
|
||||
"hits": [],
|
||||
"recall_at_5": 0.0,
|
||||
"reciprocal_rank": 0.0,
|
||||
"latency_ms": 0.095,
|
||||
"results": []
|
||||
}
|
||||
],
|
||||
"metrics": {
|
||||
"exact_keyword_recall_at_5": 1.0,
|
||||
"exact_keyword_mrr": 1.0,
|
||||
"synonym_only_recall_at_5": 0.0,
|
||||
"synonym_only_mrr": 0.0
|
||||
}
|
||||
},
|
||||
"summary": {
|
||||
"exact_keyword_recall_at_5": 1.0,
|
||||
"exact_keyword_mrr": 1.0,
|
||||
"synonym_only_recall_at_5": 0.0,
|
||||
"synonym_only_mrr": 0.0
|
||||
},
|
||||
"acceptance": {
|
||||
"uses_from_scratch_engine": true,
|
||||
"hand_score_matches": true,
|
||||
"inverted_index_df_matches": true,
|
||||
"all_exact_keyword_queries_recalled": true,
|
||||
"synonym_only_failure_observed": true,
|
||||
"transparent_tf_idf_scores_retained": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-5",
|
||||
"run_id": "20260729T183232Z-3_5-f9f70b37",
|
||||
"created_at": "2026-07-29T18:32:32.574550+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/validation/runs/20260729T183232Z-3_5-f9f70b37",
|
||||
"artifacts": {
|
||||
"evidence.json": "69bc798be2f18e814b53665d29692cb4182e70561f4dcf222f25bf36d21eae0e",
|
||||
"receipts.json": "37517e5f3dc66819f61f5a7bb8ace1921282415f10551d2defa5c3eb0985b570"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/benchmark.py",
|
||||
"sha256": "f3601a917de0b1a169691fbb5fb4323b7310220af9eb2221414ed74220a8da89",
|
||||
"bytes": 7038
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/bm25_engine.py",
|
||||
"sha256": "efcaa2eef462b05d2e1dc9d9f2c5dc0d548a529ca17315ec52e0b0e85fe636c0",
|
||||
"bytes": 20138
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/sparse-embedding/cli.py",
|
||||
"sha256": "6b76c01fd3889a029b5c37a4f589e3a438bc093c406e144e5e4209441c556ec2",
|
||||
"bytes": 16004
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"exact_keyword_recall_at_5": 1.0,
|
||||
"exact_keyword_mrr": 1.0,
|
||||
"synonym_only_recall_at_5": 0.0,
|
||||
"synonym_only_mrr": 0.0
|
||||
},
|
||||
"acceptance": {
|
||||
"uses_from_scratch_engine": true,
|
||||
"hand_score_matches": true,
|
||||
"inverted_index_df_matches": true,
|
||||
"all_exact_keyword_queries_recalled": true,
|
||||
"synonym_only_failure_observed": true,
|
||||
"transparent_tf_idf_scores_retained": true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
[]
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
Reference in New Issue
Block a user