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,370 @@
|
||||
# Hybrid Retrieval Pipeline with Neural Reranking / 混合检索流水线与神经重排序
|
||||
|
||||
> Companion material for *AI Agents in Depth*, Chapter 3 — **Experiment 3-6**: dense + sparse + fusion + rerank, with offline `evaluate.py`.
|
||||
> 配套《深入理解 AI Agent》第 3 章 **实验 3-6**：稠密 + 稀疏 + 融合 + 重排，含离线 `evaluate.py`。
|
||||
|
||||
← [Chapter 3 index / 返回第 3 章目录](../README.md)
|
||||
|
||||
## Code map
|
||||
|
||||
- **Run first:** `python evaluate.py --no-dense --no-rerank` (offline BM25 smoke; the full pipeline needs the two retrieval services and local models).
|
||||
- **Start here:** `retrieval_pipeline.py::RetrievalPipeline.search` orchestrates retrieval, fusion and reranking.
|
||||
- **Core behavior:** `retrieval_client.py::RetrievalClient.search`, `fusion.py::fuse` and `reranker.py::Reranker.rerank`.
|
||||
- **State / protocol:** `document_store.py::DocumentStore`, `SearchResult`, `PipelineConfig` and `SearchMode`.
|
||||
- **Verifier:** `evaluate.py` reports recall/MRR by stage; `test_pipeline.py` and `test_weighted_fusion_dedup.py` lock down ranking and deduplication.
|
||||
- **Experiment variable:** dense/sparse/hybrid mode, fusion method, candidate `top_k` and `rerank_top_k`.
|
||||
- **Skip on first pass:** service startup scripts, model downloads and HTTP error adapters.
|
||||
|
||||
---
|
||||
|
||||
## English
|
||||
|
||||
### Educational goals
|
||||
|
||||
1. **Dense vs sparse**: when each wins and why
|
||||
2. **Hybrid search**: combining methods
|
||||
3. **Neural reranking**: reorder candidates with transformers
|
||||
4. **Parallel processing**: multi-service index/search
|
||||
5. **Production-ish patterns**: API design and error handling
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
┌──────────────────────────────────────────────┐
|
||||
│ Client Application │
|
||||
└────────────────────┬─────────────────────────┘
|
||||
▼
|
||||
┌──────────────────────────────────────────────┐
|
||||
│ Retrieval Pipeline (Port 4242) │
|
||||
│ Document Store (In-Memory) │
|
||||
│ BGE-Reranker-v2 (Local Model) │
|
||||
└────────┬──────────────────┬─────────────────┘
|
||||
▼ ▼
|
||||
┌─────────────────┐ ┌─────────────────┐
|
||||
│ Dense Service │ │ Sparse Service │
|
||||
│ (Port 4240) │ │ (Port 4241) │
|
||||
│ BGE-M3 Model │ │ BM25 Engine │
|
||||
└─────────────────┘ └─────────────────┘
|
||||
```
|
||||
|
||||
### Key concepts
|
||||
|
||||
**Dense (BGE-M3)**: semantic / cross-lingual / synonyms; may miss exact codes; costlier.
|
||||
**Sparse (BM25)**: exact terms / IDs; no semantics; fast.
|
||||
**Fusion (`fusion.py`)**: RRF `score(d)=Σ 1/(k+rank)` with `k=60` (rank-only, scale-free) or weighted sum after min-max normalize to `[0,1]`.
|
||||
**Rerank**: BGE-Reranker-v2-M3 (service); `BAAI/bge-reranker-base` in `evaluate.py`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Python 3.12 with the root `ch3` extra, macOS M1/M2 (or adjust device), ≥8GB RAM, ~5GB disk for models.
|
||||
|
||||
### 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/retrieval-pipeline
|
||||
|
||||
# Single-project compatibility path, still supported during migration:
|
||||
# python -m pip install -r requirements.txt
|
||||
# First run downloads: BGE-M3 ~2.3GB, BGE-Reranker-v2-M3 ~1.1GB
|
||||
```
|
||||
|
||||
### Running services
|
||||
|
||||
```bash
|
||||
./start_all_services.sh
|
||||
# Dense 4240, Sparse 4241, Pipeline 4242
|
||||
```
|
||||
|
||||
Or individually:
|
||||
|
||||
```bash
|
||||
# Terminal 1
|
||||
cd ../dense-embedding && python main.py --port 4240
|
||||
# Terminal 2
|
||||
cd ../sparse-embedding && python server.py --port 4241
|
||||
# Terminal 3
|
||||
cd ../retrieval-pipeline && python main.py --port 4242
|
||||
```
|
||||
|
||||
### Testing with services
|
||||
|
||||
```bash
|
||||
python test_client.py # educational cases
|
||||
python demo.py # interactive demo
|
||||
# API docs: http://localhost:4242/docs
|
||||
```
|
||||
|
||||
### Offline evaluation CLI (`evaluate.py`)
|
||||
|
||||
`test_client.py` / `demo.py` need ports 4240–4242. **`evaluate.py` runs the full pipeline in one process — no service startup needed, and fully offline once the models are cached**. Note: the first run still downloads the dense/rerank models from HuggingFace, so initial execution requires network access.
|
||||
|
||||
```bash
|
||||
python evaluate.py --help # Chinese help
|
||||
python evaluate.py # full stage table (default)
|
||||
python evaluate.py --no-dense # BM25 only, no models
|
||||
python evaluate.py --no-rerank
|
||||
python evaluate.py --query "XR-7003"
|
||||
python evaluate.py --embed-model BAAI/bge-m3 --pooling cls
|
||||
python evaluate.py --output result.json
|
||||
```
|
||||
|
||||
| Stage | Default component | Offline? |
|
||||
|-------|-------------------|----------|
|
||||
| chunk | character-window splitter | ✅ pure Python |
|
||||
| sparse | BM25 (`rank_bm25`) | ✅ no model download |
|
||||
| dense | `sentence-transformers/all-MiniLM-L6-v2` (~90MB) | ✅ cached HF |
|
||||
| fuse | RRF + weighted (`fusion.py`) | ✅ pure Python |
|
||||
| rerank | `BAAI/bge-reranker-base` (~1.1GB first download) | ✅ once cached |
|
||||
|
||||
> `--no-dense` needs no ML model. Dense/rerank models download from HuggingFace on first run (network required); after that they run from local cache, and `--offline` forces loading from the local cache only. On Apple Silicon, MPS `NaN` is detected and falls back to CPU.
|
||||
|
||||
### Real output (reproduced)
|
||||
|
||||
Hard clusters: near-duplicate codes (`XR-7001..`, `HTTP-400..`) break dense; zero-lexical paraphrases break BM25.
|
||||
|
||||
```
|
||||
Stage / Method Recall@3 MRR nDCG@3
|
||||
------------------------------------------------------------------------------
|
||||
BM25 (sparse) 0.9000 0.8500 0.8631
|
||||
Dense 1.0000 0.9000 0.9262
|
||||
Hybrid-RRF 1.0000 1.0000 1.0000
|
||||
Hybrid-Weighted 1.0000 0.9500 0.9631
|
||||
Hybrid-RRF+Rerank 1.0000 0.9500 0.9631
|
||||
```
|
||||
|
||||
**How to read it:** BM25 nails codes, fails paraphrases; Dense is the mirror; **Hybrid-RRF** reaches perfect 1.00 (headline of Exp. 3-6). Weighted can be less robust (scale alignment). On this toy 17-doc set RRF is already strong; rerank value grows on larger pools / NL queries.
|
||||
|
||||
```
|
||||
$ python evaluate.py --query "XR-7003"
|
||||
[BM25 (sparse)]
|
||||
1. xr_7003 score= 3.2260 Product model XR-7003 is a smartphone available now.
|
||||
[Dense]
|
||||
1. xr_7001 score= 0.5247 Product model XR-7001 ...
|
||||
2. xr_7003 score= 0.5195 Product model XR-7003 ...
|
||||
[Hybrid-RRF]
|
||||
1. xr_7003 score= 0.0325 Product model XR-7003 ...
|
||||
```
|
||||
|
||||
### Educational test cases (with services)
|
||||
|
||||
1. Semantic (“kitty behavior” / feline) — dense wins
|
||||
2. Exact name (“Alexander Humphrey”) — sparse wins
|
||||
3. Multilingual (“人工智能”) — dense wins
|
||||
4. Codes (“HTTP-403”) — sparse wins
|
||||
5. Concepts (“happiness and excitement”) — dense wins
|
||||
|
||||
### API
|
||||
|
||||
```bash
|
||||
POST /index
|
||||
{"text": "Document content", "doc_id": "optional_id", "metadata": {"category": "example"}}
|
||||
|
||||
POST /search
|
||||
{"query": "search terms", "mode": "hybrid", "top_k": 20, "rerank_top_k": 10}
|
||||
|
||||
GET /stats
|
||||
GET /documents?limit=10&offset=0
|
||||
```
|
||||
|
||||
Response includes dense/sparse rankings, reranked results, rank changes, overlap stats.
|
||||
|
||||
### Project structure
|
||||
|
||||
```
|
||||
retrieval-pipeline/
|
||||
├── config.py, document_store.py, retrieval_client.py
|
||||
├── reranker.py, fusion.py, retrieval_pipeline.py
|
||||
├── evaluate.py, main.py, test_client.py, demo.py
|
||||
├── requirements.txt, start_all_services.sh, stop_all_services.sh
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### Performance / takeaways
|
||||
|
||||
- Latency ballpark: dense 50–100ms, sparse 10–30ms, rerank 100–200ms (20 docs)
|
||||
- Memory ~4GB models + docs
|
||||
- No single method wins; hybrid usually better; rerank improves relevance
|
||||
|
||||
### Troubleshooting
|
||||
|
||||
Ports 4240–4242 free; models downloaded; Python 3.12 for the root `ch3` install. OOM → smaller batches, CPU, FP16. First run slow (downloads).
|
||||
|
||||
### Further reading
|
||||
|
||||
[BGE-M3](https://arxiv.org/abs/2402.03216) · [BM25](https://en.wikipedia.org/wiki/Okapi_BM25) · [Neural IR](https://arxiv.org/abs/2301.09191)
|
||||
|
||||
### License
|
||||
|
||||
Educational project for learning purposes.
|
||||
|
||||
---
|
||||
|
||||
## 中文
|
||||
|
||||
### 教学目标
|
||||
|
||||
1. **稠密 vs 稀疏**：各自擅长场景
|
||||
2. **混合检索**：多路互补
|
||||
3. **神经重排序**：用 Transformer 重排候选
|
||||
4. **并行处理**：多服务索引/检索
|
||||
5. **工程模式**：API 与错误处理
|
||||
|
||||
### 架构
|
||||
|
||||
（与 English 节相同：Pipeline 4242，Dense 4240，Sparse 4241。）
|
||||
|
||||
### 关键概念
|
||||
|
||||
**稠密（BGE-M3）**：语义/跨语言/同义词；可能漏精确编码；计算更贵。
|
||||
**稀疏（BM25）**：精确词/ID；无语义；快。
|
||||
**融合（`fusion.py`）**：RRF（`k=60`）或 min-max 后加权求和。
|
||||
**重排**：服务用 BGE-Reranker-v2-M3；`evaluate.py` 用 `BAAI/bge-reranker-base`。
|
||||
|
||||
### 前置与安装
|
||||
|
||||
Python 3.12 与根目录 `ch3` extra，建议 ≥8GB 内存，约 5GB 模型空间。
|
||||
|
||||
```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/retrieval-pipeline
|
||||
|
||||
# 迁移期间仍支持单项目兼容路径：
|
||||
# python -m pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### 启动服务
|
||||
|
||||
```bash
|
||||
./start_all_services.sh
|
||||
```
|
||||
|
||||
或分别：
|
||||
|
||||
```bash
|
||||
cd ../dense-embedding && python main.py --port 4240
|
||||
cd ../sparse-embedding && python server.py --port 4241
|
||||
cd ../retrieval-pipeline && python main.py --port 4242
|
||||
```
|
||||
|
||||
### 带服务测试
|
||||
|
||||
```bash
|
||||
python test_client.py
|
||||
python demo.py
|
||||
# http://localhost:4242/docs
|
||||
```
|
||||
|
||||
### 离线评测 CLI（`evaluate.py`）
|
||||
|
||||
**单进程、可离线**跑通 chunk → embed → retrieve → fuse → rerank。
|
||||
|
||||
```bash
|
||||
python evaluate.py --help
|
||||
python evaluate.py
|
||||
python evaluate.py --no-dense
|
||||
python evaluate.py --no-rerank
|
||||
python evaluate.py --query "XR-7003"
|
||||
python evaluate.py --embed-model BAAI/bge-m3 --pooling cls
|
||||
python evaluate.py --output result.json
|
||||
```
|
||||
|
||||
| 阶段 | 默认组件 | 离线？ |
|
||||
|------|----------|--------|
|
||||
| chunk | 字符窗口切分 | ✅ 纯 Python |
|
||||
| sparse | BM25 | ✅ 无需下载模型 |
|
||||
| dense | MiniLM-L6-v2（~90MB） | ✅ HF 缓存 |
|
||||
| fuse | RRF + weighted | ✅ 纯 Python |
|
||||
| rerank | bge-reranker-base | ✅ 首次下载后缓存 |
|
||||
|
||||
> `--no-dense` 完全不需 ML 模型。Apple Silicon 上 MPS 出现 `NaN` 时自动回退 CPU。
|
||||
|
||||
### 真实输出解读
|
||||
|
||||
近重复编码打崩稠密；零词面重叠改写打崩 BM25。**Hybrid-RRF 全面 1.00** 是实验 3-6 的核心结论。加权融合对尺度更敏感。小语料上 RRF 已很强，重排价值在更大候选池与自然语言查询中更明显。
|
||||
|
||||
单查询追踪：
|
||||
|
||||
```
|
||||
$ python evaluate.py --query "XR-7003"
|
||||
[BM25 (sparse)]
|
||||
1. xr_7003 ...
|
||||
[Dense]
|
||||
1. xr_7001 ... # 稠密先排到兄弟编码
|
||||
2. xr_7003 ...
|
||||
[Hybrid-RRF]
|
||||
1. xr_7003 ... # 融合把精确匹配推回第 1
|
||||
```
|
||||
|
||||
### 教学测试用例（需服务）
|
||||
|
||||
语义 / 精确人名 / 多语言 / 技术编码 / 概念词——分别观察稠密或稀疏胜出。
|
||||
|
||||
### API
|
||||
|
||||
```bash
|
||||
POST /index
|
||||
{"text": "Document content", "doc_id": "optional_id", "metadata": {"category": "example"}}
|
||||
|
||||
POST /search
|
||||
{"query": "search terms", "mode": "hybrid", "top_k": 20, "rerank_top_k": 10}
|
||||
|
||||
GET /stats
|
||||
GET /documents?limit=10&offset=0
|
||||
```
|
||||
|
||||
响应含稠密/稀疏原始排名、重排结果、排名变化与重叠统计。
|
||||
|
||||
### 项目结构
|
||||
|
||||
```
|
||||
retrieval-pipeline/
|
||||
├── config.py, document_store.py, retrieval_client.py
|
||||
├── reranker.py, fusion.py, retrieval_pipeline.py
|
||||
├── evaluate.py, main.py, test_client.py, demo.py
|
||||
├── requirements.txt, start_all_services.sh, stop_all_services.sh
|
||||
└── README.md
|
||||
```
|
||||
|
||||
### 性能与要点
|
||||
|
||||
- 时延量级：稠密 50–100ms，稀疏 10–30ms，重排约 100–200ms（20 文档）
|
||||
- 模型内存约 4GB
|
||||
- 没有单一最优；混合通常更好；重排提升相关性
|
||||
|
||||
### 故障排查
|
||||
|
||||
检查 4240–4242 端口与模型下载；OOM 时减小 batch、改 CPU、开 FP16。
|
||||
|
||||
### 延伸阅读与许可
|
||||
|
||||
[BGE-M3](https://arxiv.org/abs/2402.03216) · [BM25](https://en.wikipedia.org/wiki/Okapi_BM25) · 教学项目。
|
||||
|
||||
---
|
||||
|
||||
## Notes / 说明
|
||||
|
||||
- Upstream services: [`../dense-embedding/`](../dense-embedding/) (4240), [`../sparse-embedding/`](../sparse-embedding/) (4241).
|
||||
- 上游服务：[`../dense-embedding/`](../dense-embedding/)（4240）、[`../sparse-embedding/`](../sparse-embedding/)（4241）。
|
||||
@@ -0,0 +1,136 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Canonical real-model campaign for Chapter 3 Experiment 3-6."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
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 evaluate import ( # noqa: E402
|
||||
DEFAULT_CORPUS,
|
||||
DEFAULT_QUERIES,
|
||||
METHOD_LABELS,
|
||||
Pipeline,
|
||||
build_parser,
|
||||
run_evaluation,
|
||||
)
|
||||
|
||||
|
||||
def cached_revision(model_name: str) -> str | None:
|
||||
ref = Path.home() / ".cache" / "huggingface" / "hub" / (
|
||||
"models--" + model_name.replace("/", "--")
|
||||
) / "refs" / "main"
|
||||
return ref.read_text(encoding="utf-8").strip() if ref.exists() else None
|
||||
|
||||
|
||||
def model_identity(pipeline: Pipeline, embed_name: str, reranker_name: str) -> dict:
|
||||
encoder = pipeline.dense.encoder
|
||||
reranker = pipeline.reranker
|
||||
return {
|
||||
"dense": {
|
||||
"provider": "local Hugging Face transformers",
|
||||
"model": embed_name,
|
||||
"cached_revision": cached_revision(embed_name),
|
||||
"class": type(encoder.model).__name__,
|
||||
"pooling": encoder.pooling,
|
||||
"parameters": sum(p.numel() for p in encoder.model.parameters()),
|
||||
"device": encoder.device,
|
||||
},
|
||||
"reranker": {
|
||||
"provider": "local Hugging Face transformers",
|
||||
"model": reranker_name,
|
||||
"cached_revision": cached_revision(reranker_name),
|
||||
"class": type(reranker.model).__name__,
|
||||
"parameters": sum(p.numel() for p in reranker.model.parameters()),
|
||||
"device": reranker.device,
|
||||
},
|
||||
"sparse": {"implementation": "rank_bm25.BM25Okapi"},
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = build_parser().parse_args([])
|
||||
args.embed_model = "Qwen/Qwen3-Embedding-0.6B"
|
||||
args.reranker_model = "cross-encoder/ms-marco-MiniLM-L-6-v2"
|
||||
args.pooling = "auto"
|
||||
args.device = "cpu"
|
||||
args.top_k = 10
|
||||
args.eval_k = 3
|
||||
args.rerank_pool = 10
|
||||
args.rerank_top_k = 10
|
||||
args.use_dense = True
|
||||
args.use_rerank = True
|
||||
|
||||
pipeline = Pipeline(DEFAULT_CORPUS, args)
|
||||
report = run_evaluation(pipeline, DEFAULT_QUERIES, args)
|
||||
identities = model_identity(pipeline, args.embed_model, args.reranker_model)
|
||||
methods = {key for key, _ in METHOD_LABELS}
|
||||
observed_methods = set(report["summary"])
|
||||
categories = {query.get("category") for query in DEFAULT_QUERIES}
|
||||
expected_categories = {"semantic", "exact-name", "multilingual", "technical-code"}
|
||||
all_rank_changes = [
|
||||
{"query": row["query"], **change}
|
||||
for row in report["per_query"]
|
||||
for change in row["trace"]["rank_changes"]
|
||||
]
|
||||
acceptance = {
|
||||
"real_dense_model_loaded": identities["dense"]["parameters"] > 100_000_000,
|
||||
"real_cross_encoder_loaded": identities["reranker"]["parameters"] > 1_000_000,
|
||||
"identical_labelled_queries_for_all_methods": all(
|
||||
set(row["methods"]) == methods for row in report["per_query"]
|
||||
),
|
||||
"all_required_query_categories_present": expected_categories <= categories,
|
||||
"sparse_dense_rrf_weighted_reranked_measured": observed_methods == methods,
|
||||
"recall_mrr_ndcg_and_latency_measured": all(
|
||||
{"recall@k", "mrr", "ndcg@k", "latency_ms"} <= set(metrics)
|
||||
for metrics in report["summary"].values()
|
||||
),
|
||||
"rank_changes_retained": bool(all_rank_changes),
|
||||
"hybrid_recall_not_below_best_single": report["summary"]["rrf"]["recall@k"]
|
||||
>= max(report["summary"]["sparse"]["recall@k"], report["summary"]["dense"]["recall@k"]),
|
||||
}
|
||||
evidence = {
|
||||
"status": "passed" if all(acceptance.values()) else "failed",
|
||||
"models": identities,
|
||||
"configuration": {
|
||||
"documents": len(DEFAULT_CORPUS),
|
||||
"chunks": pipeline.n_chunks,
|
||||
"queries": len(DEFAULT_QUERIES),
|
||||
"top_k": args.top_k,
|
||||
"eval_k": args.eval_k,
|
||||
"rrf_k": args.k_rrf,
|
||||
"rerank_pool": args.rerank_pool,
|
||||
"device": args.device,
|
||||
},
|
||||
"dataset": {"corpus": DEFAULT_CORPUS, "queries": DEFAULT_QUERIES},
|
||||
"report": report,
|
||||
"rank_changes": all_rank_changes,
|
||||
"summary": report["summary"],
|
||||
"acceptance": acceptance,
|
||||
}
|
||||
manifest = write_campaign_evidence(
|
||||
PROJECT_DIR,
|
||||
"3-6",
|
||||
evidence,
|
||||
receipts=[
|
||||
{
|
||||
"kind": "local-model-execution",
|
||||
"models": identities,
|
||||
"note": "No remote API or credential was used; complete rankings and timings are in evidence.json.",
|
||||
}
|
||||
],
|
||||
input_paths=[__file__, PROJECT_DIR / "evaluate.py", PROJECT_DIR / "fusion.py"],
|
||||
)
|
||||
print(json.dumps(report["summary"], indent=2))
|
||||
print(f"evidence: {manifest['run_dir']}")
|
||||
return 0 if all(acceptance.values()) else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Configuration for the retrieval pipeline."""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
class SearchMode(str, Enum):
|
||||
"""Search mode for retrieval."""
|
||||
DENSE = "dense"
|
||||
SPARSE = "sparse"
|
||||
HYBRID = "hybrid" # Both dense and sparse
|
||||
|
||||
@dataclass
|
||||
class ServiceConfig:
|
||||
"""Configuration for external services."""
|
||||
dense_service_url: str = "http://localhost:4240" # Port 4240 for dense service
|
||||
sparse_service_url: str = "http://localhost:4241" # Port 4241 for sparse service
|
||||
|
||||
@classmethod
|
||||
def from_env(cls):
|
||||
"""Create config from environment variables."""
|
||||
dense_url = os.getenv("DENSE_SERVICE_URL", "http://localhost:4240")
|
||||
sparse_url = os.getenv("SPARSE_SERVICE_URL", "http://localhost:4241")
|
||||
return cls(dense_service_url=dense_url, sparse_service_url=sparse_url)
|
||||
|
||||
@dataclass
|
||||
class RerankerConfig:
|
||||
"""Configuration for the reranker model."""
|
||||
model_name: str = "BAAI/bge-reranker-v2-m3"
|
||||
device: str = "mps" # Use MPS for Mac M1/M2
|
||||
batch_size: int = 32
|
||||
max_length: int = 8192 # Increased to match HARD_LIMIT in chunking
|
||||
use_fp16: bool = True # Use half precision for faster inference on Mac
|
||||
|
||||
@dataclass
|
||||
class PipelineConfig:
|
||||
"""Configuration for the retrieval pipeline."""
|
||||
services: ServiceConfig = field(default_factory=ServiceConfig)
|
||||
reranker: RerankerConfig = field(default_factory=RerankerConfig)
|
||||
|
||||
# Retrieval settings
|
||||
default_top_k: int = 20 # Number of candidates to retrieve from each service
|
||||
rerank_top_k: int = 10 # Number of results after reranking
|
||||
|
||||
# Fusion settings (see fusion.py)
|
||||
# "rrf": Reciprocal Rank Fusion (rank-only, robust); "weighted": weighted
|
||||
# min-max normalized score fusion; "avg_rank": legacy average-rank ordering.
|
||||
fusion_method: str = "rrf"
|
||||
rrf_k: int = 60 # RRF smoothing constant
|
||||
|
||||
# Logging
|
||||
debug: bool = True
|
||||
show_scores: bool = True # Show all scores in response for educational purposes
|
||||
|
||||
# Server settings
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 4242 # Default port for retrieval pipeline
|
||||
|
||||
@classmethod
|
||||
def from_env(cls):
|
||||
"""Create config from environment variables."""
|
||||
config = cls()
|
||||
if os.getenv("PIPELINE_PORT"):
|
||||
config.port = int(os.getenv("PIPELINE_PORT"))
|
||||
if os.getenv("PIPELINE_HOST"):
|
||||
config.host = os.getenv("PIPELINE_HOST")
|
||||
if os.getenv("DEBUG"):
|
||||
config.debug = os.getenv("DEBUG").lower() == "true"
|
||||
return config
|
||||
@@ -0,0 +1,281 @@
|
||||
"""Demo script showcasing dense vs sparse embedding strengths.
|
||||
|
||||
Service Configuration:
|
||||
- Dense Embedding: http://localhost:4240
|
||||
- Sparse Embedding: http://localhost:4241
|
||||
- Retrieval Pipeline: http://localhost:4242
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import httpx
|
||||
from typing import Dict, List
|
||||
import json
|
||||
|
||||
class RetrievalDemo:
|
||||
"""Demo for the retrieval pipeline."""
|
||||
|
||||
def __init__(self, pipeline_url: str = "http://localhost:4242"):
|
||||
self.pipeline_url = pipeline_url
|
||||
|
||||
async def index_document(self, text: str, doc_id: str, metadata: Dict = None):
|
||||
"""Index a document."""
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(
|
||||
f"{self.pipeline_url}/index",
|
||||
json={"text": text, "doc_id": doc_id, "metadata": metadata or {}}
|
||||
)
|
||||
return response.json()
|
||||
|
||||
async def search(self, query: str, mode: str = "hybrid"):
|
||||
"""Search for documents."""
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(
|
||||
f"{self.pipeline_url}/search",
|
||||
json={"query": query, "mode": mode, "top_k": 10, "rerank_top_k": 5}
|
||||
)
|
||||
return response.json()
|
||||
|
||||
async def clear(self):
|
||||
"""Clear all documents."""
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.delete(f"{self.pipeline_url}/clear")
|
||||
return response.json()
|
||||
|
||||
async def main():
|
||||
"""Run the demonstration."""
|
||||
demo = RetrievalDemo()
|
||||
|
||||
print("="*80)
|
||||
print("RETRIEVAL PIPELINE DEMONSTRATION")
|
||||
print("Showcasing Dense vs Sparse Embedding Strengths")
|
||||
print("="*80)
|
||||
|
||||
# Clear existing documents
|
||||
print("\nClearing existing documents...")
|
||||
await demo.clear()
|
||||
|
||||
# Create diverse test documents
|
||||
documents = [
|
||||
# Category 1: Programming Languages (for semantic similarity)
|
||||
{
|
||||
"doc_id": "prog_python",
|
||||
"text": "Python is renowned for its clean syntax and readability, making it ideal for beginners and experts alike.",
|
||||
"metadata": {"category": "programming", "subcategory": "languages"}
|
||||
},
|
||||
{
|
||||
"doc_id": "prog_javascript",
|
||||
"text": "JavaScript powers interactive web applications and runs in browsers worldwide.",
|
||||
"metadata": {"category": "programming", "subcategory": "languages"}
|
||||
},
|
||||
{
|
||||
"doc_id": "prog_rust",
|
||||
"text": "Rust provides memory safety without garbage collection through its ownership system.",
|
||||
"metadata": {"category": "programming", "subcategory": "languages"}
|
||||
},
|
||||
|
||||
# Category 2: Machine Learning (for concept matching)
|
||||
{
|
||||
"doc_id": "ml_intro",
|
||||
"text": "Artificial intelligence enables computers to learn from data and make decisions.",
|
||||
"metadata": {"category": "AI", "subcategory": "intro"}
|
||||
},
|
||||
{
|
||||
"doc_id": "ml_deep",
|
||||
"text": "Deep neural networks consist of multiple layers that progressively extract features.",
|
||||
"metadata": {"category": "AI", "subcategory": "deep_learning"}
|
||||
},
|
||||
{
|
||||
"doc_id": "ml_nlp",
|
||||
"text": "Natural language processing helps machines understand and generate human text.",
|
||||
"metadata": {"category": "AI", "subcategory": "NLP"}
|
||||
},
|
||||
|
||||
# Category 3: Specific Technical Terms (for exact matching)
|
||||
{
|
||||
"doc_id": "error_404",
|
||||
"text": "HTTP status code 404 indicates that the requested resource was not found on the server.",
|
||||
"metadata": {"category": "errors", "code": "404"}
|
||||
},
|
||||
{
|
||||
"doc_id": "error_500",
|
||||
"text": "HTTP status code 500 represents an internal server error that prevented request fulfillment.",
|
||||
"metadata": {"category": "errors", "code": "500"}
|
||||
},
|
||||
{
|
||||
"doc_id": "api_key",
|
||||
"text": "The API key XK9-2B4-7Q1 provides access to premium features of the service.",
|
||||
"metadata": {"category": "authentication", "type": "api_key"}
|
||||
},
|
||||
|
||||
# Category 4: Multilingual Content
|
||||
{
|
||||
"doc_id": "ml_chinese",
|
||||
"text": "机器学习是人工智能的核心技术，通过数据训练模型来解决问题。",
|
||||
"metadata": {"category": "AI", "language": "chinese"}
|
||||
},
|
||||
{
|
||||
"doc_id": "ml_spanish",
|
||||
"text": "El aprendizaje automático permite a las computadoras aprender sin programación explícita.",
|
||||
"metadata": {"category": "AI", "language": "spanish"}
|
||||
},
|
||||
|
||||
# Category 5: People and Names
|
||||
{
|
||||
"doc_id": "person_turing",
|
||||
"text": "Alan Turing pioneered computer science and artificial intelligence in the 20th century.",
|
||||
"metadata": {"category": "people", "field": "computer_science"}
|
||||
},
|
||||
{
|
||||
"doc_id": "person_lecun",
|
||||
"text": "Yann LeCun developed convolutional neural networks that revolutionized computer vision.",
|
||||
"metadata": {"category": "people", "field": "deep_learning"}
|
||||
}
|
||||
]
|
||||
|
||||
# Index all documents
|
||||
print(f"\nIndexing {len(documents)} documents...")
|
||||
for doc in documents:
|
||||
result = await demo.index_document(
|
||||
text=doc["text"],
|
||||
doc_id=doc["doc_id"],
|
||||
metadata=doc["metadata"]
|
||||
)
|
||||
success = result.get("success", False)
|
||||
status = "✓" if success else "✗"
|
||||
print(f" {status} {doc['doc_id']}: {doc['text'][:60]}...")
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("DEMONSTRATION QUERIES")
|
||||
print("="*80)
|
||||
|
||||
# Test queries demonstrating different strengths
|
||||
test_queries = [
|
||||
{
|
||||
"query": "code readability and simplicity",
|
||||
"description": "Semantic similarity - Dense should excel",
|
||||
"expected_strong": "dense",
|
||||
"explanation": "Dense embeddings understand 'readability' relates to Python even without exact match"
|
||||
},
|
||||
{
|
||||
"query": "XK9-2B4-7Q1",
|
||||
"description": "Exact code match - Sparse should excel",
|
||||
"expected_strong": "sparse",
|
||||
"explanation": "Sparse search finds exact API key string"
|
||||
},
|
||||
{
|
||||
"query": "AI learning from examples",
|
||||
"description": "Conceptual understanding - Dense should excel",
|
||||
"expected_strong": "dense",
|
||||
"explanation": "Dense understands AI/ML concepts without exact terminology"
|
||||
},
|
||||
{
|
||||
"query": "404",
|
||||
"description": "Specific error code - Sparse should excel",
|
||||
"expected_strong": "sparse",
|
||||
"explanation": "Sparse matches exact error code"
|
||||
},
|
||||
{
|
||||
"query": "人工智能",
|
||||
"description": "Cross-lingual search (Chinese for AI) - Dense should excel",
|
||||
"expected_strong": "dense",
|
||||
"explanation": "Dense embeddings (BGE-M3) handle multiple languages"
|
||||
},
|
||||
{
|
||||
"query": "Yann LeCun",
|
||||
"description": "Exact name search - Sparse should excel",
|
||||
"expected_strong": "sparse",
|
||||
"explanation": "Sparse finds exact person name"
|
||||
},
|
||||
{
|
||||
"query": "web browser programming",
|
||||
"description": "Semantic context - Dense should excel",
|
||||
"expected_strong": "dense",
|
||||
"explanation": "Dense connects 'web browser' with JavaScript"
|
||||
}
|
||||
]
|
||||
|
||||
# Run each test query
|
||||
for test in test_queries:
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Query: '{test['query']}'")
|
||||
print(f"Type: {test['description']}")
|
||||
print(f"Expected winner: {test['expected_strong']}")
|
||||
print(f"Reason: {test['explanation']}")
|
||||
print("-"*60)
|
||||
|
||||
# Run search in all three modes
|
||||
results = {}
|
||||
for mode in ["dense", "sparse", "hybrid"]:
|
||||
result = await demo.search(test["query"], mode=mode)
|
||||
|
||||
# Extract top results
|
||||
if mode == "dense":
|
||||
top_docs = result.get("dense_results", [])[:3]
|
||||
elif mode == "sparse":
|
||||
top_docs = result.get("sparse_results", [])[:3]
|
||||
else: # hybrid
|
||||
top_docs = result.get("reranked_results", [])[:3]
|
||||
|
||||
results[mode] = top_docs
|
||||
|
||||
# Print results for this mode
|
||||
print(f"\n{mode.upper()} Results:")
|
||||
for i, doc in enumerate(top_docs, 1):
|
||||
doc_id = doc.get("doc_id", "unknown")
|
||||
score = doc.get("score") or doc.get("rerank_score", 0)
|
||||
print(f" {i}. {doc_id} (score: {score:.4f})")
|
||||
|
||||
# Analyze which mode performed best
|
||||
print(f"\nAnalysis:")
|
||||
if test["expected_strong"] == "dense":
|
||||
if results["dense"] and results["sparse"]:
|
||||
dense_top = results["dense"][0]["doc_id"] if results["dense"] else None
|
||||
sparse_top = results["sparse"][0]["doc_id"] if results["sparse"] else None
|
||||
if dense_top != sparse_top:
|
||||
print(f" ✓ Dense found different (likely better) result: {dense_top}")
|
||||
print(f" ✓ Sparse found: {sparse_top}")
|
||||
elif test["expected_strong"] == "sparse":
|
||||
if results["sparse"]:
|
||||
print(f" ✓ Sparse found exact match: {results['sparse'][0]['doc_id']}")
|
||||
|
||||
# Show hybrid performance
|
||||
if results["hybrid"]:
|
||||
print(f" ✓ Hybrid (reranked) top result: {results['hybrid'][0]['doc_id']}")
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("DEMONSTRATION COMPLETE")
|
||||
print("="*80)
|
||||
print("\nKey Takeaways:")
|
||||
print("1. Dense embeddings excel at semantic similarity and concepts")
|
||||
print("2. Sparse search excels at exact matches and specific terms")
|
||||
print("3. Hybrid search with reranking combines the best of both")
|
||||
print("4. BGE-M3 dense embeddings support multilingual search")
|
||||
print("5. BM25 sparse search is unbeatable for exact string matching")
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Retrieval pipeline demonstration")
|
||||
parser.add_argument("--url", default="http://localhost:4242",
|
||||
help="Pipeline service URL")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Check if services are running
|
||||
print("Checking if services are available...")
|
||||
print(f"Pipeline URL: {args.url}")
|
||||
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except httpx.ConnectError:
|
||||
print("\nError: Could not connect to the retrieval pipeline service.")
|
||||
print("Please ensure all services are running:")
|
||||
print(" 1. Dense embedding service (port 4240)")
|
||||
print(" 2. Sparse embedding service (port 4241)")
|
||||
print(" 3. Retrieval pipeline (port 4242)")
|
||||
print("\nRun: ./restart_services.sh")
|
||||
sys.exit(1)
|
||||
except KeyboardInterrupt:
|
||||
print("\nDemo interrupted by user")
|
||||
sys.exit(0)
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Document store for the retrieval pipeline."""
|
||||
|
||||
from typing import Dict, Any, List, Optional
|
||||
from datetime import datetime
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class DocumentStore:
|
||||
"""In-memory document store for educational purposes."""
|
||||
|
||||
def __init__(self):
|
||||
self.documents: Dict[str, Dict[str, Any]] = {}
|
||||
self.metadata_index: Dict[str, List[str]] = {} # Index by metadata fields
|
||||
|
||||
def add_document(self, doc_id: str, text: str, metadata: Optional[Dict[str, Any]] = None) -> None:
|
||||
"""Add a document to the store."""
|
||||
new_metadata = metadata or {}
|
||||
if doc_id in self.documents:
|
||||
old_metadata = self.documents[doc_id].get("metadata") or {}
|
||||
for key in list(old_metadata.keys()):
|
||||
if key not in new_metadata:
|
||||
if key in self.metadata_index and doc_id in self.metadata_index[key]:
|
||||
self.metadata_index[key].remove(doc_id)
|
||||
if not self.metadata_index[key]:
|
||||
del self.metadata_index[key]
|
||||
|
||||
self.documents[doc_id] = {
|
||||
"doc_id": doc_id,
|
||||
"text": text,
|
||||
"metadata": new_metadata,
|
||||
"indexed_at": datetime.now().isoformat()
|
||||
}
|
||||
|
||||
# Update metadata index
|
||||
for key in new_metadata:
|
||||
if key not in self.metadata_index:
|
||||
self.metadata_index[key] = []
|
||||
if doc_id not in self.metadata_index[key]:
|
||||
self.metadata_index[key].append(doc_id)
|
||||
|
||||
logger.debug(f"Added document {doc_id} to store")
|
||||
|
||||
def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""Get a document by ID."""
|
||||
return self.documents.get(doc_id)
|
||||
|
||||
def get_documents(self, doc_ids: List[str]) -> List[Dict[str, Any]]:
|
||||
"""Get multiple documents by IDs."""
|
||||
docs = []
|
||||
for doc_id in doc_ids:
|
||||
doc = self.get_document(doc_id)
|
||||
if doc:
|
||||
docs.append(doc)
|
||||
return docs
|
||||
|
||||
def delete_document(self, doc_id: str) -> bool:
|
||||
"""Delete a document from the store."""
|
||||
if doc_id in self.documents:
|
||||
doc = self.documents[doc_id]
|
||||
|
||||
# Remove from metadata index
|
||||
if doc.get("metadata"):
|
||||
for key in doc["metadata"]:
|
||||
if key in self.metadata_index and doc_id in self.metadata_index[key]:
|
||||
self.metadata_index[key].remove(doc_id)
|
||||
if not self.metadata_index[key]:
|
||||
del self.metadata_index[key]
|
||||
|
||||
del self.documents[doc_id]
|
||||
logger.debug(f"Deleted document {doc_id} from store")
|
||||
return True
|
||||
return False
|
||||
|
||||
def list_documents(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
|
||||
"""List documents with pagination."""
|
||||
doc_ids = list(self.documents.keys())[offset:offset + limit]
|
||||
return [self.documents[doc_id] for doc_id in doc_ids]
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Clear all documents."""
|
||||
self.documents.clear()
|
||||
self.metadata_index.clear()
|
||||
logger.info("Cleared all documents from store")
|
||||
|
||||
def size(self) -> int:
|
||||
"""Get the number of documents."""
|
||||
return len(self.documents)
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
"""Get store statistics."""
|
||||
return {
|
||||
"total_documents": self.size(),
|
||||
"metadata_fields": list(self.metadata_index.keys()),
|
||||
"metadata_distribution": {
|
||||
key: len(values) for key, values in self.metadata_index.items()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,771 @@
|
||||
"""混合检索流水线离线评测 CLI。
|
||||
|
||||
本脚本把整条检索流水线——分块(chunk) → 嵌入(embed) → 检索(retrieve) →
|
||||
融合(fuse) → 重排(rerank)——完整地跑在**单进程、可离线**的环境里，并在一个带
|
||||
标注答案的小型评测集上，逐阶段对比各方法的检索质量。它不依赖 dense/sparse 微服务
|
||||
（4240/4241/4242 端口），因此可以脱离服务、直接用本地模型复现「每加一个阶段、指标如何
|
||||
提升」这一核心结论。
|
||||
|
||||
各阶段使用的本地组件：
|
||||
- 稀疏检索(sparse) : BM25（纯 Python，rank_bm25，无需下载模型）
|
||||
- 稠密检索(dense) : 本地句向量模型（默认 Qwen3-Embedding-0.6B，多语言，
|
||||
通过 transformers 加载；也可换成 BGE-M3 等）
|
||||
- 融合(fuse) : 见 fusion.py，RRF 与加权归一化两种策略
|
||||
- 重排(rerank) : 交叉编码器（默认 cross-encoder/ms-marco-MiniLM-L-6-v2）
|
||||
|
||||
默认行为（不带任何参数）：在内置评测集上评测
|
||||
BM25 / Dense / Hybrid-RRF / Hybrid-Weighted / Hybrid-RRF+Rerank 五种配置，
|
||||
打印 Recall@k、MRR、nDCG@k 对比表。
|
||||
|
||||
示例：
|
||||
python evaluate.py # 内置评测集，完整对比表
|
||||
python evaluate.py --top-k 10 --rerank-top-k 5
|
||||
python evaluate.py --no-rerank # 跳过重排阶段
|
||||
python evaluate.py --embed-model BAAI/bge-m3 --pooling cls
|
||||
python evaluate.py --query "怎样提升检索精度" # 单条查询、逐阶段排名追踪
|
||||
python evaluate.py --corpus my_corpus.json --queries my_queries.json --output result.json
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||||
|
||||
from fusion import fuse
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 内置评测集：直接复用 test_client.py 中的教育性测试案例（语义相似 / 精确名称 /
|
||||
# 多语言 / 技术代码四类），其 expected 字段即人工标注的相关文档，作为评测金标准。
|
||||
# 另加两篇较长文档用于演示分块(chunk)阶段。
|
||||
# ---------------------------------------------------------------------------
|
||||
DEFAULT_CORPUS: List[Dict[str, Any]] = [
|
||||
# --- 近似重复的代码簇（稀疏占优、稠密翻车）---
|
||||
# 各条文本几乎完全相同，只有型号代码不同；稠密向量几乎无法区分同一簇内的成员，
|
||||
# 稀疏检索靠精确词项匹配却能一击命中。簇越大，稠密选错的概率越高。
|
||||
{"doc_id": "xr_7001", "text": "Product model XR-7001 is a smartphone available now."},
|
||||
{"doc_id": "xr_7002", "text": "Product model XR-7002 is a smartphone available now."},
|
||||
{"doc_id": "xr_7003", "text": "Product model XR-7003 is a smartphone available now."},
|
||||
{"doc_id": "xr_7004", "text": "Product model XR-7004 is a smartphone available now."},
|
||||
{"doc_id": "xr_7005", "text": "Product model XR-7005 is a smartphone available now."},
|
||||
{"doc_id": "xr_7006", "text": "Product model XR-7006 is a smartphone available now."},
|
||||
# 近似重复的 HTTP 错误码簇（稀疏占优、稠密翻车）
|
||||
{"doc_id": "http_400", "text": "The HTTP-400 response is a client error status code."},
|
||||
{"doc_id": "http_401", "text": "The HTTP-401 response is a client error status code."},
|
||||
{"doc_id": "http_403", "text": "The HTTP-403 response is a client error status code."},
|
||||
{"doc_id": "http_404", "text": "The HTTP-404 response is a client error status code."},
|
||||
{"doc_id": "http_500", "text": "The HTTP-500 response is a server error status code."},
|
||||
# --- 语义改写簇（稠密占优、稀疏翻车）---
|
||||
# 查询与文档几乎没有共同词，稀疏 BM25 无从匹配，稠密靠语义命中。
|
||||
{"doc_id": "sem_readable", "text": "The language emphasizes clean, readable code that newcomers can pick up quickly."},
|
||||
{"doc_id": "sem_gc", "text": "Automatic memory management frees developers from manually releasing objects."},
|
||||
{"doc_id": "sem_photo", "text": "Green plants convert sunlight into chemical energy stored as sugars."},
|
||||
{"doc_id": "sem_crypto", "text": "Encryption scrambles a message so that only the intended recipient can read it."},
|
||||
# Exact proper-name cluster: sparse matching should preserve the complete
|
||||
# name while dense retrieval sees several near-duplicates.
|
||||
{"doc_id": "name_alexander_humphrey", "text": "Alexander Humphrey designed the Aurora scheduling protocol in 2019."},
|
||||
{"doc_id": "name_alexander_hughes", "text": "Alexander Hughes designed the Borealis scheduling protocol in 2019."},
|
||||
{"doc_id": "name_amelia_humphrey", "text": "Amelia Humphrey designed the Celeste scheduling protocol in 2020."},
|
||||
# 较长文档：话题彼此独立，用于演示分块阶段（会被切成多个 chunk 后再检索）
|
||||
{"doc_id": "doc_watercycle", "text": (
|
||||
"The water cycle describes how water moves continuously between the ocean, the atmosphere and the land. "
|
||||
"Heat from the sun evaporates water from the sea surface into vapor that rises high into the sky. "
|
||||
"As the vapor cools it condenses into tiny droplets that gather to form clouds. "
|
||||
"When the droplets grow heavy enough they fall back to the ground as rain or snow, "
|
||||
"and rivers eventually carry that water back to the ocean, closing the loop."
|
||||
)},
|
||||
{"doc_id": "doc_volcano", "text": (
|
||||
"A volcano forms where molten rock called magma rises from deep inside the planet toward the surface. "
|
||||
"Magma collects in a chamber beneath the crust, and mounting pressure forces it upward through cracks. "
|
||||
"During an eruption the magma bursts out as lava, ash and gas, which pile up around the vent. "
|
||||
"Layer after layer of cooled lava slowly builds the cone-shaped mountain we recognize as a volcano."
|
||||
)},
|
||||
]
|
||||
|
||||
DEFAULT_QUERIES: List[Dict[str, Any]] = [
|
||||
# 精确代码查询：稀疏一击命中，稠密难辨近似型号（expected 为唯一正确答案）
|
||||
{"query": "XR-7003", "expected": ["xr_7003"], "category": "technical-code"},
|
||||
{"query": "XR-7005", "expected": ["xr_7005"], "category": "technical-code"},
|
||||
{"query": "HTTP-403", "expected": ["http_403"], "category": "technical-code"},
|
||||
{"query": "HTTP-400", "expected": ["http_400"], "category": "technical-code"},
|
||||
{"query": "Alexander Humphrey", "expected": ["name_alexander_humphrey"], "category": "exact-name"},
|
||||
# 语义改写查询：与文档几乎无共同词，稠密靠语义命中，稀疏无从匹配
|
||||
{"query": "a beginner friendly language with tidy syntax", "expected": ["sem_readable"], "category": "semantic"},
|
||||
{"query": "reclaiming unused heap space without programmer effort", "expected": ["sem_gc"], "category": "semantic"},
|
||||
{"query": "how vegetation turns light into food", "expected": ["sem_photo"], "category": "semantic"},
|
||||
{"query": "hiding a note so eavesdroppers cannot understand it", "expected": ["sem_crypto"], "category": "semantic"},
|
||||
# Cross-lingual query with no shared lexical terms; the answer remains the
|
||||
# same English photosynthesis passage used by the semantic query above.
|
||||
{"query": "植物如何把阳光转化为食物", "expected": ["sem_photo"], "category": "multilingual"},
|
||||
# 长文档语义查询：命中的长文档会先被分块，再由某个 chunk 召回、重排
|
||||
{"query": "how does water move between the ocean and the sky", "expected": ["doc_watercycle"], "category": "semantic"},
|
||||
{"query": "how are volcanoes formed from molten rock", "expected": ["doc_volcano"], "category": "semantic"},
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 分块(chunk)
|
||||
# ---------------------------------------------------------------------------
|
||||
def chunk_text(text: str, chunk_size: int, overlap: int) -> List[str]:
|
||||
"""按字符窗口把文档切成带重叠的 chunk。
|
||||
|
||||
短文档（长度 <= chunk_size）原样返回单个 chunk。真实场景中 chunk 是检索的最小
|
||||
单元；这里用字符级滑窗保持实现简单、语言无关。
|
||||
|
||||
Args:
|
||||
text: 原始文档文本。
|
||||
chunk_size: 每个 chunk 的最大字符数。
|
||||
overlap: 相邻 chunk 的重叠字符数。
|
||||
|
||||
Returns:
|
||||
chunk 文本列表（至少一个）。
|
||||
"""
|
||||
text = text.strip()
|
||||
if chunk_size <= 0 or len(text) <= chunk_size:
|
||||
return [text]
|
||||
|
||||
step = max(1, chunk_size - overlap)
|
||||
chunks = []
|
||||
for start in range(0, len(text), step):
|
||||
piece = text[start:start + chunk_size].strip()
|
||||
if piece:
|
||||
chunks.append(piece)
|
||||
if start + chunk_size >= len(text):
|
||||
break
|
||||
return chunks or [text]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 分词（BM25 用）：保留英文词/数字/带连字符或下划线的代码，CJK 走 jieba + 单字
|
||||
# ---------------------------------------------------------------------------
|
||||
_TOKEN_RE = re.compile(r"[a-z0-9]+(?:[-_][a-z0-9]+)*|[一-鿿]+")
|
||||
|
||||
|
||||
def tokenize(text: str) -> List[str]:
|
||||
"""把文本切成 BM25 词项。
|
||||
|
||||
- 英文单词、纯数字、以及像 ``http-403`` / ``max_buffer_size`` / ``xr-7000``
|
||||
这样的技术代码会被整体保留（连字符、下划线不切开），保证精确匹配。
|
||||
- 连续 CJK 片段同时产出 jieba 分词结果与单字，增强中文召回鲁棒性。
|
||||
"""
|
||||
tokens: List[str] = []
|
||||
for match in _TOKEN_RE.finditer(text.lower()):
|
||||
span = match.group()
|
||||
if "一" <= span[0] <= "鿿":
|
||||
try:
|
||||
import jieba
|
||||
tokens.extend(w for w in jieba.cut(span) if w.strip())
|
||||
except Exception:
|
||||
pass
|
||||
tokens.extend(list(span))
|
||||
else:
|
||||
tokens.append(span)
|
||||
return tokens
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 稀疏检索：BM25
|
||||
# ---------------------------------------------------------------------------
|
||||
class BM25Retriever:
|
||||
"""基于 rank_bm25 的 BM25 检索器（chunk 级）。"""
|
||||
|
||||
def __init__(self, chunk_ids: List[str], chunk_texts: List[str]):
|
||||
from rank_bm25 import BM25Okapi
|
||||
|
||||
self.chunk_ids = chunk_ids
|
||||
self.tokenized = [tokenize(t) for t in chunk_texts]
|
||||
self.bm25 = BM25Okapi(self.tokenized)
|
||||
|
||||
def search(self, query: str, top_k: int) -> List[Tuple[str, float]]:
|
||||
"""返回 (chunk_id, score) 列表，按分数降序，只保留正分。"""
|
||||
scores = self.bm25.get_scores(tokenize(query))
|
||||
ranked = sorted(zip(self.chunk_ids, scores), key=lambda kv: kv[1], reverse=True)
|
||||
return [(cid, float(s)) for cid, s in ranked[:top_k] if s > 0]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 稠密检索：本地句向量模型（transformers）
|
||||
# ---------------------------------------------------------------------------
|
||||
class DenseEncoder:
|
||||
"""用 transformers 加载本地句向量模型，做稠密检索。"""
|
||||
|
||||
def __init__(self, model_name: str, pooling: str, device: str,
|
||||
query_instruct: str = "", max_length: int = 256):
|
||||
import torch
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
self.torch = torch
|
||||
self.device = device
|
||||
self.max_length = max_length
|
||||
self.pooling = self._resolve_pooling(pooling, model_name)
|
||||
# 指令式检索模型（如 Qwen3-Embedding，last-token 池化）要求查询侧带任务指令；
|
||||
# mean/cls 池化的模型（MiniLM / BGE-M3）不需要，自动关闭。
|
||||
self.query_instruct = query_instruct if (query_instruct and self.pooling == "last") else ""
|
||||
# last-token pooling 需要左侧 padding，才能让最后一个位置对齐真实末词
|
||||
padding_side = "left" if self.pooling == "last" else "right"
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side=padding_side)
|
||||
self.model = AutoModel.from_pretrained(model_name).to(device).eval()
|
||||
|
||||
@staticmethod
|
||||
def _resolve_pooling(pooling: str, model_name: str) -> str:
|
||||
if pooling != "auto":
|
||||
return pooling
|
||||
name = model_name.lower()
|
||||
if "qwen" in name:
|
||||
return "last"
|
||||
if "bge-m3" in name or "bge-large" in name or "bge-base" in name:
|
||||
return "cls"
|
||||
return "mean"
|
||||
|
||||
def _pool(self, last_hidden, attention_mask):
|
||||
torch = self.torch
|
||||
if self.pooling == "cls":
|
||||
return last_hidden[:, 0]
|
||||
if self.pooling == "last":
|
||||
return last_hidden[:, -1]
|
||||
# mean pooling
|
||||
mask = attention_mask.unsqueeze(-1).float()
|
||||
return (last_hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
|
||||
|
||||
def encode(self, texts: Sequence[str], is_query: bool = False, batch_size: int = 16):
|
||||
torch = self.torch
|
||||
if is_query and self.query_instruct:
|
||||
texts = [f"Instruct: {self.query_instruct}\nQuery:{t}" for t in texts]
|
||||
vectors = []
|
||||
for start in range(0, len(texts), batch_size):
|
||||
batch = list(texts[start:start + batch_size])
|
||||
pooled = self._forward(batch)
|
||||
# 某些模型在 mps 上前向会出 NaN（transformers 5.x + 某些权重）；
|
||||
# 检测到后永久退回 CPU 重算，保证向量有限、结果可复现。
|
||||
if self.device != "cpu" and torch.isnan(pooled).any():
|
||||
self.device = "cpu"
|
||||
self.model = self.model.to("cpu")
|
||||
pooled = self._forward(batch)
|
||||
pooled = torch.nn.functional.normalize(pooled.float(), p=2, dim=1)
|
||||
vectors.append(pooled.cpu())
|
||||
return torch.cat(vectors, dim=0)
|
||||
|
||||
def _forward(self, batch: List[str]):
|
||||
torch = self.torch
|
||||
enc = self.tokenizer(
|
||||
batch, padding=True, truncation=True,
|
||||
max_length=self.max_length, return_tensors="pt",
|
||||
).to(self.device)
|
||||
with torch.no_grad():
|
||||
out = self.model(**enc)
|
||||
return self._pool(out.last_hidden_state, enc["attention_mask"])
|
||||
|
||||
|
||||
class DenseRetriever:
|
||||
"""基于稠密向量余弦相似度的 chunk 级检索器。"""
|
||||
|
||||
def __init__(self, encoder: DenseEncoder, chunk_ids: List[str], chunk_texts: List[str]):
|
||||
self.encoder = encoder
|
||||
self.chunk_ids = chunk_ids
|
||||
self.matrix = encoder.encode(chunk_texts) # [N, D], 已归一化
|
||||
|
||||
def search(self, query: str, top_k: int) -> List[Tuple[str, float]]:
|
||||
q = self.encoder.encode([query], is_query=True)[0]
|
||||
sims = (self.matrix @ q).tolist()
|
||||
ranked = sorted(zip(self.chunk_ids, sims), key=lambda kv: kv[1], reverse=True)
|
||||
return [(cid, float(s)) for cid, s in ranked[:top_k]]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 重排：交叉编码器（cross-encoder）
|
||||
# ---------------------------------------------------------------------------
|
||||
class CrossEncoderReranker:
|
||||
"""用交叉编码器对候选做精排。
|
||||
|
||||
在 transformers 5.x + 部分 BERT 权重上，fp32 前向可能出现 NaN；本类检测到 NaN 后
|
||||
自动回退到 CPU + float64 重算，保证输出有限、可复现。
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str, device: str, max_length: int = 512):
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
||||
|
||||
self.torch = torch
|
||||
self.device = device
|
||||
self.max_length = max_length
|
||||
self.model_name = model_name
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
self.model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device).eval()
|
||||
|
||||
def score(self, query: str, docs: Sequence[str]) -> List[float]:
|
||||
torch = self.torch
|
||||
if not docs:
|
||||
return []
|
||||
enc = self.tokenizer(
|
||||
[query] * len(docs), list(docs),
|
||||
padding=True, truncation=True, max_length=self.max_length, return_tensors="pt",
|
||||
).to(self.device)
|
||||
with torch.no_grad():
|
||||
logits = self.model(**enc).logits.squeeze(-1).float()
|
||||
if torch.isnan(logits).any():
|
||||
# 回退：CPU + float64 重算
|
||||
enc_cpu = {k: v.to("cpu") for k, v in enc.items()}
|
||||
model64 = self.model.to("cpu").double()
|
||||
with torch.no_grad():
|
||||
logits = model64(**enc_cpu).logits.squeeze(-1)
|
||||
self.model = self.model.to(self.device).float()
|
||||
return [float(x) for x in logits.reshape(-1).tolist()]
|
||||
|
||||
def rerank(self, query: str, candidates: List[Tuple[str, str]], top_k: int) -> List[Tuple[str, float]]:
|
||||
"""candidates: [(doc_id, text)] -> [(doc_id, rerank_score)] 降序，取 top_k。"""
|
||||
scores = self.score(query, [text for _, text in candidates])
|
||||
ranked = sorted(
|
||||
((doc_id, s) for (doc_id, _), s in zip(candidates, scores)),
|
||||
key=lambda kv: kv[1], reverse=True,
|
||||
)
|
||||
return ranked[:top_k]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# chunk 级结果 -> doc 级结果（同一文档取最高分的 chunk）
|
||||
# ---------------------------------------------------------------------------
|
||||
def chunks_to_docs(ranked_chunks: List[Tuple[str, float]], chunk_to_doc: Dict[str, str]) -> List[Tuple[str, float]]:
|
||||
best: Dict[str, float] = {}
|
||||
for chunk_id, score in ranked_chunks:
|
||||
doc_id = chunk_to_doc[chunk_id]
|
||||
if doc_id not in best or score > best[doc_id]:
|
||||
best[doc_id] = score
|
||||
return sorted(best.items(), key=lambda kv: kv[1], reverse=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 评测指标
|
||||
# ---------------------------------------------------------------------------
|
||||
def recall_at_k(ranked_ids: List[str], gold: Sequence[str], k: int) -> float:
|
||||
if not gold:
|
||||
return 0.0
|
||||
topk = set(ranked_ids[:k])
|
||||
return len(topk & set(gold)) / len(gold)
|
||||
|
||||
|
||||
def reciprocal_rank(ranked_ids: List[str], gold: Sequence[str]) -> float:
|
||||
gold_set = set(gold)
|
||||
for idx, doc_id in enumerate(ranked_ids, start=1):
|
||||
if doc_id in gold_set:
|
||||
return 1.0 / idx
|
||||
return 0.0
|
||||
|
||||
|
||||
def ndcg_at_k(ranked_ids: List[str], gold: Sequence[str], k: int) -> float:
|
||||
gold_set = set(gold)
|
||||
dcg = 0.0
|
||||
for idx, doc_id in enumerate(ranked_ids[:k], start=1):
|
||||
if doc_id in gold_set:
|
||||
dcg += 1.0 / math.log2(idx + 1)
|
||||
ideal_hits = min(len(gold_set), k)
|
||||
idcg = sum(1.0 / math.log2(i + 1) for i in range(1, ideal_hits + 1))
|
||||
return dcg / idcg if idcg > 0 else 0.0
|
||||
|
||||
|
||||
def aggregate_metrics(per_query_ranked: List[Tuple[List[str], Sequence[str]]], k: int) -> Dict[str, float]:
|
||||
n = len(per_query_ranked)
|
||||
if n == 0:
|
||||
return {"recall@k": 0.0, "mrr": 0.0, "ndcg@k": 0.0}
|
||||
recall = sum(recall_at_k(r, g, k) for r, g in per_query_ranked) / n
|
||||
mrr = sum(reciprocal_rank(r, g) for r, g in per_query_ranked) / n
|
||||
ndcg = sum(ndcg_at_k(r, g, k) for r, g in per_query_ranked) / n
|
||||
return {"recall@k": recall, "mrr": mrr, "ndcg@k": ndcg}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 流水线：为一条查询产出各方法的文档级排名
|
||||
# ---------------------------------------------------------------------------
|
||||
class Pipeline:
|
||||
def __init__(self, corpus, args):
|
||||
self.args = args
|
||||
self.chunk_ids: List[str] = []
|
||||
self.chunk_texts: List[str] = []
|
||||
self.chunk_to_doc: Dict[str, str] = {}
|
||||
self.doc_text: Dict[str, str] = {}
|
||||
|
||||
# 分块
|
||||
for doc in corpus:
|
||||
self.doc_text[doc["doc_id"]] = doc["text"]
|
||||
chunks = chunk_text(doc["text"], args.chunk_size, args.chunk_overlap)
|
||||
for i, chunk in enumerate(chunks):
|
||||
cid = f"{doc['doc_id']}::c{i}" if len(chunks) > 1 else doc["doc_id"]
|
||||
self.chunk_ids.append(cid)
|
||||
self.chunk_texts.append(chunk)
|
||||
self.chunk_to_doc[cid] = doc["doc_id"]
|
||||
|
||||
self.n_docs = len(corpus)
|
||||
self.n_chunks = len(self.chunk_ids)
|
||||
|
||||
# 稀疏索引
|
||||
self.bm25 = BM25Retriever(self.chunk_ids, self.chunk_texts)
|
||||
|
||||
# 稠密索引（可选）
|
||||
self.dense: Optional[DenseRetriever] = None
|
||||
if args.use_dense:
|
||||
encoder = DenseEncoder(args.embed_model, args.pooling, args.device,
|
||||
query_instruct=args.query_instruct)
|
||||
self.dense = DenseRetriever(encoder, self.chunk_ids, self.chunk_texts)
|
||||
|
||||
# 重排器（可选）
|
||||
self.reranker: Optional[CrossEncoderReranker] = None
|
||||
if args.use_rerank:
|
||||
self.reranker = CrossEncoderReranker(args.reranker_model, args.device)
|
||||
|
||||
def run_query(self, query: str) -> Dict[str, List[Tuple[str, float]]]:
|
||||
"""返回各方法的 doc 级排名 {method: [(doc_id, score)]}。"""
|
||||
top_k = self.args.top_k
|
||||
sparse_started = time.perf_counter()
|
||||
sparse_chunks = self.bm25.search(query, top_k)
|
||||
sparse_docs = chunks_to_docs(sparse_chunks, self.chunk_to_doc)
|
||||
sparse_ms = (time.perf_counter() - sparse_started) * 1000
|
||||
|
||||
out: Dict[str, List[Tuple[str, float]]] = {"sparse": sparse_docs}
|
||||
component_ms = {"sparse": sparse_ms}
|
||||
|
||||
if self.dense is not None:
|
||||
dense_started = time.perf_counter()
|
||||
dense_chunks = self.dense.search(query, top_k)
|
||||
dense_docs = chunks_to_docs(dense_chunks, self.chunk_to_doc)
|
||||
dense_ms = (time.perf_counter() - dense_started) * 1000
|
||||
component_ms["dense"] = dense_ms
|
||||
out["dense"] = dense_docs
|
||||
|
||||
ranked_lists = {"dense": dense_docs, "sparse": sparse_docs}
|
||||
weights = {"dense": self.args.dense_weight, "sparse": self.args.sparse_weight}
|
||||
rrf_started = time.perf_counter()
|
||||
rrf = fuse(ranked_lists, method="rrf", k=self.args.k_rrf, weights=weights)
|
||||
rrf_ms = (time.perf_counter() - rrf_started) * 1000
|
||||
weighted_started = time.perf_counter()
|
||||
weighted = fuse(ranked_lists, method="weighted", weights=weights)
|
||||
weighted_ms = (time.perf_counter() - weighted_started) * 1000
|
||||
component_ms.update({"rrf_fusion": rrf_ms, "weighted_fusion": weighted_ms})
|
||||
out["rrf"] = rrf
|
||||
out["weighted"] = weighted
|
||||
|
||||
if self.reranker is not None:
|
||||
# 对 RRF 融合的候选池 top-N 精排
|
||||
pool = [doc_id for doc_id, _ in rrf[: self.args.rerank_pool]]
|
||||
candidates = [(doc_id, self.doc_text[doc_id]) for doc_id in pool]
|
||||
rerank_started = time.perf_counter()
|
||||
reranked = self.reranker.rerank(query, candidates, self.args.rerank_top_k)
|
||||
component_ms["rerank"] = (time.perf_counter() - rerank_started) * 1000
|
||||
out["rerank"] = reranked
|
||||
|
||||
end_to_end_ms = {"sparse": sparse_ms}
|
||||
if "dense" in component_ms:
|
||||
retrieval_ms = sparse_ms + component_ms["dense"]
|
||||
end_to_end_ms.update(
|
||||
{
|
||||
"dense": component_ms["dense"],
|
||||
"rrf": retrieval_ms + component_ms["rrf_fusion"],
|
||||
"weighted": retrieval_ms + component_ms["weighted_fusion"],
|
||||
}
|
||||
)
|
||||
if "rerank" in component_ms:
|
||||
end_to_end_ms["rerank"] = (
|
||||
retrieval_ms + component_ms["rrf_fusion"] + component_ms["rerank"]
|
||||
)
|
||||
rank_changes = []
|
||||
if "rerank" in out:
|
||||
before = {doc_id: rank for rank, (doc_id, _) in enumerate(out["rrf"], 1)}
|
||||
after = {doc_id: rank for rank, (doc_id, _) in enumerate(out["rerank"], 1)}
|
||||
for doc_id in sorted(set(before) | set(after)):
|
||||
rank_changes.append(
|
||||
{
|
||||
"doc_id": doc_id,
|
||||
"rrf_rank": before.get(doc_id),
|
||||
"rerank_rank": after.get(doc_id),
|
||||
"delta": (
|
||||
before[doc_id] - after[doc_id]
|
||||
if doc_id in before and doc_id in after
|
||||
else None
|
||||
),
|
||||
}
|
||||
)
|
||||
self.last_trace = {
|
||||
"component_latency_ms": component_ms,
|
||||
"end_to_end_latency_ms": end_to_end_ms,
|
||||
"rank_changes": rank_changes,
|
||||
}
|
||||
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 输出：对比表 / 单条查询追踪
|
||||
# ---------------------------------------------------------------------------
|
||||
METHOD_LABELS = [
|
||||
("sparse", "BM25 (sparse)"),
|
||||
("dense", "Dense"),
|
||||
("rrf", "Hybrid-RRF"),
|
||||
("weighted", "Hybrid-Weighted"),
|
||||
("rerank", "Hybrid-RRF+Rerank"),
|
||||
]
|
||||
|
||||
|
||||
def run_evaluation(pipeline: Pipeline, queries, args) -> Dict[str, Any]:
|
||||
k = args.eval_k
|
||||
per_method: Dict[str, List[Tuple[List[str], Sequence[str]]]] = {m: [] for m, _ in METHOD_LABELS}
|
||||
per_query_records = []
|
||||
latency_by_method: Dict[str, List[float]] = {m: [] for m, _ in METHOD_LABELS}
|
||||
|
||||
t0 = time.time()
|
||||
for spec in queries:
|
||||
query = spec["query"]
|
||||
gold = spec.get("expected", [])
|
||||
results = pipeline.run_query(query)
|
||||
record = {
|
||||
"query": query,
|
||||
"expected": gold,
|
||||
"category": spec.get("category", "unspecified"),
|
||||
"methods": {},
|
||||
"trace": pipeline.last_trace,
|
||||
}
|
||||
for method, _ in METHOD_LABELS:
|
||||
if method not in results:
|
||||
continue
|
||||
ranked_ids = [doc_id for doc_id, _ in results[method]]
|
||||
per_method[method].append((ranked_ids, gold))
|
||||
latency_by_method[method].append(pipeline.last_trace["end_to_end_latency_ms"][method])
|
||||
record["methods"][method] = {
|
||||
"top": [{"doc_id": d, "score": round(s, 4)} for d, s in results[method][:5]],
|
||||
"recall@k": round(recall_at_k(ranked_ids, gold, k), 4),
|
||||
"mrr": round(reciprocal_rank(ranked_ids, gold), 4),
|
||||
"ndcg@k": round(ndcg_at_k(ranked_ids, gold, k), 4),
|
||||
}
|
||||
per_query_records.append(record)
|
||||
elapsed = time.time() - t0
|
||||
|
||||
summary = {}
|
||||
for method, _ in METHOD_LABELS:
|
||||
if per_method[method]:
|
||||
summary[method] = aggregate_metrics(per_method[method], k)
|
||||
values = sorted(latency_by_method[method])
|
||||
p95_index = min(len(values) - 1, math.ceil(0.95 * len(values)) - 1)
|
||||
summary[method]["latency_ms"] = {
|
||||
"mean": sum(values) / len(values),
|
||||
"p50": values[len(values) // 2],
|
||||
"p95": values[p95_index],
|
||||
}
|
||||
|
||||
return {
|
||||
"summary": summary,
|
||||
"per_query": per_query_records,
|
||||
"elapsed_sec": round(elapsed, 2),
|
||||
"eval_k": k,
|
||||
}
|
||||
|
||||
|
||||
def print_table(report: Dict[str, Any], pipeline: Pipeline, args) -> None:
|
||||
k = report["eval_k"]
|
||||
print("=" * 78)
|
||||
print("混合检索流水线 · 逐阶段评测对比")
|
||||
print("=" * 78)
|
||||
print(f"语料: {pipeline.n_docs} 篇文档 → {pipeline.n_chunks} 个 chunk "
|
||||
f"(chunk_size={args.chunk_size}, overlap={args.chunk_overlap})")
|
||||
print(f"查询: {len(report['per_query'])} 条 "
|
||||
f"稠密模型: {args.embed_model if args.use_dense else '(禁用)'} "
|
||||
f"重排模型: {args.reranker_model if args.use_rerank else '(禁用)'}")
|
||||
print(f"检索 top_k={args.top_k} 融合 k(RRF)={args.k_rrf} "
|
||||
f"重排候选池={args.rerank_pool} 评测截断 k={k} 设备={args.device}")
|
||||
print(f"耗时: {report['elapsed_sec']}s")
|
||||
print("-" * 78)
|
||||
header = f"{'Stage / Method':<22}{'Recall@'+str(k):>12}{'MRR':>12}{'nDCG@'+str(k):>12}"
|
||||
print(header)
|
||||
print("-" * 78)
|
||||
for method, label in METHOD_LABELS:
|
||||
if method not in report["summary"]:
|
||||
continue
|
||||
m = report["summary"][method]
|
||||
print(f"{label:<22}{m['recall@k']:>12.4f}{m['mrr']:>12.4f}{m['ndcg@k']:>12.4f}")
|
||||
print("-" * 78)
|
||||
print("读表：从上到下逐步加入 稠密检索 / 融合 / 重排 阶段，观察指标的变化。")
|
||||
print("=" * 78)
|
||||
|
||||
|
||||
def print_per_query(report: Dict[str, Any]) -> None:
|
||||
"""逐条查询打印各方法的 MRR，直观展示「单路会翻车、融合来兜底」。"""
|
||||
methods = [m for m, _ in METHOD_LABELS]
|
||||
short = {"sparse": "BM25", "dense": "Dense", "rrf": "RRF",
|
||||
"weighted": "Wgt", "rerank": "Rerank"}
|
||||
print("\n逐条查询 MRR 明细（1.00=正确文档排在第 1 位；粗看哪一路在哪类查询上翻车）")
|
||||
print("-" * 78)
|
||||
header = f"{'Query':<42}" + "".join(f"{short[m]:>7}" for m in methods)
|
||||
print(header)
|
||||
print("-" * 78)
|
||||
for rec in report["per_query"]:
|
||||
cells = ""
|
||||
for m in methods:
|
||||
if m in rec["methods"]:
|
||||
cells += f"{rec['methods'][m]['mrr']:>7.2f}"
|
||||
else:
|
||||
cells += f"{'-':>7}"
|
||||
q = rec["query"]
|
||||
q = q if len(q) <= 41 else q[:38] + "..."
|
||||
print(f"{q:<42}{cells}")
|
||||
print("=" * 78)
|
||||
|
||||
|
||||
def print_query_trace(pipeline: Pipeline, query: str, args) -> None:
|
||||
results = pipeline.run_query(query)
|
||||
print("=" * 78)
|
||||
print(f"单条查询逐阶段排名追踪 query = {query!r}")
|
||||
print(f"语料 {pipeline.n_docs} 篇 → {pipeline.n_chunks} chunk 设备={args.device}")
|
||||
print("=" * 78)
|
||||
for method, label in METHOD_LABELS:
|
||||
if method not in results:
|
||||
continue
|
||||
print(f"\n[{label}]")
|
||||
for rank, (doc_id, score) in enumerate(results[method][:5], start=1):
|
||||
snippet = pipeline.doc_text.get(doc_id, "")[:60].replace("\n", " ")
|
||||
print(f" {rank}. {doc_id:<14} score={score:8.4f} {snippet}")
|
||||
print("=" * 78)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
def detect_device(requested: str) -> str:
|
||||
if requested != "auto":
|
||||
return requested
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
return "cuda"
|
||||
if torch.backends.mps.is_available():
|
||||
return "mps"
|
||||
except Exception:
|
||||
pass
|
||||
return "cpu"
|
||||
|
||||
|
||||
def load_json(path: str):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="混合检索流水线离线评测 CLI（chunk→embed→retrieve→fuse→rerank，逐阶段对比）。",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=(
|
||||
"示例:\n"
|
||||
" python evaluate.py # 内置评测集，完整对比表\n"
|
||||
" python evaluate.py --no-rerank # 跳过重排阶段\n"
|
||||
" python evaluate.py --no-dense # 仅 BM25（纯离线、无需模型）\n"
|
||||
" python evaluate.py --query '怎样提升检索精度' # 单条查询逐阶段排名\n"
|
||||
" python evaluate.py --embed-model BAAI/bge-m3 --pooling cls\n"
|
||||
" python evaluate.py --output result.json # 结果同时写入 JSON\n"
|
||||
),
|
||||
)
|
||||
data = parser.add_argument_group("数据")
|
||||
data.add_argument("--corpus", help="语料 JSON 文件，格式 [{'doc_id','text'}...]；缺省用内置语料")
|
||||
data.add_argument("--queries", help="查询 JSON 文件，格式 [{'query','expected':[...]}...]；缺省用内置查询")
|
||||
data.add_argument("--query", help="单条查询模式：只对该查询做逐阶段排名追踪，不跑评测")
|
||||
data.add_argument("--limit-queries", type=int, default=0, help="只评测前 N 条查询（0=全部）")
|
||||
|
||||
stages = parser.add_argument_group("流水线阶段")
|
||||
stages.add_argument("--no-dense", dest="use_dense", action="store_false",
|
||||
help="禁用稠密检索（连带禁用融合与重排；退化为纯 BM25，完全离线无需模型）")
|
||||
stages.add_argument("--no-rerank", dest="use_rerank", action="store_false",
|
||||
help="禁用神经重排阶段")
|
||||
stages.set_defaults(use_dense=True, use_rerank=True)
|
||||
|
||||
chunk = parser.add_argument_group("分块")
|
||||
chunk.add_argument("--chunk-size", type=int, default=280, help="每个 chunk 的最大字符数（默认 280）")
|
||||
chunk.add_argument("--chunk-overlap", type=int, default=40, help="相邻 chunk 的重叠字符数（默认 40）")
|
||||
|
||||
retr = parser.add_argument_group("检索与融合")
|
||||
retr.add_argument("--top-k", type=int, default=10, help="每路检索召回的候选数（默认 10）")
|
||||
retr.add_argument("--k-rrf", type=int, default=60, help="RRF 平滑常数 k（默认 60）")
|
||||
retr.add_argument("--dense-weight", type=float, default=1.0, help="融合时稠密路权重（默认 1.0）")
|
||||
retr.add_argument("--sparse-weight", type=float, default=1.0, help="融合时稀疏路权重（默认 1.0）")
|
||||
|
||||
rer = parser.add_argument_group("重排")
|
||||
rer.add_argument("--rerank-pool", type=int, default=10, help="送入重排的候选池大小（取 RRF 融合的 top-N，默认 10）")
|
||||
rer.add_argument("--rerank-top-k", type=int, default=10, help="重排后返回的结果数（默认 10）")
|
||||
|
||||
model = parser.add_argument_group("模型")
|
||||
model.add_argument("--embed-model", default="sentence-transformers/all-MiniLM-L6-v2",
|
||||
help="稠密句向量模型（默认 sentence-transformers/all-MiniLM-L6-v2，约 90MB、英文为主；"
|
||||
"多语言语料请换 Qwen/Qwen3-Embedding-0.6B 或 BAAI/bge-m3）")
|
||||
model.add_argument("--pooling", default="auto", choices=["auto", "mean", "cls", "last"],
|
||||
help="句向量池化方式（auto 会按模型名自动选择：qwen→last, bge-m3→cls, 其余→mean）")
|
||||
model.add_argument("--query-instruct",
|
||||
default="Given a search query, retrieve relevant passages that answer the query",
|
||||
help="指令式检索模型的查询侧任务指令（仅对 last-token 池化的模型如 Qwen3-Embedding 生效）")
|
||||
model.add_argument("--reranker-model", default="BAAI/bge-reranker-base",
|
||||
help="交叉编码器重排模型（默认 BAAI/bge-reranker-base，多语言、首次运行约 1.1GB；"
|
||||
"生产可换更强的 BAAI/bge-reranker-v2-m3，轻量可换 cross-encoder/ms-marco-MiniLM-L-6-v2）")
|
||||
model.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps"],
|
||||
help="推理设备（默认 auto）")
|
||||
|
||||
out = parser.add_argument_group("评测与输出")
|
||||
out.add_argument("--eval-k", type=int, default=3, help="指标截断位置 k（Recall@k / nDCG@k，默认 3）")
|
||||
out.add_argument("--no-per-query", dest="show_per_query", action="store_false",
|
||||
help="不打印逐条查询的 MRR 明细矩阵")
|
||||
out.set_defaults(show_per_query=True)
|
||||
out.add_argument("--output", help="把完整结果（含每条查询明细）写入该 JSON 文件")
|
||||
out.add_argument("--offline", action="store_true", help="设置 HF_HUB_OFFLINE=1，强制只用本地缓存模型")
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = build_parser().parse_args()
|
||||
|
||||
if args.offline:
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||||
|
||||
args.device = detect_device(args.device)
|
||||
# 单查询追踪模式不影响 use_dense/use_rerank 语义，但重排依赖稠密融合池
|
||||
if not args.use_dense:
|
||||
args.use_rerank = False
|
||||
|
||||
corpus = load_json(args.corpus) if args.corpus else DEFAULT_CORPUS
|
||||
queries = load_json(args.queries) if args.queries else DEFAULT_QUERIES
|
||||
|
||||
try:
|
||||
pipeline = Pipeline(corpus, args)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f"[错误] 流水线初始化失败: {exc}", file=sys.stderr)
|
||||
print("提示：稠密/重排阶段需要本地句向量与交叉编码器模型；"
|
||||
"可用 --no-dense 退化为纯 BM25（完全离线），或用 --embed-model 指定已缓存模型。",
|
||||
file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if args.query:
|
||||
print_query_trace(pipeline, args.query, args)
|
||||
return 0
|
||||
|
||||
if args.limit_queries > 0:
|
||||
queries = queries[: args.limit_queries]
|
||||
|
||||
report = run_evaluation(pipeline, queries, args)
|
||||
print_table(report, pipeline, args)
|
||||
if args.show_per_query:
|
||||
print_per_query(report)
|
||||
|
||||
if args.output:
|
||||
payload = {
|
||||
"config": {
|
||||
"embed_model": args.embed_model if args.use_dense else None,
|
||||
"reranker_model": args.reranker_model if args.use_rerank else None,
|
||||
"top_k": args.top_k, "k_rrf": args.k_rrf, "eval_k": args.eval_k,
|
||||
"chunk_size": args.chunk_size, "chunk_overlap": args.chunk_overlap,
|
||||
"device": args.device,
|
||||
},
|
||||
**report,
|
||||
}
|
||||
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,151 @@
|
||||
"""Result fusion for hybrid retrieval.
|
||||
|
||||
This module implements the *fusion* stage of the hybrid retrieval pipeline —
|
||||
the step that merges the separately-ranked dense and sparse candidate lists into
|
||||
a single, unified candidate pool before neural reranking.
|
||||
|
||||
Two production-grade fusion strategies are provided, matching the two approaches
|
||||
discussed in the book (第3章「混合检索流水线」):
|
||||
|
||||
1. Reciprocal Rank Fusion (RRF)
|
||||
score(d) = Σ_r 1 / (k + rank_r(d))
|
||||
Only ranks are used, original scores are discarded. Robust and scale-free,
|
||||
because it never has to compare a cosine similarity against a BM25 score.
|
||||
|
||||
2. Weighted score fusion (min-max normalized)
|
||||
score(d) = Σ_r w_r * normalize_r(score_r(d))
|
||||
Keeps the original relevance signal, at the cost of having to align the two
|
||||
score scales via per-list min-max normalization.
|
||||
|
||||
Both functions take ranked lists of ``(doc_id, score)`` tuples (sorted by score
|
||||
descending) and return a fused list of ``(doc_id, fused_score)`` tuples, also
|
||||
sorted descending. A document that appears in only one list is still fused —
|
||||
its contribution from the missing list is simply zero.
|
||||
"""
|
||||
|
||||
from typing import Dict, List, Optional, Sequence, Tuple
|
||||
|
||||
RankedList = Sequence[Tuple[str, float]]
|
||||
|
||||
# Default smoothing constant for RRF. k=60 is the value from the original
|
||||
# Cormack et al. paper and the most common choice in practice; it compresses the
|
||||
# score gap between the very top ranks.
|
||||
DEFAULT_RRF_K = 60
|
||||
|
||||
|
||||
def _ranked_to_score_map(ranked: RankedList) -> Dict[str, float]:
|
||||
scores: Dict[str, float] = {}
|
||||
for doc_id, score in ranked:
|
||||
if doc_id not in scores:
|
||||
scores[doc_id] = score
|
||||
return scores
|
||||
def min_max_normalize(scores: Dict[str, float]) -> Dict[str, float]:
|
||||
"""Min-max normalize a mapping of doc_id -> score into the [0, 1] range.
|
||||
|
||||
Args:
|
||||
scores: Mapping from document id to raw score.
|
||||
|
||||
Returns:
|
||||
Mapping from document id to normalized score. If every score is equal
|
||||
(or there is a single document), all documents receive 1.0.
|
||||
"""
|
||||
if not scores:
|
||||
return {}
|
||||
|
||||
values = list(scores.values())
|
||||
lo, hi = min(values), max(values)
|
||||
span = hi - lo
|
||||
|
||||
if span <= 0:
|
||||
# Degenerate case: all scores identical -> treat as equally relevant.
|
||||
return {doc_id: 1.0 for doc_id in scores}
|
||||
|
||||
return {doc_id: (score - lo) / span for doc_id, score in scores.items()}
|
||||
|
||||
|
||||
def reciprocal_rank_fusion(
|
||||
ranked_lists: Dict[str, RankedList],
|
||||
k: int = DEFAULT_RRF_K,
|
||||
weights: Optional[Dict[str, float]] = None,
|
||||
) -> List[Tuple[str, float]]:
|
||||
"""Fuse multiple ranked lists with Reciprocal Rank Fusion (RRF).
|
||||
|
||||
Args:
|
||||
ranked_lists: Mapping from source name (e.g. "dense", "sparse") to a
|
||||
list of ``(doc_id, score)`` tuples sorted by score descending. Only
|
||||
the *order* of each list matters; the scores are ignored.
|
||||
k: RRF smoothing constant (default 60).
|
||||
weights: Optional per-source weights. Defaults to 1.0 for every source.
|
||||
|
||||
Returns:
|
||||
Fused list of ``(doc_id, fused_score)`` tuples sorted descending.
|
||||
"""
|
||||
weights = weights or {}
|
||||
fused: Dict[str, float] = {}
|
||||
|
||||
for source, ranked in ranked_lists.items():
|
||||
weight = weights.get(source, 1.0)
|
||||
for rank, (doc_id, _score) in enumerate(ranked, start=1):
|
||||
fused[doc_id] = fused.get(doc_id, 0.0) + weight * (1.0 / (k + rank))
|
||||
|
||||
return sorted(fused.items(), key=lambda kv: kv[1], reverse=True)
|
||||
|
||||
|
||||
def weighted_score_fusion(
|
||||
ranked_lists: Dict[str, RankedList],
|
||||
weights: Optional[Dict[str, float]] = None,
|
||||
) -> List[Tuple[str, float]]:
|
||||
"""Fuse multiple ranked lists with weighted, min-max normalized scores.
|
||||
|
||||
Each source list is min-max normalized to [0, 1] independently, then the
|
||||
normalized scores are combined with a weighted sum. A document missing from
|
||||
a source contributes 0 for that source.
|
||||
|
||||
Args:
|
||||
ranked_lists: Mapping from source name to ``(doc_id, score)`` tuples.
|
||||
weights: Optional per-source weights. Defaults to 1.0 for every source.
|
||||
|
||||
Returns:
|
||||
Fused list of ``(doc_id, fused_score)`` tuples sorted descending.
|
||||
"""
|
||||
weights = weights or {}
|
||||
normalized_by_source = {
|
||||
source: min_max_normalize(_ranked_to_score_map(ranked))
|
||||
for source, ranked in ranked_lists.items()
|
||||
}
|
||||
|
||||
fused: Dict[str, float] = {}
|
||||
for source, normalized in normalized_by_source.items():
|
||||
weight = weights.get(source, 1.0)
|
||||
for doc_id, norm_score in normalized.items():
|
||||
fused[doc_id] = fused.get(doc_id, 0.0) + weight * norm_score
|
||||
|
||||
return sorted(fused.items(), key=lambda kv: kv[1], reverse=True)
|
||||
|
||||
|
||||
def fuse(
|
||||
ranked_lists: Dict[str, RankedList],
|
||||
method: str = "rrf",
|
||||
k: int = DEFAULT_RRF_K,
|
||||
weights: Optional[Dict[str, float]] = None,
|
||||
) -> List[Tuple[str, float]]:
|
||||
"""Dispatch helper: fuse ranked lists with the named method.
|
||||
|
||||
Args:
|
||||
ranked_lists: Mapping from source name to ``(doc_id, score)`` tuples.
|
||||
method: "rrf" for Reciprocal Rank Fusion, "weighted" for weighted
|
||||
min-max normalized score fusion.
|
||||
k: RRF smoothing constant (only used when method="rrf").
|
||||
weights: Optional per-source weights.
|
||||
|
||||
Returns:
|
||||
Fused list of ``(doc_id, fused_score)`` tuples sorted descending.
|
||||
|
||||
Raises:
|
||||
ValueError: If ``method`` is not recognized.
|
||||
"""
|
||||
if method == "rrf":
|
||||
return reciprocal_rank_fusion(ranked_lists, k=k, weights=weights)
|
||||
if method == "weighted":
|
||||
return weighted_score_fusion(ranked_lists, weights=weights)
|
||||
raise ValueError(f"Unknown fusion method: {method!r} (expected 'rrf' or 'weighted')")
|
||||
@@ -0,0 +1,289 @@
|
||||
"""FastAPI server for the retrieval pipeline."""
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from typing import Dict, Any, Optional, List
|
||||
from contextlib import asynccontextmanager
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
import uvicorn
|
||||
import asyncio
|
||||
|
||||
from config import PipelineConfig, SearchMode
|
||||
from retrieval_pipeline import RetrievalPipeline
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
|
||||
handlers=[
|
||||
logging.StreamHandler(sys.stdout)
|
||||
]
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize pipeline
|
||||
config = PipelineConfig()
|
||||
pipeline: Optional[RetrievalPipeline] = None
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Lifespan context manager for startup and shutdown events."""
|
||||
global pipeline
|
||||
# Startup
|
||||
try:
|
||||
logger.info("Starting retrieval pipeline...")
|
||||
pipeline = RetrievalPipeline(config)
|
||||
logger.info("Pipeline initialized successfully")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to initialize pipeline: {e}")
|
||||
raise
|
||||
|
||||
yield
|
||||
|
||||
# Shutdown (cleanup if needed)
|
||||
logger.info("Shutting down retrieval pipeline...")
|
||||
|
||||
# Create FastAPI app with lifespan
|
||||
app = FastAPI(
|
||||
title="Hybrid Retrieval Pipeline",
|
||||
description="Educational retrieval pipeline combining dense embeddings, sparse search, and neural reranking",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
# Add CORS middleware
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Pydantic models
|
||||
class IndexRequest(BaseModel):
|
||||
"""Request model for document indexing."""
|
||||
text: str = Field(..., description="Document text to index")
|
||||
doc_id: Optional[str] = Field(None, description="Optional document ID")
|
||||
metadata: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Optional metadata")
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
"""Request model for search."""
|
||||
query: str = Field(..., description="Search query")
|
||||
mode: SearchMode = Field(SearchMode.HYBRID, description="Search mode: dense, sparse, or hybrid")
|
||||
top_k: int = Field(20, ge=1, le=100, description="Number of candidates to retrieve")
|
||||
rerank_top_k: int = Field(10, ge=1, le=50, description="Number of results after reranking")
|
||||
skip_reranking: bool = Field(False, description="Skip reranking for comparison")
|
||||
|
||||
class DeleteRequest(BaseModel):
|
||||
"""Request model for document deletion."""
|
||||
doc_id: str = Field(..., description="Document ID to delete")
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
"""Root endpoint with service information."""
|
||||
return {
|
||||
"service": "Hybrid Retrieval Pipeline",
|
||||
"status": "running" if pipeline else "not initialized",
|
||||
"endpoints": {
|
||||
"index": "/index",
|
||||
"search": "/search",
|
||||
"delete": "/delete",
|
||||
"stats": "/stats",
|
||||
"health": "/health"
|
||||
},
|
||||
"modes": ["dense", "sparse", "hybrid"],
|
||||
"features": [
|
||||
"Parallel dense and sparse indexing",
|
||||
"Hybrid search with both embedding types",
|
||||
"Neural reranking with BGE-Reranker-v2",
|
||||
"Educational score visualization",
|
||||
"Rank change analysis"
|
||||
]
|
||||
}
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
"""Health check endpoint."""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
return {"status": "healthy"}
|
||||
|
||||
@app.post("/index")
|
||||
async def index_document(request: IndexRequest):
|
||||
"""Index a document in both dense and sparse services.
|
||||
|
||||
This endpoint:
|
||||
1. Stores the document locally
|
||||
2. Indexes in dense embedding service (semantic search)
|
||||
3. Indexes in sparse service (BM25 keyword search)
|
||||
4. Returns status from both services
|
||||
"""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
|
||||
try:
|
||||
logger.info(f"Indexing document: {request.doc_id or 'auto-generated'}")
|
||||
result = await pipeline.index_document(
|
||||
text=request.text,
|
||||
doc_id=request.doc_id,
|
||||
metadata=request.metadata
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Indexing failed: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/search")
|
||||
async def search_documents(request: SearchRequest):
|
||||
"""Search for documents using specified mode with optional reranking.
|
||||
|
||||
Educational features:
|
||||
- Shows original rankings from dense and sparse search
|
||||
- Displays similarity scores from each method
|
||||
- Shows final reranked results with score changes
|
||||
- Provides statistics on rank changes and overlap
|
||||
|
||||
Modes:
|
||||
- dense: Semantic search using dense embeddings
|
||||
- sparse: Keyword search using BM25
|
||||
- hybrid: Both methods combined with reranking
|
||||
"""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
|
||||
try:
|
||||
logger.info(f"Search request: mode={request.mode}, query='{request.query[:50]}...'")
|
||||
result = await pipeline.search(
|
||||
query=request.query,
|
||||
mode=request.mode,
|
||||
top_k=request.top_k,
|
||||
rerank_top_k=request.rerank_top_k,
|
||||
skip_reranking=request.skip_reranking
|
||||
)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Search failed: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.delete("/delete")
|
||||
async def delete_document(request: DeleteRequest):
|
||||
"""Delete a document from all services."""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
|
||||
try:
|
||||
logger.info(f"Deleting document: {request.doc_id}")
|
||||
result = await pipeline.delete_document(request.doc_id)
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Deletion failed: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get("/stats")
|
||||
async def get_statistics():
|
||||
"""Get pipeline statistics and configuration."""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
|
||||
try:
|
||||
stats = pipeline.get_statistics()
|
||||
return stats
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get statistics: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get("/documents")
|
||||
async def list_documents(limit: int = 100, offset: int = 0):
|
||||
"""List indexed documents."""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
|
||||
try:
|
||||
docs = pipeline.document_store.list_documents(limit=limit, offset=offset)
|
||||
return {
|
||||
"documents": docs,
|
||||
"total": pipeline.document_store.size(),
|
||||
"limit": limit,
|
||||
"offset": offset
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to list documents: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get("/documents/{doc_id}")
|
||||
async def get_document(doc_id: str):
|
||||
"""Get a specific document by ID."""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
|
||||
try:
|
||||
doc = pipeline.document_store.get_document(doc_id)
|
||||
if doc is None:
|
||||
raise HTTPException(status_code=404, detail=f"Document {doc_id} not found")
|
||||
|
||||
# Return in the format expected by agentic RAG
|
||||
return {
|
||||
"doc_id": doc_id,
|
||||
"content": doc.get("text", ""),
|
||||
"metadata": doc.get("metadata", {})
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get document {doc_id}: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.delete("/clear")
|
||||
async def clear_all():
|
||||
"""Clear all documents from the pipeline (for testing)."""
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=503, detail="Pipeline not initialized")
|
||||
|
||||
try:
|
||||
pipeline.document_store.clear()
|
||||
# Note: This only clears local store, not remote services
|
||||
return {
|
||||
"success": True,
|
||||
"message": "Local document store cleared. Note: Remote services not cleared."
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to clear documents: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
def main():
|
||||
"""Run the server."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Retrieval Pipeline Server")
|
||||
parser.add_argument("--host", default="0.0.0.0", help="Host to bind")
|
||||
parser.add_argument("--port", type=int, default=4242, help="Port to bind (default: 4242)")
|
||||
parser.add_argument("--debug", action="store_true", help="Enable debug mode")
|
||||
parser.add_argument("--dense-url", default="http://localhost:4240", help="Dense service URL")
|
||||
parser.add_argument("--sparse-url", default="http://localhost:4241", help="Sparse service URL")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Update config
|
||||
config.services.dense_service_url = args.dense_url
|
||||
config.services.sparse_service_url = args.sparse_url
|
||||
config.debug = args.debug
|
||||
|
||||
if args.debug:
|
||||
logging.getLogger().setLevel(logging.DEBUG)
|
||||
|
||||
logger.info(f"Starting server on {args.host}:{args.port}")
|
||||
logger.info(f"Dense service: {config.services.dense_service_url}")
|
||||
logger.info(f"Sparse service: {config.services.sparse_service_url}")
|
||||
|
||||
uvicorn.run(
|
||||
app,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
log_level="debug" if args.debug else "info"
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,17 @@
|
||||
fastapi==0.109.0
|
||||
uvicorn==0.27.0
|
||||
httpx==0.26.0
|
||||
pydantic==2.5.3
|
||||
transformers==4.36.2
|
||||
torch==2.1.2
|
||||
sentence-transformers==2.2.2
|
||||
numpy==1.26.3
|
||||
aiohttp==3.9.3
|
||||
typing-extensions==4.9.0
|
||||
python-multipart==0.0.6
|
||||
FlagEmbedding==1.2.10
|
||||
huggingface-hub==0.20.3
|
||||
tqdm==4.66.1
|
||||
# Used by the offline evaluation CLI (evaluate.py)
|
||||
rank-bm25==0.2.2
|
||||
jieba==0.42.1
|
||||
@@ -0,0 +1,267 @@
|
||||
"""Reranker module using BGE-Reranker-v2 model."""
|
||||
|
||||
import torch
|
||||
from typing import List, Tuple, Dict, Any, Optional
|
||||
from dataclasses import dataclass
|
||||
from FlagEmbedding import FlagReranker
|
||||
import logging
|
||||
import time
|
||||
import numpy as np
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from huggingface_hub import snapshot_download
|
||||
from tqdm import tqdm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class RerankResult:
|
||||
"""Result from reranking."""
|
||||
doc_id: str
|
||||
rerank_score: float
|
||||
original_dense_score: Optional[float] = None
|
||||
original_sparse_score: Optional[float] = None
|
||||
original_dense_rank: Optional[int] = None
|
||||
original_sparse_rank: Optional[int] = None
|
||||
text: Optional[str] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
debug_info: Optional[Dict[str, Any]] = None
|
||||
|
||||
class Reranker:
|
||||
"""Reranker using BGE-Reranker-v2 model."""
|
||||
|
||||
def _ensure_model_downloaded(self, model_name: str):
|
||||
"""Check if model is cached and download if needed with progress.
|
||||
|
||||
Args:
|
||||
model_name: HuggingFace model name
|
||||
"""
|
||||
# Check cache directory
|
||||
cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
|
||||
model_id = model_name.replace("/", "--")
|
||||
model_cache_path = cache_dir / f"models--{model_id}"
|
||||
|
||||
if model_cache_path.exists() and any(model_cache_path.iterdir()):
|
||||
logger.info(f"Model already cached at {model_cache_path}")
|
||||
return
|
||||
|
||||
logger.info(f"Model not found in cache. Downloading {model_name}...")
|
||||
logger.info("This is a one-time download. The model will be cached for future use.")
|
||||
|
||||
try:
|
||||
# Use huggingface_hub to download with progress
|
||||
class DownloadProgressBar:
|
||||
def __init__(self):
|
||||
self.pbar = None
|
||||
self.total_size = 0
|
||||
self.downloaded = 0
|
||||
|
||||
def __call__(self, chunk_size: int):
|
||||
if self.pbar is None:
|
||||
return
|
||||
self.downloaded += chunk_size
|
||||
self.pbar.update(chunk_size)
|
||||
|
||||
# Download the model with progress tracking
|
||||
logger.info("Downloading model files...")
|
||||
snapshot_download(
|
||||
repo_id=model_name,
|
||||
cache_dir=cache_dir,
|
||||
resume_download=True,
|
||||
local_files_only=False
|
||||
)
|
||||
logger.info("Model download completed!")
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not pre-download model: {e}")
|
||||
logger.info("Model will be downloaded automatically during initialization...")
|
||||
|
||||
def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3",
|
||||
device: str = None,
|
||||
use_fp16: bool = True,
|
||||
max_length: int = 512):
|
||||
"""Initialize the reranker.
|
||||
|
||||
Args:
|
||||
model_name: HuggingFace model name
|
||||
device: Device to use (mps for Mac, cuda for GPU, cpu)
|
||||
use_fp16: Use half precision for faster inference
|
||||
max_length: Maximum sequence length
|
||||
"""
|
||||
self.model_name = model_name
|
||||
|
||||
# Auto-detect device if not specified
|
||||
if device is None:
|
||||
if torch.backends.mps.is_available():
|
||||
device = "mps"
|
||||
elif torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
self.device = device
|
||||
self.use_fp16 = use_fp16 and device != "cpu"
|
||||
self.max_length = max_length
|
||||
|
||||
logger.info(f"Initializing reranker with model: {model_name}")
|
||||
logger.info(f"Device: {device}, FP16: {self.use_fp16}")
|
||||
|
||||
# Check if model needs to be downloaded
|
||||
self._ensure_model_downloaded(model_name)
|
||||
|
||||
# Initialize the model
|
||||
logger.info("Loading reranker model into memory...")
|
||||
start_time = time.time()
|
||||
self.model = FlagReranker(
|
||||
model_name,
|
||||
use_fp16=self.use_fp16,
|
||||
device=device
|
||||
)
|
||||
elapsed = time.time() - start_time
|
||||
logger.info(f"Reranker initialized successfully in {elapsed:.2f}s")
|
||||
|
||||
def rerank(self,
|
||||
query: str,
|
||||
documents: List[Dict[str, Any]],
|
||||
top_k: int = 10,
|
||||
return_scores: bool = True) -> List[RerankResult]:
|
||||
"""Rerank documents for a query.
|
||||
|
||||
Args:
|
||||
query: The search query
|
||||
documents: List of documents with text and metadata
|
||||
top_k: Number of top results to return
|
||||
return_scores: Whether to return all scores for educational purposes
|
||||
|
||||
Returns:
|
||||
List of reranked results
|
||||
"""
|
||||
if not documents:
|
||||
return []
|
||||
|
||||
start_time = time.time()
|
||||
logger.info(f"Reranking {len(documents)} documents for query: '{query[:50]}...'")
|
||||
|
||||
# Prepare texts for reranking
|
||||
texts = []
|
||||
doc_info = []
|
||||
|
||||
for doc in documents:
|
||||
text = doc.get("text", "")
|
||||
if not text:
|
||||
continue
|
||||
|
||||
texts.append(text)
|
||||
doc_info.append({
|
||||
"doc_id": doc.get("doc_id"),
|
||||
"original_dense_score": doc.get("dense_score"),
|
||||
"original_sparse_score": doc.get("sparse_score"),
|
||||
"original_dense_rank": doc.get("dense_rank"),
|
||||
"original_sparse_rank": doc.get("sparse_rank"),
|
||||
"text": text,
|
||||
"metadata": doc.get("metadata", {})
|
||||
})
|
||||
|
||||
if not texts:
|
||||
logger.warning("No valid texts to rerank")
|
||||
return []
|
||||
|
||||
# Create query-document pairs
|
||||
pairs = [[query, text] for text in texts]
|
||||
|
||||
# Get reranking scores
|
||||
try:
|
||||
scores = self.model.compute_score(pairs, max_length=self.max_length)
|
||||
|
||||
# Convert to numpy array if needed
|
||||
if not isinstance(scores, np.ndarray):
|
||||
scores = np.array(scores)
|
||||
|
||||
# Ensure scores is 1D. FlagReranker.compute_score returns a bare
|
||||
# float when exactly one pair is scored, which becomes a 0-d array
|
||||
# here — atleast_1d keeps the single-candidate case iterable.
|
||||
scores = np.atleast_1d(np.asarray(scores).squeeze())
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Reranking failed: {e}")
|
||||
return []
|
||||
|
||||
# Create results with scores
|
||||
results = []
|
||||
for i, score in enumerate(scores):
|
||||
info = doc_info[i]
|
||||
|
||||
result = RerankResult(
|
||||
doc_id=info["doc_id"],
|
||||
rerank_score=float(score),
|
||||
original_dense_score=info["original_dense_score"],
|
||||
original_sparse_score=info["original_sparse_score"],
|
||||
original_dense_rank=info["original_dense_rank"],
|
||||
original_sparse_rank=info["original_sparse_rank"],
|
||||
text=info["text"] if return_scores else None,
|
||||
metadata=info["metadata"],
|
||||
debug_info={
|
||||
"rerank_model": self.model_name,
|
||||
"max_length": self.max_length,
|
||||
"device": self.device
|
||||
}
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
# Sort by rerank score (descending)
|
||||
results.sort(key=lambda x: x.rerank_score, reverse=True)
|
||||
|
||||
# Add final ranks
|
||||
for i, result in enumerate(results):
|
||||
if result.debug_info:
|
||||
result.debug_info["final_rank"] = i + 1
|
||||
|
||||
elapsed_time = time.time() - start_time
|
||||
logger.info(f"Reranking completed in {elapsed_time:.2f}s")
|
||||
|
||||
# Log score distribution for educational purposes
|
||||
if return_scores and results:
|
||||
scores_array = [r.rerank_score for r in results]
|
||||
logger.info(f"Rerank score distribution: min={min(scores_array):.3f}, "
|
||||
f"max={max(scores_array):.3f}, mean={np.mean(scores_array):.3f}")
|
||||
|
||||
# Log rank changes for top results
|
||||
for i, result in enumerate(results[:5]):
|
||||
changes = []
|
||||
if result.original_dense_rank:
|
||||
dense_change = result.original_dense_rank - (i + 1)
|
||||
changes.append(f"dense: {result.original_dense_rank}→{i+1} ({dense_change:+d})")
|
||||
if result.original_sparse_rank:
|
||||
sparse_change = result.original_sparse_rank - (i + 1)
|
||||
changes.append(f"sparse: {result.original_sparse_rank}→{i+1} ({sparse_change:+d})")
|
||||
|
||||
if changes:
|
||||
logger.debug(f"Doc {result.doc_id} rank changes: {', '.join(changes)}")
|
||||
|
||||
# Return top_k results
|
||||
return results[:top_k]
|
||||
|
||||
def batch_rerank(self,
|
||||
queries: List[str],
|
||||
documents_list: List[List[Dict[str, Any]]],
|
||||
top_k: int = 10,
|
||||
batch_size: int = 32) -> List[List[RerankResult]]:
|
||||
"""Rerank multiple queries in batch.
|
||||
|
||||
Args:
|
||||
queries: List of queries
|
||||
documents_list: List of document lists (one per query)
|
||||
top_k: Number of top results per query
|
||||
batch_size: Batch size for processing
|
||||
|
||||
Returns:
|
||||
List of reranked results for each query
|
||||
"""
|
||||
all_results = []
|
||||
|
||||
for query, documents in zip(queries, documents_list):
|
||||
results = self.rerank(query, documents, top_k)
|
||||
all_results.append(results)
|
||||
|
||||
return all_results
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Script to restart all services for the retrieval pipeline
|
||||
# Run this from the retrieval-pipeline directory
|
||||
|
||||
echo "Restarting all retrieval pipeline services..."
|
||||
echo "============================================"
|
||||
|
||||
# Kill existing services if running
|
||||
echo "Stopping existing services..."
|
||||
pkill -f "python.*server.py" 2>/dev/null
|
||||
pkill -f "python.*main.py" 2>/dev/null
|
||||
# Kill old ports if any still running
|
||||
pkill -f "uvicorn.*8000" 2>/dev/null
|
||||
pkill -f "uvicorn.*8001" 2>/dev/null
|
||||
pkill -f "uvicorn.*4242" 2>/dev/null
|
||||
pkill -f "uvicorn.*8003" 2>/dev/null
|
||||
# Kill new ports
|
||||
pkill -f "uvicorn.*4240" 2>/dev/null
|
||||
pkill -f "uvicorn.*4241" 2>/dev/null
|
||||
pkill -f "uvicorn.*4242" 2>/dev/null
|
||||
|
||||
sleep 2
|
||||
|
||||
# Start dense embedding service
|
||||
echo ""
|
||||
echo "Starting Dense Embedding Service (port 4240)..."
|
||||
cd ../dense-embedding
|
||||
python main.py --port 4240 > dense.log 2>&1 &
|
||||
DENSE_PID=$!
|
||||
echo "Dense service started with PID: $DENSE_PID"
|
||||
|
||||
sleep 3
|
||||
|
||||
# Start sparse embedding service
|
||||
echo ""
|
||||
echo "Starting Sparse Embedding Service (port 4241)..."
|
||||
cd ../sparse-embedding
|
||||
python server.py 4241 > sparse.log 2>&1 &
|
||||
SPARSE_PID=$!
|
||||
echo "Sparse service started with PID: $SPARSE_PID"
|
||||
|
||||
sleep 3
|
||||
|
||||
# Start retrieval pipeline service
|
||||
echo ""
|
||||
echo "Starting Retrieval Pipeline Service (port 4242)..."
|
||||
cd ../retrieval-pipeline
|
||||
python main.py --port 4242 > pipeline.log 2>&1 &
|
||||
PIPELINE_PID=$!
|
||||
echo "Pipeline service started with PID: $PIPELINE_PID"
|
||||
|
||||
sleep 5
|
||||
|
||||
echo ""
|
||||
echo "All services started!"
|
||||
echo "====================="
|
||||
echo "Dense service: http://localhost:4240 (PID: $DENSE_PID)"
|
||||
echo "Sparse service: http://localhost:4241 (PID: $SPARSE_PID)"
|
||||
echo "Pipeline service: http://localhost:4242 (PID: $PIPELINE_PID)"
|
||||
echo ""
|
||||
echo "Logs are being written to:"
|
||||
echo " - dense.log (in dense-embedding/)"
|
||||
echo " - sparse.log (in sparse-embedding/)"
|
||||
echo " - pipeline.log (in retrieval-pipeline/)"
|
||||
echo ""
|
||||
echo "To test the pipeline, run: python test_pipeline.py"
|
||||
echo "To stop all services, run: pkill -f 'python.*server.py|python.*main.py'"
|
||||
@@ -0,0 +1,201 @@
|
||||
"""Client for communicating with dense and sparse embedding services."""
|
||||
|
||||
import httpx
|
||||
import asyncio
|
||||
from typing import Dict, Any, List, Optional, Tuple
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""Unified search result from embedding services."""
|
||||
doc_id: str
|
||||
score: float
|
||||
text: Optional[str] = None
|
||||
metadata: Optional[Dict[str, Any]] = None
|
||||
source: str = "" # "dense" or "sparse"
|
||||
rank: Optional[int] = None
|
||||
debug_info: Optional[Dict[str, Any]] = None
|
||||
|
||||
class RetrievalClient:
|
||||
"""Client for parallel retrieval from dense and sparse services."""
|
||||
|
||||
def __init__(self, dense_url: str, sparse_url: str, timeout: float = 30.0):
|
||||
self.dense_url = dense_url.rstrip('/')
|
||||
self.sparse_url = sparse_url.rstrip('/')
|
||||
self.timeout = timeout
|
||||
|
||||
async def index_document_dense(self, text: str, doc_id: str, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Index a document in the dense embedding service."""
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
try:
|
||||
payload = {
|
||||
"text": text,
|
||||
"doc_id": doc_id,
|
||||
"metadata": metadata or {}
|
||||
}
|
||||
response = await client.post(f"{self.dense_url}/index", json=payload)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
logger.debug(f"Dense indexing successful for doc {doc_id}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Dense indexing failed for doc {doc_id}: {e}")
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
async def index_document_sparse(self, text: str, doc_id: str, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Index a document in the sparse embedding service."""
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
try:
|
||||
# Sparse service now accepts doc_id directly
|
||||
payload = {
|
||||
"text": text,
|
||||
"doc_id": doc_id, # Pass doc_id directly
|
||||
"metadata": metadata or {}
|
||||
}
|
||||
|
||||
response = await client.post(f"{self.sparse_url}/index", json=payload)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
logger.debug(f"Sparse indexing successful for doc {doc_id}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Sparse indexing failed for doc {doc_id}: {e}")
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
async def index_document(self, text: str, doc_id: str, metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Index a document in both services in parallel."""
|
||||
logger.info(f"Indexing document {doc_id} in parallel...")
|
||||
|
||||
# Run both indexing operations in parallel
|
||||
dense_task = self.index_document_dense(text, doc_id, metadata)
|
||||
sparse_task = self.index_document_sparse(text, doc_id, metadata)
|
||||
|
||||
dense_result, sparse_result = await asyncio.gather(dense_task, sparse_task)
|
||||
|
||||
return {
|
||||
"doc_id": doc_id,
|
||||
"dense": dense_result,
|
||||
"sparse": sparse_result,
|
||||
"success": dense_result.get("success", False) and sparse_result.get("success", False)
|
||||
}
|
||||
|
||||
async def search_dense(self, query: str, top_k: int = 20) -> List[SearchResult]:
|
||||
"""Search using dense embeddings."""
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
try:
|
||||
payload = {
|
||||
"query": query,
|
||||
"top_k": top_k,
|
||||
"return_documents": True
|
||||
}
|
||||
response = await client.post(f"{self.dense_url}/search", json=payload)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
results = []
|
||||
for item in data.get("results", []):
|
||||
results.append(SearchResult(
|
||||
doc_id=item["doc_id"],
|
||||
score=item["score"],
|
||||
text=item.get("text"),
|
||||
metadata=item.get("metadata"),
|
||||
source="dense",
|
||||
rank=item.get("rank"),
|
||||
debug_info={
|
||||
"original_score": item["score"],
|
||||
"original_rank": item.get("rank", 0)
|
||||
}
|
||||
))
|
||||
|
||||
logger.debug(f"Dense search returned {len(results)} results")
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Dense search failed: {e}")
|
||||
return []
|
||||
|
||||
async def search_sparse(self, query: str, top_k: int = 20) -> List[SearchResult]:
|
||||
"""Search using sparse embeddings (BM25)."""
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
try:
|
||||
payload = {
|
||||
"query": query,
|
||||
"top_k": top_k
|
||||
}
|
||||
response = await client.post(f"{self.sparse_url}/search", json=payload)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
results = []
|
||||
for idx, item in enumerate(data):
|
||||
# Now doc_id is returned directly from the sparse service
|
||||
doc_id = item.get("doc_id", f"doc_{idx}")
|
||||
|
||||
results.append(SearchResult(
|
||||
doc_id=doc_id,
|
||||
score=item["score"],
|
||||
text=item.get("text"),
|
||||
metadata=item.get("metadata"),
|
||||
source="sparse",
|
||||
rank=idx + 1,
|
||||
debug_info={
|
||||
"bm25_score": item["score"],
|
||||
"matched_terms": item.get("debug", {}).get("matched_terms", []) if item.get("debug") else [],
|
||||
"doc_length": item.get("debug", {}).get("doc_length", 0) if item.get("debug") else 0,
|
||||
"original_rank": idx + 1
|
||||
}
|
||||
))
|
||||
|
||||
logger.debug(f"Sparse search returned {len(results)} results")
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Sparse search failed: {e}")
|
||||
return []
|
||||
|
||||
async def search(self, query: str, top_k: int = 20, mode: str = "hybrid") -> Tuple[List[SearchResult], List[SearchResult]]:
|
||||
"""Search using specified mode (dense, sparse, or hybrid)."""
|
||||
logger.info(f"Searching with mode: {mode}, query: '{query[:50]}...'")
|
||||
|
||||
dense_results = []
|
||||
sparse_results = []
|
||||
|
||||
if mode == "dense":
|
||||
dense_results = await self.search_dense(query, top_k)
|
||||
elif mode == "sparse":
|
||||
sparse_results = await self.search_sparse(query, top_k)
|
||||
elif mode == "hybrid":
|
||||
# Run both searches in parallel
|
||||
dense_task = self.search_dense(query, top_k)
|
||||
sparse_task = self.search_sparse(query, top_k)
|
||||
dense_results, sparse_results = await asyncio.gather(dense_task, sparse_task)
|
||||
else:
|
||||
raise ValueError(f"Invalid search mode: {mode}")
|
||||
|
||||
return dense_results, sparse_results
|
||||
|
||||
async def delete_document(self, doc_id: str) -> Dict[str, Any]:
|
||||
"""Delete a document from both services."""
|
||||
logger.info(f"Deleting document {doc_id} from both services...")
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
# Delete from dense service
|
||||
dense_task = client.delete(f"{self.dense_url}/index", json={"doc_id": doc_id})
|
||||
|
||||
# For sparse service, we need to check if it supports deletion
|
||||
# If not, we'll need to rebuild the index
|
||||
sparse_task = client.delete(f"{self.sparse_url}/index") # Clear all for now
|
||||
|
||||
try:
|
||||
dense_response, sparse_response = await asyncio.gather(dense_task, sparse_task)
|
||||
return {
|
||||
"doc_id": doc_id,
|
||||
"dense": dense_response.json() if dense_response.status_code == 200 else {"success": False},
|
||||
"sparse": sparse_response.json() if sparse_response.status_code == 200 else {"success": False}
|
||||
}
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete document {doc_id}: {e}")
|
||||
return {"success": False, "error": str(e)}
|
||||
@@ -0,0 +1,421 @@
|
||||
"""Main retrieval pipeline combining dense, sparse, and reranking."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any, List, Optional, Tuple
|
||||
from datetime import datetime
|
||||
import uuid
|
||||
|
||||
from config import PipelineConfig, SearchMode
|
||||
from document_store import DocumentStore
|
||||
from retrieval_client import RetrievalClient, SearchResult
|
||||
from reranker import Reranker, RerankResult
|
||||
from fusion import fuse
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class RetrievalPipeline:
|
||||
"""Main retrieval pipeline orchestrating dense, sparse, and reranking."""
|
||||
|
||||
def __init__(self, config: Optional[PipelineConfig] = None):
|
||||
self.config = config or PipelineConfig()
|
||||
|
||||
# Initialize components
|
||||
self.document_store = DocumentStore()
|
||||
self.retrieval_client = RetrievalClient(
|
||||
dense_url=self.config.services.dense_service_url,
|
||||
sparse_url=self.config.services.sparse_service_url
|
||||
)
|
||||
self.reranker = Reranker(
|
||||
model_name=self.config.reranker.model_name,
|
||||
device=self.config.reranker.device,
|
||||
use_fp16=self.config.reranker.use_fp16,
|
||||
max_length=self.config.reranker.max_length
|
||||
)
|
||||
|
||||
logger.info("Retrieval pipeline initialized")
|
||||
|
||||
async def index_document(self,
|
||||
text: str,
|
||||
doc_id: Optional[str] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Index a document in both dense and sparse services.
|
||||
|
||||
Args:
|
||||
text: Document text
|
||||
doc_id: Optional document ID (generated if not provided)
|
||||
metadata: Optional metadata
|
||||
|
||||
Returns:
|
||||
Indexing result with status from both services
|
||||
"""
|
||||
# Generate doc_id if not provided
|
||||
if not doc_id:
|
||||
doc_id = f"doc_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
logger.info(f"Indexing document {doc_id}")
|
||||
|
||||
# Store document locally
|
||||
self.document_store.add_document(doc_id, text, metadata)
|
||||
|
||||
# Index in both services in parallel
|
||||
result = await self.retrieval_client.index_document(text, doc_id, metadata)
|
||||
|
||||
# Add timestamp and document info
|
||||
result["timestamp"] = datetime.now().isoformat()
|
||||
result["text_length"] = len(text)
|
||||
|
||||
if self.config.debug:
|
||||
logger.debug(f"Indexing result: dense={result['dense'].get('success')}, "
|
||||
f"sparse={result['sparse'].get('success')}")
|
||||
|
||||
return result
|
||||
|
||||
async def search(self,
|
||||
query: str,
|
||||
mode: SearchMode = SearchMode.HYBRID,
|
||||
top_k: Optional[int] = None,
|
||||
rerank_top_k: Optional[int] = None,
|
||||
skip_reranking: bool = False) -> Dict[str, Any]:
|
||||
"""Search for documents using specified mode.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
mode: Search mode (dense, sparse, or hybrid)
|
||||
top_k: Number of candidates to retrieve from each service
|
||||
rerank_top_k: Number of results after reranking
|
||||
skip_reranking: Skip reranking step (for comparison)
|
||||
|
||||
Returns:
|
||||
Search results with scores and rankings
|
||||
"""
|
||||
top_k = top_k or self.config.default_top_k
|
||||
rerank_top_k = rerank_top_k or self.config.rerank_top_k
|
||||
|
||||
logger.info(f"Searching with mode={mode}, top_k={top_k}, rerank_top_k={rerank_top_k}")
|
||||
|
||||
start_time = datetime.now()
|
||||
|
||||
# Retrieve from services
|
||||
dense_results, sparse_results = await self.retrieval_client.search(
|
||||
query, top_k, mode.value
|
||||
)
|
||||
|
||||
retrieval_time = (datetime.now() - start_time).total_seconds()
|
||||
|
||||
# Prepare response structure
|
||||
response = {
|
||||
"query": query,
|
||||
"mode": mode.value,
|
||||
"timestamp": start_time.isoformat(),
|
||||
"retrieval_time_ms": retrieval_time * 1000,
|
||||
"dense_results": [],
|
||||
"sparse_results": [],
|
||||
"combined_results": [],
|
||||
"reranked_results": [],
|
||||
"statistics": {}
|
||||
}
|
||||
|
||||
# Process dense results
|
||||
if dense_results:
|
||||
response["dense_results"] = [
|
||||
{
|
||||
"doc_id": r.doc_id,
|
||||
"score": r.score,
|
||||
"rank": r.rank or idx + 1,
|
||||
"text": r.text[:200] if r.text else None
|
||||
}
|
||||
for idx, r in enumerate(dense_results[:10]) # Show top 10 for educational purposes
|
||||
]
|
||||
|
||||
# Process sparse results
|
||||
if sparse_results:
|
||||
response["sparse_results"] = [
|
||||
{
|
||||
"doc_id": r.doc_id,
|
||||
"score": r.score,
|
||||
"rank": r.rank or idx + 1,
|
||||
"text": r.text[:200] if r.text else None,
|
||||
"matched_terms": r.debug_info.get("matched_terms", []) if r.debug_info else []
|
||||
}
|
||||
for idx, r in enumerate(sparse_results[:10]) # Show top 10 for educational purposes
|
||||
]
|
||||
|
||||
# Combine results for reranking
|
||||
combined_docs = self._combine_results(dense_results, sparse_results)
|
||||
|
||||
if not combined_docs:
|
||||
logger.warning("No documents to rerank")
|
||||
return response
|
||||
|
||||
response["combined_results"] = [
|
||||
{
|
||||
"doc_id": doc["doc_id"],
|
||||
"dense_score": doc.get("dense_score"),
|
||||
"sparse_score": doc.get("sparse_score"),
|
||||
"dense_rank": doc.get("dense_rank"),
|
||||
"sparse_rank": doc.get("sparse_rank")
|
||||
}
|
||||
for doc in combined_docs[:20] # Show top 20 combined
|
||||
]
|
||||
|
||||
# Perform reranking if not skipped
|
||||
if not skip_reranking and combined_docs:
|
||||
rerank_start = datetime.now()
|
||||
|
||||
reranked = self.reranker.rerank(
|
||||
query=query,
|
||||
documents=combined_docs,
|
||||
top_k=rerank_top_k,
|
||||
return_scores=self.config.show_scores
|
||||
)
|
||||
|
||||
rerank_time = (datetime.now() - rerank_start).total_seconds()
|
||||
response["rerank_time_ms"] = rerank_time * 1000
|
||||
|
||||
# Format reranked results
|
||||
response["reranked_results"] = self._format_reranked_results(reranked)
|
||||
|
||||
# Calculate statistics
|
||||
response["statistics"] = self._calculate_statistics(
|
||||
dense_results, sparse_results, response.get("reranked_results", [])
|
||||
)
|
||||
|
||||
total_time = (datetime.now() - start_time).total_seconds()
|
||||
response["total_time_ms"] = total_time * 1000
|
||||
|
||||
return response
|
||||
|
||||
def _combine_results(self,
|
||||
dense_results: List[SearchResult],
|
||||
sparse_results: List[SearchResult]) -> List[Dict[str, Any]]:
|
||||
"""Combine dense and sparse results for reranking.
|
||||
|
||||
Args:
|
||||
dense_results: Results from dense search
|
||||
sparse_results: Results from sparse search
|
||||
|
||||
Returns:
|
||||
Combined document list with scores and ranks from both sources
|
||||
"""
|
||||
combined = {}
|
||||
|
||||
# Add dense results
|
||||
for idx, result in enumerate(dense_results):
|
||||
doc_id = result.doc_id
|
||||
if doc_id not in combined:
|
||||
# Get full document from store
|
||||
doc = self.document_store.get_document(doc_id)
|
||||
combined[doc_id] = {
|
||||
"doc_id": doc_id,
|
||||
"text": doc["text"] if doc else result.text,
|
||||
"metadata": doc["metadata"] if doc else result.metadata,
|
||||
"dense_score": result.score,
|
||||
"dense_rank": idx + 1,
|
||||
"sparse_score": None,
|
||||
"sparse_rank": None
|
||||
}
|
||||
else:
|
||||
combined[doc_id]["dense_score"] = result.score
|
||||
combined[doc_id]["dense_rank"] = idx + 1
|
||||
|
||||
# Add sparse results
|
||||
for idx, result in enumerate(sparse_results):
|
||||
doc_id = result.doc_id
|
||||
if doc_id not in combined:
|
||||
# Get full document from store
|
||||
doc = self.document_store.get_document(doc_id)
|
||||
combined[doc_id] = {
|
||||
"doc_id": doc_id,
|
||||
"text": doc["text"] if doc else result.text,
|
||||
"metadata": doc["metadata"] if doc else result.metadata,
|
||||
"dense_score": None,
|
||||
"dense_rank": None,
|
||||
"sparse_score": result.score,
|
||||
"sparse_rank": idx + 1
|
||||
}
|
||||
else:
|
||||
combined[doc_id]["sparse_score"] = result.score
|
||||
combined[doc_id]["sparse_rank"] = idx + 1
|
||||
|
||||
# Convert to list and keep the legacy average-rank field for reference.
|
||||
combined_list = list(combined.values())
|
||||
for doc in combined_list:
|
||||
ranks = []
|
||||
if doc["dense_rank"] is not None:
|
||||
ranks.append(doc["dense_rank"])
|
||||
if doc["sparse_rank"] is not None:
|
||||
ranks.append(doc["sparse_rank"])
|
||||
doc["avg_rank"] = sum(ranks) / len(ranks) if ranks else float('inf')
|
||||
|
||||
# Fuse the two ranked lists into one unified candidate pool. This is the
|
||||
# dedicated fusion stage described in the book (RRF / weighted). The order
|
||||
# produced here is the candidate pool that neural reranking then refines.
|
||||
method = getattr(self.config, "fusion_method", "rrf")
|
||||
if method == "avg_rank":
|
||||
combined_list.sort(key=lambda x: x["avg_rank"])
|
||||
return combined_list
|
||||
|
||||
dense_ranked = [(r.doc_id, r.score) for r in dense_results]
|
||||
sparse_ranked = [(r.doc_id, r.score) for r in sparse_results]
|
||||
fused = fuse(
|
||||
{"dense": dense_ranked, "sparse": sparse_ranked},
|
||||
method=method,
|
||||
k=getattr(self.config, "rrf_k", 60),
|
||||
)
|
||||
fused_order = {doc_id: rank for rank, (doc_id, _) in enumerate(fused)}
|
||||
for doc in combined_list:
|
||||
doc["fusion_score"] = dict(fused).get(doc["doc_id"])
|
||||
# Sort by fused rank; any doc not in the fused output (shouldn't happen)
|
||||
# falls back to the end, then to its average rank for stability.
|
||||
combined_list.sort(
|
||||
key=lambda x: (fused_order.get(x["doc_id"], len(fused)), x["avg_rank"])
|
||||
)
|
||||
return combined_list
|
||||
|
||||
def _format_reranked_results(self, reranked: List[RerankResult]) -> List[Dict[str, Any]]:
|
||||
"""Format reranked results for response.
|
||||
|
||||
Args:
|
||||
reranked: List of reranked results
|
||||
|
||||
Returns:
|
||||
Formatted results with educational information
|
||||
"""
|
||||
formatted = []
|
||||
|
||||
for idx, result in enumerate(reranked):
|
||||
item = {
|
||||
"rank": idx + 1,
|
||||
"doc_id": result.doc_id,
|
||||
"rerank_score": result.rerank_score,
|
||||
"text": result.text,
|
||||
"metadata": result.metadata
|
||||
}
|
||||
|
||||
# Add educational information about rank changes
|
||||
if self.config.show_scores:
|
||||
item["original_scores"] = {
|
||||
"dense": result.original_dense_score,
|
||||
"sparse": result.original_sparse_score
|
||||
}
|
||||
item["original_ranks"] = {
|
||||
"dense": result.original_dense_rank,
|
||||
"sparse": result.original_sparse_rank
|
||||
}
|
||||
|
||||
# Calculate rank changes
|
||||
rank_changes = []
|
||||
if result.original_dense_rank:
|
||||
change = result.original_dense_rank - (idx + 1)
|
||||
rank_changes.append(f"dense: {change:+d}")
|
||||
if result.original_sparse_rank:
|
||||
change = result.original_sparse_rank - (idx + 1)
|
||||
rank_changes.append(f"sparse: {change:+d}")
|
||||
|
||||
item["rank_changes"] = rank_changes
|
||||
|
||||
formatted.append(item)
|
||||
|
||||
return formatted
|
||||
|
||||
def _calculate_statistics(self,
|
||||
dense_results: List[SearchResult],
|
||||
sparse_results: List[SearchResult],
|
||||
reranked_results: List[Dict[str, Any]]) -> Dict[str, Any]:
|
||||
"""Calculate statistics for educational purposes.
|
||||
|
||||
Args:
|
||||
dense_results: Dense search results
|
||||
sparse_results: Sparse search results
|
||||
reranked_results: Reranked results
|
||||
|
||||
Returns:
|
||||
Statistics dictionary
|
||||
"""
|
||||
stats = {
|
||||
"dense_retrieved": len(dense_results),
|
||||
"sparse_retrieved": len(sparse_results),
|
||||
"total_unique_documents": 0,
|
||||
"reranked_count": len(reranked_results)
|
||||
}
|
||||
|
||||
# Count unique documents
|
||||
unique_docs = set()
|
||||
for r in dense_results:
|
||||
unique_docs.add(r.doc_id)
|
||||
for r in sparse_results:
|
||||
unique_docs.add(r.doc_id)
|
||||
|
||||
stats["total_unique_documents"] = len(unique_docs)
|
||||
|
||||
# Calculate overlap
|
||||
if dense_results and sparse_results:
|
||||
dense_ids = {r.doc_id for r in dense_results}
|
||||
sparse_ids = {r.doc_id for r in sparse_results}
|
||||
overlap = dense_ids & sparse_ids
|
||||
stats["overlap_count"] = len(overlap)
|
||||
stats["overlap_percentage"] = (len(overlap) / len(unique_docs)) * 100 if unique_docs else 0
|
||||
|
||||
# Analyze rank changes if available
|
||||
if reranked_results and self.config.show_scores:
|
||||
avg_dense_change = []
|
||||
avg_sparse_change = []
|
||||
|
||||
for r in reranked_results:
|
||||
if "original_ranks" in r:
|
||||
if r["original_ranks"]["dense"]:
|
||||
avg_dense_change.append(r["original_ranks"]["dense"] - r["rank"])
|
||||
if r["original_ranks"]["sparse"]:
|
||||
avg_sparse_change.append(r["original_ranks"]["sparse"] - r["rank"])
|
||||
|
||||
if avg_dense_change:
|
||||
stats["avg_dense_rank_change"] = sum(avg_dense_change) / len(avg_dense_change)
|
||||
if avg_sparse_change:
|
||||
stats["avg_sparse_rank_change"] = sum(avg_sparse_change) / len(avg_sparse_change)
|
||||
|
||||
return stats
|
||||
|
||||
async def delete_document(self, doc_id: str) -> Dict[str, Any]:
|
||||
"""Delete a document from all services.
|
||||
|
||||
Args:
|
||||
doc_id: Document ID to delete
|
||||
|
||||
Returns:
|
||||
Deletion result
|
||||
"""
|
||||
logger.info(f"Deleting document {doc_id}")
|
||||
|
||||
# Delete from local store
|
||||
local_deleted = self.document_store.delete_document(doc_id)
|
||||
|
||||
# Delete from remote services
|
||||
remote_result = await self.retrieval_client.delete_document(doc_id)
|
||||
|
||||
return {
|
||||
"doc_id": doc_id,
|
||||
"local_deleted": local_deleted,
|
||||
"remote_result": remote_result,
|
||||
"success": local_deleted
|
||||
}
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""Get pipeline statistics.
|
||||
|
||||
Returns:
|
||||
Statistics dictionary
|
||||
"""
|
||||
return {
|
||||
"document_store": self.document_store.get_stats(),
|
||||
"pipeline_config": {
|
||||
"default_top_k": self.config.default_top_k,
|
||||
"rerank_top_k": self.config.rerank_top_k,
|
||||
"reranker_model": self.config.reranker.model_name,
|
||||
"device": self.config.reranker.device
|
||||
},
|
||||
"services": {
|
||||
"dense_url": self.config.services.dense_service_url,
|
||||
"sparse_url": self.config.services.sparse_service_url
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,417 @@
|
||||
"""Stage-level evaluation for the hybrid retrieval pipeline.
|
||||
|
||||
The book's 第3章「混合检索流水线」(experiment 3-6) chains four retrieval
|
||||
stages — dense retrieval, sparse retrieval, fusion of the two, and neural
|
||||
reranking — and judges the whole pipeline with aggregate metrics. What it does
|
||||
not provide is an *automated, stage-level* evaluator that answers two questions
|
||||
the experiment raises but leaves to the reader:
|
||||
|
||||
1. **Contribution**: how much does each stage actually add to final quality?
|
||||
(dense alone vs. dense+sparse vs. dense+sparse+rerank)
|
||||
2. **Diminishing returns**: at what point does an extra stage stop paying for
|
||||
itself?
|
||||
|
||||
``RetrievalStageEvaluator`` takes a query set with ground-truth relevant
|
||||
document IDs and pre-computed ranked results from each pipeline stage, computes
|
||||
precision@k / recall@k / NDCG@k / MRR per stage, measures the marginal
|
||||
improvement of adding each stage over the previous one, flags stages whose
|
||||
marginal gain falls below a configurable threshold, and emits a
|
||||
``StageContributionReport`` with both per-query and aggregate analysis.
|
||||
|
||||
The metric definitions match ``evaluate.py`` (binary relevance): NDCG uses a
|
||||
``1 / log2(rank + 1)`` gain, MRR is ``1 / rank`` of the first relevant hit, and
|
||||
recall is ``|topk ∩ gold| / |gold|``. Precision@k is the standard
|
||||
``|topk ∩ gold| / k``. No network, no model, no GPU — the evaluator is a pure
|
||||
function of the ranked lists and the gold sets it is handed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Sequence
|
||||
|
||||
__all__ = [
|
||||
"StageMetrics",
|
||||
"StageContributionReport",
|
||||
"RetrievalStageEvaluator",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Result-shape helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Keys tried, in priority order, when extracting a document id from a result
|
||||
# dict. The pipeline's ``SearchResult`` / reranker output uses ``doc_id``; some
|
||||
# fixtures use ``id``. Both are accepted so the evaluator is shape-tolerant.
|
||||
_DOC_ID_KEYS = ("doc_id", "id", "_id", "document_id")
|
||||
|
||||
|
||||
def _extract_doc_id(result: dict[str, Any]) -> str:
|
||||
"""Pull the document id out of a single result dict."""
|
||||
for key in _DOC_ID_KEYS:
|
||||
value = result.get(key)
|
||||
if value is not None:
|
||||
return str(value)
|
||||
raise KeyError(
|
||||
"result dict has no document id under any of "
|
||||
f"{_DOC_ID_KEYS!r}; got keys {list(result)!r}"
|
||||
)
|
||||
|
||||
|
||||
def _extract_ranked_ids(results: Sequence[dict[str, Any]]) -> list[str]:
|
||||
"""Turn a list of result dicts into an ordered list of document ids.
|
||||
|
||||
Ordering priority: an explicit ``rank`` field (ascending), then an explicit
|
||||
``score`` field (descending), then the original list order. Deduplicates
|
||||
while preserving the first (best) occurrence — a doc should only be counted
|
||||
once even if a buggy stage emits it twice.
|
||||
"""
|
||||
if not results:
|
||||
return []
|
||||
|
||||
items: list[tuple[str, float]] = []
|
||||
has_rank = any("rank" in r for r in results)
|
||||
has_score = any("score" in r for r in results)
|
||||
|
||||
for idx, result in enumerate(results):
|
||||
doc_id = _extract_doc_id(result)
|
||||
if has_rank:
|
||||
# rank is 1-indexed ascending; missing rank sorts last by index.
|
||||
sort_key = float(result.get("rank", idx))
|
||||
elif has_score:
|
||||
# score is descending; negate so larger score sorts first.
|
||||
sort_key = -float(result.get("score", 0.0))
|
||||
else:
|
||||
# Preserve insertion order with a stable key.
|
||||
sort_key = float(idx)
|
||||
items.append((doc_id, sort_key))
|
||||
|
||||
items.sort(key=lambda kv: kv[1])
|
||||
|
||||
seen: set[str] = set()
|
||||
ranked: list[str] = []
|
||||
for doc_id, _ in items:
|
||||
if doc_id in seen:
|
||||
continue
|
||||
seen.add(doc_id)
|
||||
ranked.append(doc_id)
|
||||
return ranked
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataclasses
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageMetrics:
|
||||
"""Metrics for one stage, aggregated (mean) over the query set.
|
||||
|
||||
``precision_at_k`` / ``recall_at_k`` / ``ndcg_at_k`` map each evaluated k to
|
||||
its aggregate value; ``mrr`` is the mean reciprocal rank over all queries.
|
||||
"""
|
||||
|
||||
stage_name: str
|
||||
precision_at_k: dict[int, float] = field(default_factory=dict)
|
||||
recall_at_k: dict[int, float] = field(default_factory=dict)
|
||||
ndcg_at_k: dict[int, float] = field(default_factory=dict)
|
||||
mrr: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageContributionReport:
|
||||
"""Full stage-contribution analysis for a multi-stage retrieval pipeline."""
|
||||
|
||||
total_queries: int
|
||||
stage_metrics: list[StageMetrics]
|
||||
# stage_name -> k -> improvement of that stage over the previous stage
|
||||
# (the first stage's "improvement" is measured over an empty baseline, so
|
||||
# its values equal its own NDCG@k).
|
||||
marginal_improvement: dict[str, dict[int, float]] = field(default_factory=dict)
|
||||
# Stages (after the first) whose mean marginal NDCG@k gain across all k is
|
||||
# below ``diminishing_threshold``.
|
||||
diminishing_return_stages: list[str] = field(default_factory=list)
|
||||
# Cumulative prefix (e.g. "dense+sparse+rerank") with the highest mean
|
||||
# NDCG@k across k; ties resolve to the shorter (cheaper) combination.
|
||||
best_stage_combination: str = ""
|
||||
# One entry per query with per-stage metrics and per-query marginal gains.
|
||||
per_query_analysis: list[dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Evaluator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class RetrievalStageEvaluator:
|
||||
"""Measure how much each stage of a retrieval pipeline contributes.
|
||||
|
||||
The evaluator is stage-agnostic: it works with any ordered set of stages
|
||||
whose results are supplied as ranked lists of result dicts. The stage order
|
||||
in ``stage_results`` defines the accumulation chain used for marginal
|
||||
improvement and best-combination selection (e.g. for the book's pipeline the
|
||||
caller passes ``{"dense": ..., "sparse": ..., "fusion": ..., "rerank": ...}``
|
||||
and the cumulative combinations become ``dense``, ``dense+sparse``,
|
||||
``dense+sparse+fusion``, ``dense+sparse+fusion+rerank``).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
k_values: list[int] | None = None,
|
||||
diminishing_threshold: float = 0.01,
|
||||
) -> None:
|
||||
self.k_values: list[int] = sorted(set(k_values or [1, 5, 10, 20]))
|
||||
if any(k < 1 for k in self.k_values):
|
||||
raise ValueError("k_values must be positive integers")
|
||||
self.diminishing_threshold = float(diminishing_threshold)
|
||||
|
||||
# -- atomic metric primitives -------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def compute_precision_at_k(
|
||||
ranked_ids: list[str], relevant_ids: set[str], k: int
|
||||
) -> float:
|
||||
"""Precision@k = |top-k ∩ relevant| / k (0.0 when k <= 0)."""
|
||||
if k <= 0:
|
||||
return 0.0
|
||||
topk = set(ranked_ids[:k])
|
||||
return len(topk & relevant_ids) / k
|
||||
|
||||
@staticmethod
|
||||
def compute_recall_at_k(
|
||||
ranked_ids: list[str], relevant_ids: set[str], k: int
|
||||
) -> float:
|
||||
"""Recall@k = |top-k ∩ relevant| / |relevant| (0.0 when no gold)."""
|
||||
if not relevant_ids:
|
||||
return 0.0
|
||||
topk = set(ranked_ids[:k])
|
||||
return len(topk & relevant_ids) / len(relevant_ids)
|
||||
|
||||
@staticmethod
|
||||
def compute_ndcg_at_k(
|
||||
ranked_ids: list[str], relevant_ids: set[str], k: int
|
||||
) -> float:
|
||||
"""NDCG@k with binary relevance (matches ``evaluate.ndcg_at_k``).
|
||||
|
||||
Gain per relevant hit at rank ``i`` (1-indexed) is ``1 / log2(i + 1)``;
|
||||
the ideal DCG places the ``min(|gold|, k)`` relevant docs in the top
|
||||
positions. Returns 0.0 when there is no ideal ranking (no gold).
|
||||
"""
|
||||
if k <= 0 or not relevant_ids:
|
||||
return 0.0
|
||||
dcg = 0.0
|
||||
for idx, doc_id in enumerate(ranked_ids[:k], start=1):
|
||||
if doc_id in relevant_ids:
|
||||
dcg += 1.0 / math.log2(idx + 1)
|
||||
ideal_hits = min(len(relevant_ids), k)
|
||||
idcg = sum(1.0 / math.log2(i + 1) for i in range(1, ideal_hits + 1))
|
||||
return dcg / idcg if idcg > 0 else 0.0
|
||||
|
||||
@staticmethod
|
||||
def compute_mrr(ranked_ids: list[str], relevant_ids: set[str]) -> float:
|
||||
"""Mean reciprocal rank source: 1 / rank of first relevant hit, else 0.0."""
|
||||
for idx, doc_id in enumerate(ranked_ids, start=1):
|
||||
if doc_id in relevant_ids:
|
||||
return 1.0 / idx
|
||||
return 0.0
|
||||
|
||||
# -- per-(query, stage) evaluation --------------------------------------
|
||||
|
||||
def evaluate_stage(
|
||||
self,
|
||||
results: list[dict[str, Any]],
|
||||
relevant_ids: set[str],
|
||||
k_values: list[int],
|
||||
) -> StageMetrics:
|
||||
"""Evaluate one stage's results for a single query.
|
||||
|
||||
``results`` is the ranked result-dict list for one (query, stage) pair;
|
||||
``relevant_ids`` is that query's gold set. Returns a ``StageMetrics``
|
||||
with the per-k metrics and MRR for this single query. The
|
||||
``stage_name`` is left blank here — ``evaluate_pipeline`` sets the real
|
||||
name when it aggregates across queries.
|
||||
"""
|
||||
ranked_ids = _extract_ranked_ids(results)
|
||||
precision: dict[int, float] = {}
|
||||
recall: dict[int, float] = {}
|
||||
ndcg: dict[int, float] = {}
|
||||
for k in k_values:
|
||||
precision[k] = self.compute_precision_at_k(ranked_ids, relevant_ids, k)
|
||||
recall[k] = self.compute_recall_at_k(ranked_ids, relevant_ids, k)
|
||||
ndcg[k] = self.compute_ndcg_at_k(ranked_ids, relevant_ids, k)
|
||||
mrr = self.compute_mrr(ranked_ids, relevant_ids)
|
||||
return StageMetrics(
|
||||
stage_name="",
|
||||
precision_at_k=precision,
|
||||
recall_at_k=recall,
|
||||
ndcg_at_k=ndcg,
|
||||
mrr=mrr,
|
||||
)
|
||||
|
||||
# -- whole-pipeline evaluation ------------------------------------------
|
||||
|
||||
def evaluate_pipeline(
|
||||
self,
|
||||
stage_results: dict[str, list[dict[str, Any]]],
|
||||
ground_truth: dict[str, set[str]],
|
||||
) -> StageContributionReport:
|
||||
"""Evaluate every stage across the query set and build the report.
|
||||
|
||||
``stage_results`` maps stage name -> flat list of result dicts for that
|
||||
stage across *all* queries; each result dict must carry a ``query_id``
|
||||
(so it can be grouped back to its query) and a document id. Insertion
|
||||
order of ``stage_results`` defines the accumulation chain.
|
||||
|
||||
``ground_truth`` maps query id -> set of relevant document ids. The
|
||||
query set is exactly ``ground_truth.keys()``; queries appearing only in
|
||||
stage results are ignored, and queries with no results for a stage score
|
||||
zero for that stage.
|
||||
"""
|
||||
stage_names = list(stage_results.keys())
|
||||
query_ids = list(ground_truth.keys())
|
||||
|
||||
# Group each stage's flat result list by query_id, once.
|
||||
grouped: dict[str, dict[str, list[dict[str, Any]]]] = {
|
||||
stage: {} for stage in stage_names
|
||||
}
|
||||
for stage, results in stage_results.items():
|
||||
for result in results:
|
||||
qid = result.get("query_id")
|
||||
if qid is None:
|
||||
raise KeyError(
|
||||
f"stage {stage!r} result missing 'query_id'; "
|
||||
f"got keys {list(result)!r}"
|
||||
)
|
||||
grouped[stage].setdefault(qid, []).append(result)
|
||||
|
||||
# Per-query, per-stage metrics.
|
||||
per_query_stage: dict[str, dict[str, StageMetrics]] = {
|
||||
qid: {} for qid in query_ids
|
||||
}
|
||||
for stage in stage_names:
|
||||
for qid in query_ids:
|
||||
relevant = ground_truth.get(qid, set())
|
||||
results = grouped[stage].get(qid, [])
|
||||
per_query_stage[qid][stage] = self.evaluate_stage(
|
||||
results, relevant, self.k_values
|
||||
)
|
||||
|
||||
# Aggregate each stage by mean over queries.
|
||||
stage_metrics: list[StageMetrics] = []
|
||||
for stage in stage_names:
|
||||
n = max(len(query_ids), 1)
|
||||
agg_precision: dict[int, float] = {
|
||||
k: sum(per_query_stage[qid][stage].precision_at_k[k] for qid in query_ids) / n
|
||||
for k in self.k_values
|
||||
}
|
||||
agg_recall: dict[int, float] = {
|
||||
k: sum(per_query_stage[qid][stage].recall_at_k[k] for qid in query_ids) / n
|
||||
for k in self.k_values
|
||||
}
|
||||
agg_ndcg: dict[int, float] = {
|
||||
k: sum(per_query_stage[qid][stage].ndcg_at_k[k] for qid in query_ids) / n
|
||||
for k in self.k_values
|
||||
}
|
||||
agg_mrr = sum(per_query_stage[qid][stage].mrr for qid in query_ids) / n
|
||||
stage_metrics.append(
|
||||
StageMetrics(
|
||||
stage_name=stage,
|
||||
precision_at_k=agg_precision,
|
||||
recall_at_k=agg_recall,
|
||||
ndcg_at_k=agg_ndcg,
|
||||
mrr=agg_mrr,
|
||||
)
|
||||
)
|
||||
|
||||
# Marginal improvement over the previous stage (NDCG@k as the
|
||||
# contribution signal). The first stage is measured against an empty
|
||||
# baseline, so its marginal equals its own NDCG@k.
|
||||
marginal_improvement: dict[str, dict[int, float]] = {}
|
||||
for i, stage in enumerate(stage_names):
|
||||
current = stage_metrics[i].ndcg_at_k
|
||||
if i == 0:
|
||||
marginal_improvement[stage] = {k: current[k] for k in self.k_values}
|
||||
else:
|
||||
previous = stage_metrics[i - 1].ndcg_at_k
|
||||
marginal_improvement[stage] = {
|
||||
k: current[k] - previous[k] for k in self.k_values
|
||||
}
|
||||
|
||||
# Diminishing returns: a non-first stage whose mean marginal NDCG@k
|
||||
# gain across all k is below the threshold.
|
||||
diminishing_return_stages: list[str] = []
|
||||
for i, stage in enumerate(stage_names):
|
||||
if i == 0:
|
||||
continue
|
||||
mean_gain = sum(marginal_improvement[stage].values()) / max(
|
||||
len(self.k_values), 1
|
||||
)
|
||||
if mean_gain < self.diminishing_threshold:
|
||||
diminishing_return_stages.append(stage)
|
||||
|
||||
# Best cumulative combination: highest mean NDCG@k across k; ties go to
|
||||
# the shorter (cheaper) prefix.
|
||||
best_label = ""
|
||||
best_score = -1.0
|
||||
for i, stage in enumerate(stage_names):
|
||||
label = "+".join(stage_names[: i + 1])
|
||||
score = sum(stage_metrics[i].ndcg_at_k.values()) / max(
|
||||
len(self.k_values), 1
|
||||
)
|
||||
if score > best_score + 1e-12:
|
||||
best_score = score
|
||||
best_label = label
|
||||
if not stage_names:
|
||||
best_label = ""
|
||||
|
||||
# Per-query analysis.
|
||||
per_query_analysis: list[dict[str, Any]] = []
|
||||
for qid in query_ids:
|
||||
stage_block: dict[str, dict[str, Any]] = {}
|
||||
for stage in stage_names:
|
||||
sm = per_query_stage[qid][stage]
|
||||
stage_block[stage] = {
|
||||
"precision_at_k": dict(sm.precision_at_k),
|
||||
"recall_at_k": dict(sm.recall_at_k),
|
||||
"ndcg_at_k": dict(sm.ndcg_at_k),
|
||||
"mrr": sm.mrr,
|
||||
}
|
||||
# Best stage for this query by mean NDCG@k (ties -> first stage).
|
||||
best_stage = ""
|
||||
best_q_score = -1.0
|
||||
for stage in stage_names:
|
||||
sm = per_query_stage[qid][stage]
|
||||
score = sum(sm.ndcg_at_k.values()) / max(len(self.k_values), 1)
|
||||
if score > best_q_score + 1e-12:
|
||||
best_q_score = score
|
||||
best_stage = stage
|
||||
# Per-query marginal improvement (NDCG@k).
|
||||
q_marginal: dict[str, dict[int, float]] = {}
|
||||
for i, stage in enumerate(stage_names):
|
||||
current = per_query_stage[qid][stage].ndcg_at_k
|
||||
if i == 0:
|
||||
q_marginal[stage] = {k: current[k] for k in self.k_values}
|
||||
else:
|
||||
previous = per_query_stage[qid][stage_names[i - 1]].ndcg_at_k
|
||||
q_marginal[stage] = {
|
||||
k: current[k] - previous[k] for k in self.k_values
|
||||
}
|
||||
per_query_analysis.append(
|
||||
{
|
||||
"query_id": qid,
|
||||
"stage_metrics": stage_block,
|
||||
"best_stage": best_stage,
|
||||
"marginal_improvement": q_marginal,
|
||||
}
|
||||
)
|
||||
|
||||
return StageContributionReport(
|
||||
total_queries=len(query_ids),
|
||||
stage_metrics=stage_metrics,
|
||||
marginal_improvement=marginal_improvement,
|
||||
diminishing_return_stages=diminishing_return_stages,
|
||||
best_stage_combination=best_label,
|
||||
per_query_analysis=per_query_analysis,
|
||||
)
|
||||
+125
@@ -0,0 +1,125 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Start all services for the retrieval pipeline
|
||||
|
||||
echo "========================================="
|
||||
echo "Starting Retrieval Pipeline Services"
|
||||
echo "========================================="
|
||||
|
||||
# Colors for output
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
# Function to check if port is in use
|
||||
check_port() {
|
||||
if lsof -Pi :$1 -sTCP:LISTEN -t >/dev/null ; then
|
||||
return 0
|
||||
else
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Kill existing services on ports
|
||||
echo -e "${YELLOW}Checking for existing services...${NC}"
|
||||
for port in 4240 4241 4242; do
|
||||
if check_port $port; then
|
||||
echo -e "${YELLOW}Killing existing service on port $port${NC}"
|
||||
lsof -ti:$port | xargs kill -9 2>/dev/null
|
||||
sleep 1
|
||||
fi
|
||||
done
|
||||
|
||||
# Start dense embedding service
|
||||
echo -e "\n${GREEN}Starting Dense Embedding Service (port 8000)...${NC}"
|
||||
cd ../dense-embedding
|
||||
python main.py --port 8000 > dense.log 2>&1 &
|
||||
DENSE_PID=$!
|
||||
echo "Dense service PID: $DENSE_PID"
|
||||
|
||||
# Wait for dense service to start
|
||||
echo "Waiting for dense service to initialize..."
|
||||
for i in {1..30}; do
|
||||
if check_port 8000; then
|
||||
echo -e "${GREEN}✓ Dense service ready${NC}"
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
# Start sparse embedding service
|
||||
echo -e "\n${GREEN}Starting Sparse Embedding Service (port 8001)...${NC}"
|
||||
cd ../sparse-embedding
|
||||
python server.py --port 8001 > sparse.log 2>&1 &
|
||||
SPARSE_PID=$!
|
||||
echo "Sparse service PID: $SPARSE_PID"
|
||||
|
||||
# Wait for sparse service to start
|
||||
echo "Waiting for sparse service to initialize..."
|
||||
for i in {1..30}; do
|
||||
if check_port 8001; then
|
||||
echo -e "${GREEN}✓ Sparse service ready${NC}"
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
# Start retrieval pipeline
|
||||
echo -e "\n${GREEN}Starting Retrieval Pipeline (port 4242)...${NC}"
|
||||
cd ../retrieval-pipeline
|
||||
python main.py --port 4242 > pipeline.log 2>&1 &
|
||||
PIPELINE_PID=$!
|
||||
echo "Pipeline service PID: $PIPELINE_PID"
|
||||
|
||||
# Wait for pipeline to start
|
||||
echo "Waiting for pipeline to initialize (loading reranker model)..."
|
||||
for i in {1..60}; do
|
||||
if check_port 4242; then
|
||||
echo -e "${GREEN}✓ Pipeline ready${NC}"
|
||||
break
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
# Check all services
|
||||
echo -e "\n========================================="
|
||||
echo "Service Status:"
|
||||
echo "========================================="
|
||||
|
||||
if check_port 8000; then
|
||||
echo -e "${GREEN}✓ Dense Embedding Service: http://localhost:8000${NC}"
|
||||
else
|
||||
echo -e "${RED}✗ Dense Embedding Service failed to start${NC}"
|
||||
fi
|
||||
|
||||
if check_port 8001; then
|
||||
echo -e "${GREEN}✓ Sparse Embedding Service: http://localhost:8001${NC}"
|
||||
else
|
||||
echo -e "${RED}✗ Sparse Embedding Service failed to start${NC}"
|
||||
fi
|
||||
|
||||
if check_port 4242; then
|
||||
echo -e "${GREEN}✓ Retrieval Pipeline: http://localhost:4242${NC}"
|
||||
echo -e "${GREEN}✓ API Documentation: http://localhost:4242/docs${NC}"
|
||||
else
|
||||
echo -e "${RED}✗ Retrieval Pipeline failed to start${NC}"
|
||||
fi
|
||||
|
||||
echo -e "\n========================================="
|
||||
echo "All services started!"
|
||||
echo "========================================="
|
||||
echo ""
|
||||
echo "To test the pipeline:"
|
||||
echo " python test_client.py"
|
||||
echo ""
|
||||
echo "To run the demo:"
|
||||
echo " python demo.py"
|
||||
echo ""
|
||||
echo "To stop all services:"
|
||||
echo " ./stop_all_services.sh"
|
||||
echo ""
|
||||
echo "Service logs:"
|
||||
echo " Dense: dense.log"
|
||||
echo " Sparse: sparse.log"
|
||||
echo " Pipeline: pipeline.log"
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Stop all retrieval pipeline services
|
||||
|
||||
echo "Stopping all retrieval pipeline services..."
|
||||
|
||||
# Kill processes on ports
|
||||
for port in 4240 4241 4242; do
|
||||
if lsof -Pi :$port -sTCP:LISTEN -t >/dev/null ; then
|
||||
echo "Stopping service on port $port..."
|
||||
lsof -ti:$port | xargs kill -9 2>/dev/null
|
||||
fi
|
||||
done
|
||||
|
||||
echo "All services stopped."
|
||||
@@ -0,0 +1,497 @@
|
||||
"""Test client with educational test cases for dense vs sparse retrieval."""
|
||||
|
||||
import asyncio
|
||||
import httpx
|
||||
import json
|
||||
from typing import List, Dict, Any
|
||||
import logging
|
||||
from datetime import datetime
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class TestClient:
|
||||
"""Test client for the retrieval pipeline."""
|
||||
|
||||
def __init__(self, base_url: str = "http://localhost:4242"):
|
||||
self.base_url = base_url.rstrip('/')
|
||||
self.test_results = []
|
||||
|
||||
async def index_document(self, text: str, doc_id: str = None, metadata: Dict = None) -> Dict:
|
||||
"""Index a document."""
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/index",
|
||||
json={"text": text, "doc_id": doc_id, "metadata": metadata or {}}
|
||||
)
|
||||
return response.json()
|
||||
|
||||
async def search(self, query: str, mode: str = "hybrid", top_k: int = 20, rerank_top_k: int = 10) -> Dict:
|
||||
"""Search for documents."""
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(
|
||||
f"{self.base_url}/search",
|
||||
json={
|
||||
"query": query,
|
||||
"mode": mode,
|
||||
"top_k": top_k,
|
||||
"rerank_top_k": rerank_top_k
|
||||
}
|
||||
)
|
||||
return response.json()
|
||||
|
||||
async def clear_documents(self) -> Dict:
|
||||
"""Clear all documents."""
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.delete(f"{self.base_url}/clear")
|
||||
return response.json()
|
||||
|
||||
def print_results(self, results: Dict, title: str = "Search Results"):
|
||||
"""Pretty print search results."""
|
||||
print(f"\n{'='*80}")
|
||||
print(f"{title}")
|
||||
print(f"{'='*80}")
|
||||
print(f"Query: {results.get('query', 'N/A')}")
|
||||
print(f"Mode: {results.get('mode', 'N/A')}")
|
||||
print(f"Times: Retrieval={results.get('retrieval_time_ms', 0):.1f}ms, "
|
||||
f"Rerank={results.get('rerank_time_ms', 0):.1f}ms, "
|
||||
f"Total={results.get('total_time_ms', 0):.1f}ms")
|
||||
|
||||
# Show top dense results
|
||||
if results.get('dense_results'):
|
||||
print(f"\nTop Dense Results:")
|
||||
for r in results['dense_results'][:5]:
|
||||
print(f" #{r['rank']}: {r['doc_id']} (score: {r['score']:.4f})")
|
||||
|
||||
# Show top sparse results
|
||||
if results.get('sparse_results'):
|
||||
print(f"\nTop Sparse Results:")
|
||||
for r in results['sparse_results'][:5]:
|
||||
matched = r.get('matched_terms', [])
|
||||
print(f" #{r['rank']}: {r['doc_id']} (score: {r['score']:.4f}, matched: {matched})")
|
||||
|
||||
# Show reranked results
|
||||
if results.get('reranked_results'):
|
||||
print(f"\nReranked Results:")
|
||||
for r in results['reranked_results'][:5]:
|
||||
changes = r.get('rank_changes', [])
|
||||
print(f" #{r['rank']}: {r['doc_id']} (score: {r['rerank_score']:.4f})")
|
||||
if changes:
|
||||
print(f" Rank changes: {', '.join(changes)}")
|
||||
|
||||
# Show statistics
|
||||
if results.get('statistics'):
|
||||
stats = results['statistics']
|
||||
print(f"\nStatistics:")
|
||||
print(f" Dense retrieved: {stats.get('dense_retrieved', 0)}")
|
||||
print(f" Sparse retrieved: {stats.get('sparse_retrieved', 0)}")
|
||||
print(f" Overlap: {stats.get('overlap_count', 0)} ({stats.get('overlap_percentage', 0):.1f}%)")
|
||||
|
||||
async def run_test_case(self, name: str, documents: List[Dict], queries: List[Dict]) -> Dict:
|
||||
"""Run a complete test case."""
|
||||
print(f"\n{'='*80}")
|
||||
print(f"TEST CASE: {name}")
|
||||
print(f"{'='*80}")
|
||||
|
||||
test_result = {
|
||||
"name": name,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"documents": len(documents),
|
||||
"queries": len(queries),
|
||||
"results": []
|
||||
}
|
||||
|
||||
# Index documents
|
||||
print(f"\nIndexing {len(documents)} documents...")
|
||||
for doc in documents:
|
||||
result = await self.index_document(
|
||||
text=doc["text"],
|
||||
doc_id=doc.get("doc_id"),
|
||||
metadata=doc.get("metadata", {})
|
||||
)
|
||||
print(f" Indexed: {doc.get('doc_id', 'auto')} - {doc['text'][:50]}...")
|
||||
|
||||
# Run queries
|
||||
print(f"\nRunning {len(queries)} queries...")
|
||||
for query_spec in queries:
|
||||
query = query_spec["query"]
|
||||
expected = query_spec.get("expected", [])
|
||||
explanation = query_spec.get("explanation", "")
|
||||
|
||||
print(f"\nQuery: '{query}'")
|
||||
if explanation:
|
||||
print(f"Explanation: {explanation}")
|
||||
if expected:
|
||||
print(f"Expected top results: {expected}")
|
||||
|
||||
# Test all modes
|
||||
for mode in ["dense", "sparse", "hybrid"]:
|
||||
print(f"\n--- Mode: {mode} ---")
|
||||
result = await self.search(query, mode=mode, top_k=10, rerank_top_k=5)
|
||||
|
||||
# Extract top results
|
||||
top_results = []
|
||||
if mode == "hybrid" and result.get("reranked_results"):
|
||||
top_results = [r["doc_id"] for r in result["reranked_results"][:3]]
|
||||
elif mode == "dense" and result.get("dense_results"):
|
||||
top_results = [r["doc_id"] for r in result["dense_results"][:3]]
|
||||
elif mode == "sparse" and result.get("sparse_results"):
|
||||
top_results = [r["doc_id"] for r in result["sparse_results"][:3]]
|
||||
|
||||
print(f"Top 3: {top_results}")
|
||||
|
||||
# Check if expected results are in top positions
|
||||
if expected:
|
||||
matches = [doc_id in top_results for doc_id in expected]
|
||||
accuracy = sum(matches) / len(expected) * 100
|
||||
print(f"Accuracy: {accuracy:.0f}% ({sum(matches)}/{len(expected)} expected found)")
|
||||
|
||||
test_result["results"].append({
|
||||
"query": query,
|
||||
"mode": mode,
|
||||
"top_results": top_results,
|
||||
"expected": expected,
|
||||
"time_ms": result.get("total_time_ms", 0)
|
||||
})
|
||||
|
||||
self.test_results.append(test_result)
|
||||
return test_result
|
||||
|
||||
# Test cases demonstrating dense vs sparse strengths
|
||||
async def run_educational_tests():
|
||||
"""Run educational test cases."""
|
||||
client = TestClient()
|
||||
|
||||
# Clear existing documents
|
||||
await client.clear_documents()
|
||||
|
||||
# Test Case 1: Semantic Similarity (Dense is better)
|
||||
semantic_docs = [
|
||||
{
|
||||
"doc_id": "cat_1",
|
||||
"text": "The feline jumped onto the couch and purred contentedly.",
|
||||
"metadata": {"category": "animals", "type": "behavior"}
|
||||
},
|
||||
{
|
||||
"doc_id": "cat_2",
|
||||
"text": "A tabby cat sleeps on the windowsill in the afternoon sun.",
|
||||
"metadata": {"category": "animals", "type": "description"}
|
||||
},
|
||||
{
|
||||
"doc_id": "dog_1",
|
||||
"text": "The puppy barked excitedly and wagged its tail.",
|
||||
"metadata": {"category": "animals", "type": "behavior"}
|
||||
},
|
||||
{
|
||||
"doc_id": "car_1",
|
||||
"text": "The vehicle accelerated down the highway.",
|
||||
"metadata": {"category": "transportation", "type": "action"}
|
||||
}
|
||||
]
|
||||
|
||||
semantic_queries = [
|
||||
{
|
||||
"query": "kitty behavior", # Uses different words but same concept
|
||||
"expected": ["cat_1", "cat_2"],
|
||||
"explanation": "Dense should find cat documents despite using 'kitty' instead of 'cat/feline'"
|
||||
},
|
||||
{
|
||||
"query": "automobile speed", # Semantic similarity to car/vehicle
|
||||
"expected": ["car_1"],
|
||||
"explanation": "Dense should match 'automobile' to 'vehicle' and 'speed' to 'accelerated'"
|
||||
}
|
||||
]
|
||||
|
||||
await client.run_test_case(
|
||||
"Semantic Similarity (Dense Advantage)",
|
||||
semantic_docs,
|
||||
semantic_queries
|
||||
)
|
||||
|
||||
# Test Case 2: Exact Terms and Names (Sparse is better)
|
||||
exact_docs = [
|
||||
{
|
||||
"doc_id": "person_1",
|
||||
"text": "Dr. Alexander Humphrey published groundbreaking research on quantum computing.",
|
||||
"metadata": {"type": "person", "field": "science"}
|
||||
},
|
||||
{
|
||||
"doc_id": "person_2",
|
||||
"text": "Professor Smith teaches computer science at the university.",
|
||||
"metadata": {"type": "person", "field": "education"}
|
||||
},
|
||||
{
|
||||
"doc_id": "company_1",
|
||||
"text": "XR-7000 is a new model released by TechCorp Industries.",
|
||||
"metadata": {"type": "product", "company": "TechCorp"}
|
||||
},
|
||||
{
|
||||
"doc_id": "company_2",
|
||||
"text": "The latest smartphone features advanced technology.",
|
||||
"metadata": {"type": "product", "category": "electronics"}
|
||||
}
|
||||
]
|
||||
|
||||
exact_queries = [
|
||||
{
|
||||
"query": "Alexander Humphrey", # Exact name match
|
||||
"expected": ["person_1"],
|
||||
"explanation": "Sparse should excel at finding exact name 'Alexander Humphrey'"
|
||||
},
|
||||
{
|
||||
"query": "XR-7000", # Specific product code
|
||||
"expected": ["company_1"],
|
||||
"explanation": "Sparse should find exact product code 'XR-7000'"
|
||||
}
|
||||
]
|
||||
|
||||
await client.run_test_case(
|
||||
"Exact Terms and Names (Sparse Advantage)",
|
||||
exact_docs,
|
||||
exact_queries
|
||||
)
|
||||
|
||||
# Test Case 3: Multilingual (Dense is better)
|
||||
multilingual_docs = [
|
||||
{
|
||||
"doc_id": "ml_en_1",
|
||||
"text": "Machine learning is a subset of artificial intelligence.",
|
||||
"metadata": {"language": "english", "topic": "AI"}
|
||||
},
|
||||
{
|
||||
"doc_id": "ml_zh_1",
|
||||
"text": "机器学习是人工智能的一个子集。", # Same content in Chinese
|
||||
"metadata": {"language": "chinese", "topic": "AI"}
|
||||
},
|
||||
{
|
||||
"doc_id": "ml_es_1",
|
||||
"text": "El aprendizaje automático es un subconjunto de la inteligencia artificial.", # Spanish
|
||||
"metadata": {"language": "spanish", "topic": "AI"}
|
||||
},
|
||||
{
|
||||
"doc_id": "other_1",
|
||||
"text": "Database systems store and retrieve information efficiently.",
|
||||
"metadata": {"language": "english", "topic": "database"}
|
||||
}
|
||||
]
|
||||
|
||||
multilingual_queries = [
|
||||
{
|
||||
"query": "AI learning", # English query
|
||||
"expected": ["ml_en_1", "ml_zh_1", "ml_es_1"],
|
||||
"explanation": "Dense embeddings (BGE-M3) should find similar content across languages"
|
||||
},
|
||||
{
|
||||
"query": "人工智能", # Chinese query for "artificial intelligence"
|
||||
"expected": ["ml_zh_1", "ml_en_1"],
|
||||
"explanation": "Dense should match Chinese query to related documents in any language"
|
||||
}
|
||||
]
|
||||
|
||||
await client.run_test_case(
|
||||
"Multilingual Matching (Dense Advantage)",
|
||||
multilingual_docs,
|
||||
multilingual_queries
|
||||
)
|
||||
|
||||
# Test Case 4: Technical Terms and Codes (Sparse is better)
|
||||
technical_docs = [
|
||||
{
|
||||
"doc_id": "error_1",
|
||||
"text": "Error code HTTP-403 indicates forbidden access to the resource.",
|
||||
"metadata": {"type": "error", "category": "http"}
|
||||
},
|
||||
{
|
||||
"doc_id": "error_2",
|
||||
"text": "The system returned status 500 for internal server problems.",
|
||||
"metadata": {"type": "error", "category": "http"}
|
||||
},
|
||||
{
|
||||
"doc_id": "config_1",
|
||||
"text": "Set parameter MAX_BUFFER_SIZE=8192 in the configuration file.",
|
||||
"metadata": {"type": "configuration"}
|
||||
},
|
||||
{
|
||||
"doc_id": "generic_1",
|
||||
"text": "The application encountered an issue during startup.",
|
||||
"metadata": {"type": "error", "category": "general"}
|
||||
}
|
||||
]
|
||||
|
||||
technical_queries = [
|
||||
{
|
||||
"query": "HTTP-403", # Exact error code
|
||||
"expected": ["error_1"],
|
||||
"explanation": "Sparse should match exact error code 'HTTP-403'"
|
||||
},
|
||||
{
|
||||
"query": "MAX_BUFFER_SIZE", # Exact parameter name
|
||||
"expected": ["config_1"],
|
||||
"explanation": "Sparse should find exact configuration parameter"
|
||||
}
|
||||
]
|
||||
|
||||
await client.run_test_case(
|
||||
"Technical Terms and Codes (Sparse Advantage)",
|
||||
technical_docs,
|
||||
technical_queries
|
||||
)
|
||||
|
||||
# Test Case 5: Conceptual Understanding (Dense is better)
|
||||
conceptual_docs = [
|
||||
{
|
||||
"doc_id": "happy_1",
|
||||
"text": "She was filled with joy and couldn't stop smiling.",
|
||||
"metadata": {"emotion": "positive"}
|
||||
},
|
||||
{
|
||||
"doc_id": "happy_2",
|
||||
"text": "His elation was evident as he celebrated the victory.",
|
||||
"metadata": {"emotion": "positive"}
|
||||
},
|
||||
{
|
||||
"doc_id": "sad_1",
|
||||
"text": "Tears rolled down her face as she felt overwhelmed with sorrow.",
|
||||
"metadata": {"emotion": "negative"}
|
||||
},
|
||||
{
|
||||
"doc_id": "neutral_1",
|
||||
"text": "The meeting proceeded according to the scheduled agenda.",
|
||||
"metadata": {"emotion": "neutral"}
|
||||
}
|
||||
]
|
||||
|
||||
conceptual_queries = [
|
||||
{
|
||||
"query": "happiness and excitement", # Concept not exact words
|
||||
"expected": ["happy_1", "happy_2"],
|
||||
"explanation": "Dense should understand happiness concept despite different words (joy, elation)"
|
||||
},
|
||||
{
|
||||
"query": "melancholy mood", # Related to sadness
|
||||
"expected": ["sad_1"],
|
||||
"explanation": "Dense should connect 'melancholy' with 'sorrow' conceptually"
|
||||
}
|
||||
]
|
||||
|
||||
await client.run_test_case(
|
||||
"Conceptual Understanding (Dense Advantage)",
|
||||
conceptual_docs,
|
||||
conceptual_queries
|
||||
)
|
||||
|
||||
# Print summary
|
||||
print(f"\n{'='*80}")
|
||||
print("TEST SUMMARY")
|
||||
print(f"{'='*80}")
|
||||
print(f"Total test cases: {len(client.test_results)}")
|
||||
|
||||
for test in client.test_results:
|
||||
print(f"\n{test['name']}:")
|
||||
print(f" Documents: {test['documents']}")
|
||||
print(f" Queries: {test['queries']}")
|
||||
print(f" Total searches: {len(test['results'])}")
|
||||
|
||||
async def run_interactive_demo():
|
||||
"""Run an interactive demonstration."""
|
||||
client = TestClient()
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("INTERACTIVE RETRIEVAL PIPELINE DEMO")
|
||||
print("="*80)
|
||||
print("\nThis demo shows how dense and sparse retrieval work differently.")
|
||||
print("Dense is better for: semantic similarity, concepts, multilingual")
|
||||
print("Sparse is better for: exact names, codes, technical terms")
|
||||
|
||||
# Sample documents for interactive demo
|
||||
sample_docs = [
|
||||
{
|
||||
"doc_id": "python_intro",
|
||||
"text": "Python is a high-level programming language known for its simplicity and readability.",
|
||||
"metadata": {"category": "programming", "language": "english"}
|
||||
},
|
||||
{
|
||||
"doc_id": "python_syntax",
|
||||
"text": "def hello_world(): print('Hello, World!') is a simple Python function.",
|
||||
"metadata": {"category": "code", "language": "english"}
|
||||
},
|
||||
{
|
||||
"doc_id": "ml_basics",
|
||||
"text": "Machine learning algorithms learn patterns from data without explicit programming.",
|
||||
"metadata": {"category": "AI", "language": "english"}
|
||||
},
|
||||
{
|
||||
"doc_id": "深度学习",
|
||||
"text": "深度学习是机器学习的一个分支，使用神经网络处理复杂数据。",
|
||||
"metadata": {"category": "AI", "language": "chinese"}
|
||||
},
|
||||
{
|
||||
"doc_id": "api_error",
|
||||
"text": "API returned error code E-2001: Invalid authentication token provided.",
|
||||
"metadata": {"category": "error", "type": "api"}
|
||||
}
|
||||
]
|
||||
|
||||
print("\nIndexing sample documents...")
|
||||
await client.clear_documents()
|
||||
|
||||
for doc in sample_docs:
|
||||
await client.index_document(
|
||||
text=doc["text"],
|
||||
doc_id=doc["doc_id"],
|
||||
metadata=doc.get("metadata", {})
|
||||
)
|
||||
print(f" ✓ {doc['doc_id']}: {doc['text'][:60]}...")
|
||||
|
||||
# Interactive queries
|
||||
queries = [
|
||||
("coding simplicity", "Should find Python docs via semantic similarity"),
|
||||
("E-2001", "Should find exact error code via sparse search"),
|
||||
("neural networks", "Should find ML/DL docs including Chinese via dense"),
|
||||
("hello_world", "Should find exact function name via sparse")
|
||||
]
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("RUNNING COMPARISON QUERIES")
|
||||
print("="*80)
|
||||
|
||||
for query, explanation in queries:
|
||||
print(f"\nQuery: '{query}'")
|
||||
print(f"Expected: {explanation}")
|
||||
print("-" * 40)
|
||||
|
||||
# Compare all three modes
|
||||
modes_results = {}
|
||||
for mode in ["dense", "sparse", "hybrid"]:
|
||||
result = await client.search(query, mode=mode, top_k=5, rerank_top_k=3)
|
||||
|
||||
# Get top results based on mode
|
||||
if mode == "hybrid" and result.get("reranked_results"):
|
||||
top = [r["doc_id"] for r in result["reranked_results"][:3]]
|
||||
elif mode == "dense" and result.get("dense_results"):
|
||||
top = [r["doc_id"] for r in result["dense_results"][:3]]
|
||||
elif mode == "sparse" and result.get("sparse_results"):
|
||||
top = [r["doc_id"] for r in result["sparse_results"][:3]]
|
||||
else:
|
||||
top = []
|
||||
|
||||
modes_results[mode] = top
|
||||
print(f"{mode:8}: {top}")
|
||||
|
||||
# Show which mode performed best
|
||||
print("-" * 40)
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Test client for retrieval pipeline")
|
||||
parser.add_argument("--url", default="http://localhost:4242", help="Pipeline service URL")
|
||||
parser.add_argument("--mode", choices=["test", "demo"], default="test",
|
||||
help="Run mode: test (all test cases) or demo (interactive)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.mode == "test":
|
||||
asyncio.run(run_educational_tests())
|
||||
else:
|
||||
asyncio.run(run_interactive_demo())
|
||||
@@ -0,0 +1,80 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test script to verify the improvements made to the retrieval pipeline."""
|
||||
|
||||
import subprocess
|
||||
import time
|
||||
import requests
|
||||
import sys
|
||||
import signal
|
||||
|
||||
def test_server_startup():
|
||||
"""Test that the server starts without deprecation warnings."""
|
||||
print("=" * 60)
|
||||
print("Testing Server Startup (No Deprecation Warnings)")
|
||||
print("=" * 60)
|
||||
|
||||
# Start the server
|
||||
process = subprocess.Popen(
|
||||
[sys.executable, "main.py", "--port", "8004"],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True
|
||||
)
|
||||
|
||||
# Collect output for 5 seconds
|
||||
output_lines = []
|
||||
start_time = time.time()
|
||||
|
||||
while time.time() - start_time < 5:
|
||||
line = process.stdout.readline()
|
||||
if line:
|
||||
output_lines.append(line.strip())
|
||||
print(f" {line.strip()}")
|
||||
|
||||
# Check for deprecation warning
|
||||
has_warning = any("DeprecationWarning" in line or "on_event is deprecated" in line
|
||||
for line in output_lines)
|
||||
|
||||
if has_warning:
|
||||
print("\n❌ FAILED: Deprecation warning still present!")
|
||||
else:
|
||||
print("\n✅ PASSED: No deprecation warnings found!")
|
||||
|
||||
# Check for model loading messages
|
||||
has_model_info = any(
|
||||
"Model already cached" in line or
|
||||
"Downloading model" in line or
|
||||
"Reranker initialized successfully" in line
|
||||
for line in output_lines
|
||||
)
|
||||
|
||||
if has_model_info:
|
||||
print("✅ PASSED: Model loading information displayed!")
|
||||
else:
|
||||
print("❌ FAILED: No model loading information found!")
|
||||
|
||||
# Check for loading time display
|
||||
has_timing = any("initialized successfully in" in line for line in output_lines)
|
||||
|
||||
if has_timing:
|
||||
print("✅ PASSED: Model loading time displayed!")
|
||||
else:
|
||||
print("❌ FAILED: No loading time information found!")
|
||||
|
||||
# Clean up
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=2)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("Summary of Improvements:")
|
||||
print("=" * 60)
|
||||
print("1. FastAPI deprecation warning: FIXED ✅" if not has_warning else "1. FastAPI deprecation warning: NOT FIXED ❌")
|
||||
print("2. Model loading progress: ADDED ✅" if has_model_info else "2. Model loading progress: NOT ADDED ❌")
|
||||
print("3. Loading time display: ADDED ✅" if has_timing else "3. Loading time display: NOT ADDED ❌")
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_server_startup()
|
||||
print("\nTest completed!")
|
||||
@@ -0,0 +1,200 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Test script for the retrieval pipeline with external doc_id support."""
|
||||
|
||||
import httpx
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Service URLs
|
||||
DENSE_URL = "http://localhost:4240"
|
||||
SPARSE_URL = "http://localhost:4241"
|
||||
PIPELINE_URL = "http://localhost:4242"
|
||||
|
||||
async def test_sparse_service():
|
||||
"""Test the sparse service directly to ensure it handles external doc_ids."""
|
||||
logger.info("Testing sparse service with external doc_id...")
|
||||
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
# Test indexing with external doc_id
|
||||
test_doc = {
|
||||
"text": "Python is a high-level programming language known for its simplicity and readability.",
|
||||
"doc_id": "test_python_doc_001",
|
||||
"metadata": {"category": "programming", "language": "Python"}
|
||||
}
|
||||
|
||||
try:
|
||||
response = await client.post(f"{SPARSE_URL}/index", json=test_doc)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
logger.info(f"Sparse indexing result: {json.dumps(result, indent=2)}")
|
||||
|
||||
# Verify the doc_id matches what we sent
|
||||
if result.get("doc_id") == "test_python_doc_001":
|
||||
logger.info("✅ Sparse service correctly preserved external doc_id")
|
||||
else:
|
||||
logger.error(f"❌ Sparse service returned different doc_id: {result.get('doc_id')}")
|
||||
|
||||
# Test search
|
||||
search_query = {"query": "Python programming", "top_k": 5}
|
||||
response = await client.post(f"{SPARSE_URL}/search", json=search_query)
|
||||
response.raise_for_status()
|
||||
search_results = response.json()
|
||||
|
||||
if search_results:
|
||||
logger.info(f"✅ Sparse search returned {len(search_results)} results")
|
||||
first_result = search_results[0]
|
||||
logger.info(f"First result doc_id: {first_result.get('doc_id')}")
|
||||
if first_result.get('doc_id') == "test_python_doc_001":
|
||||
logger.info("✅ Search correctly returned our document with external doc_id")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Sparse service test failed: {e}")
|
||||
return False
|
||||
|
||||
async def test_pipeline():
|
||||
"""Test the complete retrieval pipeline."""
|
||||
logger.info("Testing retrieval pipeline...")
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
try:
|
||||
# First, clear the pipeline
|
||||
logger.info("Clearing pipeline...")
|
||||
response = await client.delete(f"{PIPELINE_URL}/clear")
|
||||
logger.info(f"Clear response: {response.json()}")
|
||||
|
||||
# Test documents
|
||||
test_documents = [
|
||||
{
|
||||
"text": "Python is renowned for its clean syntax and readability, making it ideal for beginners and experts alike.",
|
||||
"doc_id": "prog_python",
|
||||
"metadata": {"category": "programming", "subcategory": "languages"}
|
||||
},
|
||||
{
|
||||
"text": "Machine learning with Python involves libraries like scikit-learn, TensorFlow, and PyTorch for building AI models.",
|
||||
"doc_id": "ml_python",
|
||||
"metadata": {"category": "machine_learning", "subcategory": "tools"}
|
||||
},
|
||||
{
|
||||
"text": "JavaScript is the language of the web, enabling dynamic and interactive user interfaces in browsers.",
|
||||
"doc_id": "prog_javascript",
|
||||
"metadata": {"category": "programming", "subcategory": "web"}
|
||||
}
|
||||
]
|
||||
|
||||
# Index documents
|
||||
for doc in test_documents:
|
||||
logger.info(f"Indexing document: {doc['doc_id']}")
|
||||
response = await client.post(f"{PIPELINE_URL}/index", json=doc)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# Check both services succeeded
|
||||
dense_success = result.get("dense", {}).get("success", False)
|
||||
sparse_success = result.get("sparse", {}).get("success", False)
|
||||
|
||||
if dense_success and sparse_success:
|
||||
logger.info(f"✅ Document {doc['doc_id']} indexed successfully in both services")
|
||||
else:
|
||||
logger.error(f"❌ Indexing failed for {doc['doc_id']}")
|
||||
logger.error(f" Dense: {result.get('dense')}")
|
||||
logger.error(f" Sparse: {result.get('sparse')}")
|
||||
|
||||
# Wait a moment for indexing to complete
|
||||
await asyncio.sleep(1)
|
||||
|
||||
# Test search in different modes
|
||||
search_query = "Python programming language"
|
||||
logger.info(f"\nTesting search with query: '{search_query}'")
|
||||
|
||||
for mode in ["dense", "sparse", "hybrid"]:
|
||||
logger.info(f"\n--- Testing {mode} search ---")
|
||||
search_request = {
|
||||
"query": search_query,
|
||||
"mode": mode,
|
||||
"top_k": 10,
|
||||
"rerank_top_k": 5,
|
||||
"skip_reranking": False if mode == "hybrid" else True
|
||||
}
|
||||
|
||||
response = await client.post(f"{PIPELINE_URL}/search", json=search_request)
|
||||
response.raise_for_status()
|
||||
results = response.json()
|
||||
|
||||
# Log results summary
|
||||
if mode == "dense":
|
||||
dense_results = results.get("dense_results", [])
|
||||
if dense_results:
|
||||
logger.info(f"✅ Dense search returned {len(dense_results)} results")
|
||||
logger.info(f" Top result: {dense_results[0]['doc_id']} (score: {dense_results[0]['score']:.4f})")
|
||||
else:
|
||||
logger.error("❌ No dense results returned")
|
||||
|
||||
elif mode == "sparse":
|
||||
sparse_results = results.get("sparse_results", [])
|
||||
if sparse_results:
|
||||
logger.info(f"✅ Sparse search returned {len(sparse_results)} results")
|
||||
logger.info(f" Top result: {sparse_results[0]['doc_id']} (score: {sparse_results[0]['score']:.4f})")
|
||||
else:
|
||||
logger.error("❌ No sparse results returned")
|
||||
|
||||
elif mode == "hybrid":
|
||||
dense_results = results.get("dense_results", [])
|
||||
sparse_results = results.get("sparse_results", [])
|
||||
reranked_results = results.get("reranked_results", [])
|
||||
|
||||
logger.info(f"✅ Hybrid search results:")
|
||||
logger.info(f" Dense: {len(dense_results)} results")
|
||||
logger.info(f" Sparse: {len(sparse_results)} results")
|
||||
logger.info(f" Reranked: {len(reranked_results)} results")
|
||||
|
||||
if reranked_results:
|
||||
logger.info(f" Top reranked result: {reranked_results[0]['doc_id']} (score: {reranked_results[0]['rerank_score']:.4f})")
|
||||
|
||||
# Check statistics
|
||||
stats = results.get("statistics", {})
|
||||
if stats:
|
||||
logger.info(f" Overlap: {stats.get('overlap_count', 0)} documents ({stats.get('overlap_percentage', 0):.1f}%)")
|
||||
|
||||
logger.info("\n✅ All pipeline tests completed successfully!")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Pipeline test failed: {e}")
|
||||
import traceback
|
||||
logger.error(traceback.format_exc())
|
||||
return False
|
||||
|
||||
async def main():
|
||||
"""Run all tests."""
|
||||
logger.info("Starting retrieval pipeline tests...")
|
||||
logger.info("Make sure all three services are running:")
|
||||
logger.info(" - Dense service on port 4240")
|
||||
logger.info(" - Sparse service on port 4241")
|
||||
logger.info(" - Pipeline service on port 4242")
|
||||
logger.info("")
|
||||
|
||||
# Test sparse service first
|
||||
sparse_ok = await test_sparse_service()
|
||||
|
||||
if sparse_ok:
|
||||
logger.info("\n" + "="*50 + "\n")
|
||||
# Test full pipeline
|
||||
pipeline_ok = await test_pipeline()
|
||||
|
||||
if pipeline_ok:
|
||||
logger.info("\n🎉 All tests passed successfully!")
|
||||
else:
|
||||
logger.info("\n⚠️ Some pipeline tests failed")
|
||||
else:
|
||||
logger.error("\n⚠️ Sparse service test failed - skipping pipeline tests")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,15 @@
|
||||
import pytest
|
||||
from fusion import weighted_score_fusion
|
||||
|
||||
def test_weighted_score_fusion_preserves_top_score_on_duplicate_doc_id():
|
||||
"""Verify weighted score fusion preserves top score when duplicate doc_ids exist in ranked list."""
|
||||
ranked_lists = {
|
||||
"dense": [("doc1", 0.95), ("doc2", 0.80), ("doc1", 0.10)],
|
||||
"sparse": [("doc1", 10.0), ("doc2", 5.0)]
|
||||
}
|
||||
|
||||
results = weighted_score_fusion(ranked_lists)
|
||||
|
||||
assert results[0][0] == "doc1"
|
||||
assert results[1][0] == "doc2"
|
||||
assert results[0][1] > results[1][1]
|
||||
@@ -0,0 +1,92 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-6",
|
||||
"run_id": "20260729T183512Z-3_6-a398d0f7",
|
||||
"created_at": "2026-07-29T18:35:12.571244+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/validation/runs/20260729T183512Z-3_6-a398d0f7",
|
||||
"artifacts": {
|
||||
"evidence.json": "280aae37e6da541881550a29ce799b81e9d381db3ea541eebbbddb6a72295788",
|
||||
"receipts.json": "3db117d3cc0779290333a8da33db529dccd88af2434c966a7930e2dad9434e1e",
|
||||
"manifest.json": "c305591efaeb1f85d6bab118e6bc8761a6e805feb0abf2b8146403be151dbcbf"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/campaign.py",
|
||||
"sha256": "59fdde0060261dd89f8aeffa8a2e3fed0887510c1b4b29bee50a47698e514dde",
|
||||
"bytes": 5019
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/evaluate.py",
|
||||
"sha256": "5ec2d32d386d396b909b1709cac66d55717753dccfcbb66c2d4e0e48b3049371",
|
||||
"bytes": 37506
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/fusion.py",
|
||||
"sha256": "79824352be064e5059cf90ba58bab81f27c4ca94b059ab463b0cbb6f38b07c69",
|
||||
"bytes": 5483
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"sparse": {
|
||||
"recall@k": 0.8333333333333334,
|
||||
"mrr": 0.7916666666666666,
|
||||
"ndcg@k": 0.8025774794642881,
|
||||
"latency_ms": {
|
||||
"mean": 29.646694504966337,
|
||||
"p50": 0.10129157453775406,
|
||||
"p95": 354.74016703665257
|
||||
}
|
||||
},
|
||||
"dense": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 1.0,
|
||||
"ndcg@k": 1.0,
|
||||
"latency_ms": {
|
||||
"mean": 164.16664922144264,
|
||||
"p50": 162.6114579848945,
|
||||
"p95": 183.55158297345042
|
||||
}
|
||||
},
|
||||
"rrf": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 1.0,
|
||||
"ndcg@k": 1.0,
|
||||
"latency_ms": {
|
||||
"mean": 193.8208055216819,
|
||||
"p50": 165.03987507894635,
|
||||
"p95": 517.3575421795249
|
||||
}
|
||||
},
|
||||
"weighted": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 0.9583333333333334,
|
||||
"ndcg@k": 0.9692441461309548,
|
||||
"latency_ms": {
|
||||
"mean": 193.82462510839105,
|
||||
"p50": 165.04266718402505,
|
||||
"p95": 517.3604171723127
|
||||
}
|
||||
},
|
||||
"rerank": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 0.9027777777777778,
|
||||
"ndcg@k": 0.9275774794642881,
|
||||
"latency_ms": {
|
||||
"mean": 236.57109351673475,
|
||||
"p50": 210.11475007981062,
|
||||
"p95": 552.1125011146069
|
||||
}
|
||||
}
|
||||
},
|
||||
"acceptance": {
|
||||
"real_dense_model_loaded": true,
|
||||
"real_cross_encoder_loaded": true,
|
||||
"identical_labelled_queries_for_all_methods": true,
|
||||
"all_required_query_categories_present": true,
|
||||
"sparse_dense_rrf_weighted_reranked_measured": true,
|
||||
"recall_mrr_ndcg_and_latency_measured": true,
|
||||
"rank_changes_retained": true,
|
||||
"hybrid_recall_not_below_best_single": true
|
||||
}
|
||||
}
|
||||
+3804
File diff suppressed because it is too large
Load Diff
+91
@@ -0,0 +1,91 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-6",
|
||||
"run_id": "20260729T183512Z-3_6-a398d0f7",
|
||||
"created_at": "2026-07-29T18:35:12.571244+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/validation/runs/20260729T183512Z-3_6-a398d0f7",
|
||||
"artifacts": {
|
||||
"evidence.json": "280aae37e6da541881550a29ce799b81e9d381db3ea541eebbbddb6a72295788",
|
||||
"receipts.json": "3db117d3cc0779290333a8da33db529dccd88af2434c966a7930e2dad9434e1e"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/campaign.py",
|
||||
"sha256": "59fdde0060261dd89f8aeffa8a2e3fed0887510c1b4b29bee50a47698e514dde",
|
||||
"bytes": 5019
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/evaluate.py",
|
||||
"sha256": "5ec2d32d386d396b909b1709cac66d55717753dccfcbb66c2d4e0e48b3049371",
|
||||
"bytes": 37506
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/retrieval-pipeline/fusion.py",
|
||||
"sha256": "79824352be064e5059cf90ba58bab81f27c4ca94b059ab463b0cbb6f38b07c69",
|
||||
"bytes": 5483
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"sparse": {
|
||||
"recall@k": 0.8333333333333334,
|
||||
"mrr": 0.7916666666666666,
|
||||
"ndcg@k": 0.8025774794642881,
|
||||
"latency_ms": {
|
||||
"mean": 29.646694504966337,
|
||||
"p50": 0.10129157453775406,
|
||||
"p95": 354.74016703665257
|
||||
}
|
||||
},
|
||||
"dense": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 1.0,
|
||||
"ndcg@k": 1.0,
|
||||
"latency_ms": {
|
||||
"mean": 164.16664922144264,
|
||||
"p50": 162.6114579848945,
|
||||
"p95": 183.55158297345042
|
||||
}
|
||||
},
|
||||
"rrf": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 1.0,
|
||||
"ndcg@k": 1.0,
|
||||
"latency_ms": {
|
||||
"mean": 193.8208055216819,
|
||||
"p50": 165.03987507894635,
|
||||
"p95": 517.3575421795249
|
||||
}
|
||||
},
|
||||
"weighted": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 0.9583333333333334,
|
||||
"ndcg@k": 0.9692441461309548,
|
||||
"latency_ms": {
|
||||
"mean": 193.82462510839105,
|
||||
"p50": 165.04266718402505,
|
||||
"p95": 517.3604171723127
|
||||
}
|
||||
},
|
||||
"rerank": {
|
||||
"recall@k": 1.0,
|
||||
"mrr": 0.9027777777777778,
|
||||
"ndcg@k": 0.9275774794642881,
|
||||
"latency_ms": {
|
||||
"mean": 236.57109351673475,
|
||||
"p50": 210.11475007981062,
|
||||
"p95": 552.1125011146069
|
||||
}
|
||||
}
|
||||
},
|
||||
"acceptance": {
|
||||
"real_dense_model_loaded": true,
|
||||
"real_cross_encoder_loaded": true,
|
||||
"identical_labelled_queries_for_all_methods": true,
|
||||
"all_required_query_categories_present": true,
|
||||
"sparse_dense_rrf_weighted_reranked_measured": true,
|
||||
"recall_mrr_ndcg_and_latency_measured": true,
|
||||
"rank_changes_retained": true,
|
||||
"hybrid_recall_not_below_best_single": true
|
||||
}
|
||||
}
|
||||
+28
@@ -0,0 +1,28 @@
|
||||
[
|
||||
{
|
||||
"kind": "local-model-execution",
|
||||
"models": {
|
||||
"dense": {
|
||||
"provider": "local Hugging Face transformers",
|
||||
"model": "Qwen/Qwen3-Embedding-0.6B",
|
||||
"cached_revision": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
|
||||
"class": "Qwen3Model",
|
||||
"pooling": "last",
|
||||
"parameters": 595776512,
|
||||
"device": "cpu"
|
||||
},
|
||||
"reranker": {
|
||||
"provider": "local Hugging Face transformers",
|
||||
"model": "cross-encoder/ms-marco-MiniLM-L-6-v2",
|
||||
"cached_revision": "c5ee24cb16019beea0893ab7796b1df96625c6b8",
|
||||
"class": "BertForSequenceClassification",
|
||||
"parameters": 22713601,
|
||||
"device": "cpu"
|
||||
},
|
||||
"sparse": {
|
||||
"implementation": "rank_bm25.BM25Okapi"
|
||||
}
|
||||
},
|
||||
"note": "No remote API or credential was used; complete rankings and timings are in evidence.json."
|
||||
}
|
||||
]
|
||||
Reference in New Issue
Block a user