retrieval-pipeline¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/retrieval-pipeline
项目说明¶
Hybrid Retrieval Pipeline with Neural Reranking¶
An educational retrieval pipeline that combines dense embeddings, sparse search, and neural reranking to demonstrate the strengths and weaknesses of different retrieval methods.
🎯 Educational Goals¶
This project demonstrates: 1. Dense vs Sparse Retrieval: When each method excels and why 2. Hybrid Search: Combining multiple retrieval methods for better results 3. Neural Reranking: Using transformer models to reorder search results 4. Parallel Processing: Efficient indexing and searching across multiple services 5. Real-world Patterns: Production-ready 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 Embeddings (Semantic Search)¶
- Model: BGE-M3 (multilingual, 1024-dim vectors)
- Strengths:
- Semantic similarity (finds related concepts)
- Cross-lingual search (works across languages)
- Conceptual understanding (handles synonyms)
- Context awareness (understands meaning)
- Weaknesses:
- May miss exact strings
- Less effective for codes/IDs
- Computationally expensive
Sparse Search (BM25)¶
- Algorithm: BM25 (Best Matching 25)
- Strengths:
- Exact term matching
- Specific names and codes
- Technical identifiers
- Fast and efficient
- Weaknesses:
- No semantic understanding
- Language-specific
- Requires exact terms
Result Fusion (RRF / Weighted)¶
- Module:
fusion.py - Purpose: Merge the separately-ranked dense and sparse candidate lists into one unified pool before reranking (the fusion stage)
- Two strategies (both discussed in the book):
- Reciprocal Rank Fusion (RRF):
score(d) = Σ 1/(k + rank_r(d)),k=60. Uses only ranks, so it never has to compare a cosine similarity against a BM25 score — robust and scale-free. - Weighted score fusion: min-max normalize each list to
[0,1], then take a weighted sum. Keeps the raw relevance signal, at the cost of tuning the scale alignment.
Neural Reranking¶
- Model: BGE-Reranker-v2-M3 (production);
BAAI/bge-reranker-base(lighter, used byevaluate.py) - Purpose: Re-score and reorder the fused candidate pool
- Benefits:
- Better relevance ranking
- Combines signals from both methods
- Context-aware scoring
🚀 Quick Start¶
Prerequisites¶
- Python 3.8+
- macOS with M1/M2 chip (or modify device settings for other platforms)
- At least 8GB RAM
- About 5GB disk space for models
Installation¶
# Clone the repository
cd projects/week3/retrieval-pipeline
# Install dependencies
pip install -r requirements.txt
# The models will be downloaded automatically on first run:
# - BGE-M3: ~2.3GB
# - BGE-Reranker-v2-M3: ~1.1GB
Running the Services¶
- Start all services (recommended):
This will start: - Dense embedding service on port 4240 - Sparse embedding service on port 4241 - Retrieval pipeline on port 4242
- Or start individually:
Testing the Pipeline¶
- Run educational tests:
This runs comprehensive test cases showing when dense vs sparse excels.
- Run interactive demo:
This demonstrates real queries with explanations.
- Access API documentation:
🧪 Offline Evaluation CLI (evaluate.py)¶
test_client.py / demo.py above require the three microservices (ports
4240/4241/4242) to be running. evaluate.py runs the entire pipeline in a
single process, fully offline, so you can reproduce the "each stage improves
the ranking" story with local models and no services.
It walks the complete pipeline — chunk → embed → retrieve → fuse → rerank —
on a small labelled eval set and prints a stage-by-stage comparison table plus a
per-query breakdown. The CLI has a full Chinese --help:
python evaluate.py --help # 中文帮助：语料/查询/阶段/top-k/模型/输出等
python evaluate.py # 内置评测集，完整对比表（默认）
python evaluate.py --no-dense # 仅 BM25，纯离线、无需任何模型
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 # 结果写入 JSON
Local components (each stage is a real model / algorithm, not a mock):
| Stage | Component (default) | Offline? |
|---|---|---|
| chunk | character-window splitter | ✅ pure Python |
| sparse | BM25 (rank_bm25) |
✅ no model download |
| dense | sentence-transformers/all-MiniLM-L6-v2 (~90MB) |
✅ via transformers (multilingual: swap in Qwen/Qwen3-Embedding-0.6B / BAAI/bge-m3) |
| fuse | RRF + weighted (fusion.py) |
✅ pure Python |
| rerank | BAAI/bge-reranker-base (~1.1GB, first run downloads) |
✅ once cached |
Note:
--no-denseneeds no ML model at all (BM25 only). The dense and rerank stages need local models; on first run they download from HuggingFace, after which--offlinekeeps everything on the local cache. On Apple Silicon, sometransformersbuilds emitNaNon the MPS device — the CLI detects this and automatically falls back to CPU, so results are always finite.
Real output (reproduced on this machine)¶
The built-in eval set has two deliberately hard clusters: near-duplicate codes
(XR-7001..XR-7006, HTTP-400..HTTP-500) that break dense retrieval (the vectors
are almost identical, only exact term matching finds the right one), and
zero-lexical-overlap paraphrases (query "reclaiming unused heap space without
programmer effort" → doc "Automatic memory management frees developers…") that
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
逐条查询 MRR 明细（1.00=正确文档排在第 1 位）
Query BM25 Dense RRF Wgt Rerank
------------------------------------------------------------------------------
XR-7003 1.00 0.50 1.00 1.00 1.00
XR-7005 1.00 0.50 1.00 1.00 1.00
HTTP-403 1.00 1.00 1.00 1.00 1.00
HTTP-400 1.00 1.00 1.00 1.00 0.50
a beginner friendly language with tidy... 1.00 1.00 1.00 1.00 1.00
reclaiming unused heap space without p... 0.00 1.00 1.00 1.00 1.00
how vegetation turns light into food 0.50 1.00 1.00 0.50 1.00
hiding a note so eavesdroppers cannot ... 1.00 1.00 1.00 1.00 1.00
how does water move between the ocean ... 1.00 1.00 1.00 1.00 1.00
how are volcanoes formed from molten rock 1.00 1.00 1.00 1.00 1.00
How to read it:
- BM25 (sparse) nails every exact code but collapses on the two paraphrase
queries (reclaiming…=0.00, vegetation…=0.50).
- Dense is the mirror image: perfect on paraphrases, but confuses the
near-duplicate codes (XR-7003/XR-7005=0.50 — it ranks a sibling code first).
- Hybrid-RRF combines the two and reaches a perfect 1.00 across the board:
fusion rescues both single-method failures. This is the headline of Exp. 3-6.
- Weighted fusion is strong too but less robust than RRF here (the vegetation
query drops to 0.50 because a lexical near-match distorts the normalized scores —
exactly the "scale alignment is hard to tune" caveat).
- Reranking: on this 17-doc toy corpus RRF is already optimal, so reranking
has no headroom; a cross-encoder is a semantic matcher and can even slightly
reorder a trivial exact-code lookup (HTTP-400). Its value grows on larger
candidate pools and natural-language queries — see the single-query trace below.
The single-query trace makes the fusion mechanism concrete:
$ 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 ... # dense ranks the wrong sibling first
2. xr_7003 score= 0.5195 Product model XR-7003 ...
[Hybrid-RRF]
1. xr_7003 score= 0.0325 Product model XR-7003 ... # fusion promotes the exact match to #1
📊 Educational Test Cases¶
Test Case 1: Semantic Similarity (Dense Wins)¶
# Query: "kitty behavior"
# Documents contain: "feline", "cat"
# Dense finds semantic match, sparse misses
Test Case 2: Exact Names (Sparse Wins)¶
Test Case 3: Multilingual (Dense Wins)¶
# Query: "人工智能" (Chinese for AI)
# Dense finds AI docs in any language
# Sparse only finds Chinese text
Test Case 4: Technical Codes (Sparse Wins)¶
Test Case 5: Concepts (Dense Wins)¶
# Query: "happiness and excitement"
# Documents contain: "joy", "elation"
# Dense understands emotional concepts
🔧 API Endpoints¶
Index Document¶
POST /index
{
"text": "Document content",
"doc_id": "optional_id",
"metadata": {"category": "example"}
}
Search¶
POST /search
{
"query": "search terms",
"mode": "hybrid", # or "dense" or "sparse"
"top_k": 20,
"rerank_top_k": 10
}
Response includes: - Original dense rankings and scores - Original sparse rankings and scores - Final reranked results - Rank change statistics - Performance metrics
Statistics¶
List Documents¶
📈 Understanding the Results¶
The search response provides educational insights:
{
"dense_results": [...], # Top results from semantic search
"sparse_results": [...], # Top results from BM25
"reranked_results": [ # Final reranked results
{
"rank": 1,
"doc_id": "doc_1",
"rerank_score": 0.95,
"original_ranks": {
"dense": 3, # Was rank 3 in dense
"sparse": 5 # Was rank 5 in sparse
},
"rank_changes": [
"dense: +2", # Moved up 2 positions
"sparse: +4" # Moved up 4 positions
]
}
],
"statistics": {
"overlap_percentage": 30.0, # How much dense/sparse agree
"avg_dense_rank_change": 1.5,
"avg_sparse_rank_change": 2.1
}
}
🎓 Learning Exercises¶
- Experiment with queries:
- Try synonyms vs exact terms
- Test multilingual queries
-
Use technical codes
-
Modify parameters:
- Change
top_kto retrieve more/fewer candidates - Skip reranking to see raw results
-
Try different search modes
-
Analyze patterns:
- When does hybrid outperform single methods?
- How much do dense and sparse results overlap?
- Which queries benefit most from reranking?
🔍 Troubleshooting¶
Services won't start¶
- Check ports 4240-4242 are free
- Ensure models downloaded properly
- Check Python version (3.8+)
Out of memory¶
- Reduce batch sizes in config.py
- Use CPU instead of MPS/CUDA
- Enable FP16 mode
Slow performance¶
- First run downloads models (be patient)
- Subsequent runs use cached models
- Consider using GPU if available
📝 Configuration¶
Edit config.py to adjust:
- Service URLs
- Model parameters
- Retrieval settings
- Reranking parameters
🏛️ Project Structure¶
retrieval-pipeline/
├── config.py # Configuration settings (incl. fusion_method / rrf_k)
├── document_store.py # In-memory document storage
├── retrieval_client.py # Client for dense/sparse services
├── reranker.py # BGE-Reranker implementation
├── fusion.py # Result fusion: RRF + weighted score fusion
├── retrieval_pipeline.py # Main pipeline orchestration (uses fusion.py)
├── evaluate.py # Offline single-process eval CLI (Chinese --help)
├── main.py # FastAPI server
├── test_client.py # Educational test cases
├── demo.py # Interactive demonstration
├── requirements.txt # Python dependencies
├── start_all_services.sh # Start script
├── stop_all_services.sh # Stop script
└── README.md # This file
🚦 Performance Considerations¶
- Indexing: Parallel indexing to both services
- Search: Parallel retrieval, then sequential reranking
- Memory: ~4GB for models, plus document storage
- Latency:
- Dense: ~50-100ms per query
- Sparse: ~10-30ms per query
- Reranking: ~100-200ms for 20 documents
🎯 Key Takeaways¶
- No single method is best: Dense and sparse have complementary strengths
- Hybrid search wins: Combining methods typically improves results
- Reranking matters: Neural reranking significantly improves relevance
- Parallel processing: Essential for production performance
- Educational value: Understanding trade-offs helps choose the right approach
📚 Further Reading¶
📄 License¶
This is an educational project for learning purposes.
源代码¶
config.py¶
"""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
demo.py¶
"""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)
document_store.py¶
"""Document store for the retrieval pipeline."""
import json
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."""
self.documents[doc_id] = {
"doc_id": doc_id,
"text": text,
"metadata": metadata or {},
"indexed_at": datetime.now().isoformat()
}
# Update metadata index
if metadata:
for key, value in metadata.items():
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()
}
}
evaluate.py¶
"""混合检索流水线离线评测 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."},
# 较长文档：话题彼此独立，用于演示分块阶段（会被切成多个 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"]},
{"query": "XR-7005", "expected": ["xr_7005"]},
{"query": "HTTP-403", "expected": ["http_403"]},
{"query": "HTTP-400", "expected": ["http_400"]},
# 语义改写查询：与文档几乎无共同词，稠密靠语义命中，稀疏无从匹配
{"query": "a beginner friendly language with tidy syntax", "expected": ["sem_readable"]},
{"query": "reclaiming unused heap space without programmer effort", "expected": ["sem_gc"]},
{"query": "how vegetation turns light into food", "expected": ["sem_photo"]},
{"query": "hiding a note so eavesdroppers cannot understand it", "expected": ["sem_crypto"]},
# 长文档语义查询：命中的长文档会先被分块，再由某个 chunk 召回、重排
{"query": "how does water move between the ocean and the sky", "expected": ["doc_watercycle"]},
{"query": "how are volcanoes formed from molten rock", "expected": ["doc_volcano"]},
]
# ---------------------------------------------------------------------------
# 分块(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_chunks = self.bm25.search(query, top_k)
sparse_docs = chunks_to_docs(sparse_chunks, self.chunk_to_doc)
out: Dict[str, List[Tuple[str, float]]] = {"sparse": sparse_docs}
if self.dense is not None:
dense_chunks = self.dense.search(query, top_k)
dense_docs = chunks_to_docs(dense_chunks, self.chunk_to_doc)
out["dense"] = dense_docs
ranked_lists = {"dense": dense_docs, "sparse": sparse_docs}
weights = {"dense": self.args.dense_weight, "sparse": self.args.sparse_weight}
rrf = fuse(ranked_lists, method="rrf", k=self.args.k_rrf, weights=weights)
weighted = fuse(ranked_lists, method="weighted", weights=weights)
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]
reranked = self.reranker.rerank(query, candidates, self.args.rerank_top_k)
out["rerank"] = reranked
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 = []
t0 = time.time()
for spec in queries:
query = spec["query"]
gold = spec.get("expected", [])
results = pipeline.run_query(query)
record = {"query": query, "expected": gold, "methods": {}}
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))
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)
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())
fusion.py¶
"""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 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(dict(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')")
main.py¶
"""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()
reranker.py¶
"""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
if len(scores.shape) > 1:
scores = 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
retrieval_client.py¶
"""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)}
retrieval_pipeline.py¶
"""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
}
}
test_client.py¶
"""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())
test_improvements.py¶
#!/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!")
test_pipeline.py¶
#!/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())
restart_services.sh¶
#!/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'"
start_all_services.sh¶
#!/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"
stop_all_services.sh¶
#!/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."