dense-embedding¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/dense-embedding
项目说明¶
Vector Similarity Search Service¶
An educational HTTP service for vector similarity search using BGE-M3 embeddings with configurable ANNOY or HNSW indexing backends.
命令行工具：稠密检索与 ANN 对比（cli.py，实验 3-4）¶
除了上面的 HTTP 服务，本项目还提供一个开箱即用、可离线复现的命令行工具 cli.py，
把书中实验 3-4 的两个观察点直接跑成可量化的数字，无需先启动服务：
- 稠密嵌入检索的语义能力——在带标注的小型语料上计算
recall@k / precision@k / MRR； - ANN 索引后端对比（实验 3-4 的重点）——复用服务端
indexing.py里的 ANNOY / HNSW 实现，测量二者相对精确暴力检索的召回率、建索引耗时与查询延迟。
用法¶
# 1) 单条稠密查询（默认查询 "a cat playing"，需要嵌入模型）
python cli.py -q "model distillation" -k 3
# 2) 检索质量评测：recall@k / precision@k / MRR
python cli.py --eval
# 2') 离线复现：用已缓存的小模型（无需下载 2.3GB 的 BGE-M3）
python cli.py --embedding-model sentence-transformers/all-MiniLM-L6-v2 --eval
# 3) ANN 后端对比（合成向量，完全离线、无需任何模型）
python cli.py --compare-ann -k 10
python cli.py --compare-ann --backend hnsw --hnsw-ef-search 200 -k 10 # 调 ef_search 看召回随之上升
# 自定义语料 / 标注 / 输出
python cli.py --corpus my.json --labels my_labels.json --eval -o result.json
python cli.py --help 提供完整的中文参数说明（--corpus / --query / --embedding-model /
--top-k / --output，以及 ANN 对比的各项索引超参）。
常用参数¶
| 参数 | 说明 |
|---|---|
-q, --query |
查询字符串（默认 a cat playing） |
-c, --corpus |
语料文件（.json 数组 或 .jsonl 每行一篇）；缺省用内置示例语料 |
-k, --top-k |
返回前 k 条结果（默认 5） |
-o, --output |
把结果 / 评测指标写入 JSON 文件 |
--embedding-model |
嵌入模型名（默认 BAAI/bge-m3；离线可用 sentence-transformers/all-MiniLM-L6-v2） |
--pooling |
池化方式 auto（bge* 用 cls，其余 mean）/ mean / cls |
--eval |
在标注集上评测 recall@k / precision@k / MRR |
--compare-ann |
对比 ANNOY / HNSW（合成向量，无需模型） |
--ann-base / --ann-dim / --ann-queries |
合成底库规模 / 维度 / 查询数（默认 3000 / 128 / 100） |
--annoy-n-trees / --hnsw-M / --hnsw-ef-search |
两类 ANN 的关键索引超参 |
实测结果（真实运行，非杜撰）¶
稠密检索质量（内置 12 篇语料，all-MiniLM-L6-v2，离线）：
其中查询 a cat playing 的相关文档只用 kitten / feline 表达、不含字面 "cat"，
稠密检索仍把它们排到第 1、2 名——这正是稠密相对稀疏 BM25（实验 3-5 会漏召回）的语义优势。
ANN 后端对比（3000 条 128 维随机单位向量，100 条查询，top-10）：HNSW 的召回率随
ef_search 单调上升，体现"精度 / 速度"取舍：
| 配置 | recall@10 | 平均查询延迟 |
|---|---|---|
HNSW ef_search=20 |
0.562 | 0.05 ms |
HNSW ef_search=200 |
0.991 | 0.25 ms |
环境提示：本工具对每个后端会先做"自查询自身向量"的健康检查。在部分 macOS/arm64 环境下，
annoy==1.17.3的预编译轮子存在缺陷（连查询库中已有向量都只返回它自己），此时工具会打印[警告] ...疑似当前环境下损坏并把该后端的数字标记为不可信。HNSW 不受影响。若要复现完整的 ANNOY vs HNSW 对比，请在 annoy 正常工作的环境（如 Linux x86_64）中运行。
Features¶
- BGE-M3 Model: State-of-the-art multilingual embedding model supporting:
- Dense embeddings for semantic search
- Multi-language support (100+ languages)
-
Long context (up to 8192 tokens)
-
Dual Indexing Backends:
- ANNOY (Approximate Nearest Neighbors Oh Yeah): Fast, memory-efficient tree-based index
-
HNSW (Hierarchical Navigable Small World): High-precision graph-based index
-
Educational Logging: Extensive debug logs showing:
- Embedding generation process
- Index operations (insert/delete/search)
- Performance metrics
-
Vector statistics
-
RESTful API: Clean HTTP endpoints for:
- Document indexing (insert/update)
- Document deletion
- Similarity search
-
Statistics and monitoring
-
In-Memory Storage: Pure in-memory operation for simplicity (no persistence)
Architecture¶
┌──────────────────┐
│ HTTP Client │
└────────┬─────────┘
│
▼
┌──────────────────┐
│ FastAPI Server │
└────────┬─────────┘
│
▼
┌────┴────┐
│ │
▼ ▼
┌──────────┐ ┌──────────────┐
│ Document │ │ Embedding │
│ Store │ │ Service │
└──────────┘ │ (BGE-M3) │
└──────┬─────────┘
│
▼
┌─────────┴──────────┐
│ │
▼ ▼
┌──────────┐ ┌──────────┐
│ ANNOY │ │ HNSW │
│ Index │ │ Index │
└──────────┘ └──────────┘
Installation¶
Prerequisites¶
- Python 3.8 or higher
- macOS (optimized for M1/M2 chips) or Linux
- At least 4GB RAM (8GB recommended)
- CUDA-compatible GPU (optional, for faster embedding generation)
Setup¶
- Install dependencies:
- Download BGE-M3 model (will be downloaded automatically on first run):
- Model size: ~2.3GB
- Will be cached in HuggingFace cache directory
Usage¶
Starting the Service¶
With HNSW Index (Default)¶
With ANNOY Index¶
With Custom Configuration¶
Available Options¶
--index-type: Choose index backend (annoyorhnsw, default:hnsw)--host: Server host (default:0.0.0.0)--port: Server port (default:4240)--debug: Enable debug mode with verbose logging--show-embeddings: Show embedding vectors in logs (educational)
API Documentation¶
Once the service is running, visit: - Interactive docs: http://localhost:4240/docs - OpenAPI schema: http://localhost:4240/openapi.json
API Endpoints¶
1. Index Document¶
Index a new document or update existing one.
Request:
{
"text": "Machine learning is a subset of artificial intelligence.",
"doc_id": "doc_001", // Optional, auto-generated if not provided
"metadata": { // Optional metadata
"category": "AI",
"author": "John Doe"
}
}
Response:
{
"success": true,
"doc_id": "doc_001",
"message": "Document indexed successfully using hnsw",
"index_size": 1
}
2. Search Documents¶
Search for similar documents.
Request:
Response:
{
"success": true,
"query": "What is deep learning?",
"results": [
{
"doc_id": "doc_001",
"score": 0.8543,
"text": "Machine learning is a subset...",
"metadata": {"category": "AI"},
"rank": 1
}
],
"total_results": 5,
"search_time_ms": 12.5
}
3. Delete Document¶
Delete a document from the index.
Request:
Response:
4. Get Statistics¶
Get service statistics.
Response:
{
"index_type": "hnsw",
"index_size": 100,
"document_count": 100,
"embedding_dimension": 1024,
"model_name": "BAAI/bge-m3"
}
5. List Documents¶
List documents in the store.
Testing¶
Run Demo Client¶
The demo client showcases all features with sample documents:
Run Performance Test¶
Test indexing and search performance with synthetic data:
Manual Testing with curl¶
Index a document:
curl -X POST http://localhost:4240/index \
-H "Content-Type: application/json" \
-d '{"text": "This is a test document about machine learning."}'
Search for similar documents:
curl -X POST http://localhost:4240/search \
-H "Content-Type: application/json" \
-d '{"query": "artificial intelligence", "top_k": 5}'
Index Comparison¶
ANNOY (Approximate Nearest Neighbors Oh Yeah)¶
Pros: - Very fast indexing - Low memory footprint - Good for static datasets - Supports multiple distance metrics
Cons: - Requires rebuild for deletion - No incremental updates after build - Trade-off between speed and accuracy (controlled by n_trees)
Best for: - Large-scale similarity search - Read-heavy workloads - Memory-constrained environments
HNSW (Hierarchical Navigable Small World)¶
Pros: - High recall accuracy - Supports incremental updates - Fast search with good precision - Supports soft deletion
Cons: - Higher memory usage - Slower indexing than ANNOY - More complex parameter tuning
Best for: - Dynamic datasets - High-precision requirements - Balanced read/write workloads
Configuration¶
Environment Variables¶
You can configure the service using environment variables with the VEC_ prefix:
export VEC_INDEX_TYPE=hnsw
export VEC_MODEL_NAME=BAAI/bge-m3
export VEC_USE_FP16=true
export VEC_MAX_SEQ_LENGTH=512
export VEC_MAX_DOCUMENTS=100000
export VEC_LOG_LEVEL=DEBUG
# ANNOY specific
export VEC_ANNOY_N_TREES=50
export VEC_ANNOY_METRIC=angular
# HNSW specific
export VEC_HNSW_EF_CONSTRUCTION=200
export VEC_HNSW_M=32
export VEC_HNSW_EF_SEARCH=100
export VEC_HNSW_SPACE=cosine
Educational Features¶
This service includes extensive logging for educational purposes:
- Embedding Generation Logs: Shows the process of converting text to vectors
- Index Operation Logs: Detailed information about index updates
- Search Process Logs: Step-by-step search execution
- Performance Metrics: Timing information for all operations
- Vector Statistics: Min/max/mean values of embeddings (when enabled)
Enable full educational logging:
Performance Considerations¶
Memory Usage¶
- BGE-M3 model: ~2.3GB
- Per document overhead: ~4KB (1024-dim float32 embedding)
- ANNOY index: ~(4 * dimension * n_items * n_trees / 2) bytes
- HNSW index: ~(4 * dimension * n_items * M * 2) bytes
Optimization Tips¶
- For ANNOY:
- Increase
n_treesfor better accuracy (slower build) - Use
angularmetric for normalized vectors -
Build index after batch insertions
-
For HNSW:
- Increase
Mfor better recall (more memory) - Increase
ef_constructionfor better index quality (slower build) -
Adjust
ef_searchfor speed/accuracy trade-off -
General:
- Use FP16 for faster inference (slight accuracy loss)
- Batch document insertions when possible
- Limit
max_seq_lengthbased on your documents
Troubleshooting¶
Common Issues¶
- Out of Memory:
- Reduce batch size
- Use FP16 mode
- Lower max_seq_length
-
Use ANNOY instead of HNSW
-
Slow Indexing:
- Reduce HNSW ef_construction
- Reduce ANNOY n_trees
-
Use GPU if available
-
Poor Search Quality:
- Increase ANNOY n_trees
- Increase HNSW M and ef_search
- Check if documents are too short/long
References¶
License¶
This is an educational project for learning purposes.
源代码¶
cli.py¶
#!/usr/bin/env python3
"""
稠密检索命令行工具（实验 3-4）
在一个小型示例语料上运行稠密嵌入检索，支持：
- 自定义语料 / 查询 / top-k / 输出文件
- --eval：在带标注的小型评测集上计算 recall@k / precision@k / MRR，
直观展示"稠密嵌入读得懂同义表达"这一核心卖点
- --compare-ann：复现书中实验 3-4 的重点——对比 ANNOY 与 HNSW 两种 ANN 后端
相对精确暴力检索的召回率、建索引耗时与查询延迟（复用服务端 indexing.py）
- --embedding-model：可切换嵌入模型；默认 BAAI/bge-m3，离线可用已缓存的
sentence-transformers/all-MiniLM-L6-v2
不带任何参数运行时，等价于书中实验 3-4 的默认演示（查询 "a cat playing"）。
--compare-ann 使用合成向量、无需任何模型，可在完全离线环境下复现 ANN 对比。
"""
import argparse
import json
import sys
import time
from typing import Dict, List, Optional, Set
import numpy as np
from indexing import AnnoyIndex, HNSWIndex
# ---------------------------------------------------------------------------
# 内置示例语料与标注（英文，与常见句向量模型能力一致，可完全离线复现）
# 语料刻意加入了"同义表达"文档（kitten / feline 表示 cat，distillation 的两种写法），
# 用来展示稠密检索在语义匹配上的强项——这些正是稀疏 BM25（实验 3-5）会漏召回的场景。
# ---------------------------------------------------------------------------
DEFAULT_CORPUS: List[Dict] = [
{"doc_id": "doc_1", "title": "Python Language",
"text": "Python is a high-level programming language known for readability and a simple syntax."},
{"doc_id": "doc_2", "title": "JavaScript Runtime",
"text": "JavaScript runs in the browser and on servers via Node.js for full-stack web development."},
{"doc_id": "doc_3", "title": "Model Distillation",
"text": "Model distillation compresses a large teacher model into a smaller student model while preserving accuracy."},
{"doc_id": "doc_4", "title": "Knowledge Distillation",
"text": "Knowledge distillation transfers knowledge from a big neural network to a compact model for efficient inference."},
{"doc_id": "doc_5", "title": "BM25 Ranking",
"text": "BM25 is a probabilistic ranking function using term frequency and inverse document frequency."},
{"doc_id": "doc_6", "title": "HTTP Errors",
"text": "The HTTP 404 error code means the requested resource was not found on the web server."},
{"doc_id": "doc_7", "title": "A Playful Kitten",
"text": "A cute kitten chased a ball of yarn across the living room floor all afternoon."},
{"doc_id": "doc_8", "title": "Silent Hunter",
"text": "The feline predator stalked its prey silently through the tall grass at dusk."},
{"doc_id": "doc_9", "title": "Hardware Fault",
"text": "Error code XK9-2B4-7Q1 indicates a hardware fault in the storage controller board."},
{"doc_id": "doc_10", "title": "Transformers",
"text": "Transformer models use self-attention to process input sequences in parallel efficiently."},
{"doc_id": "doc_11", "title": "Deep Learning",
"text": "Deep learning stacks many layers of neurons to extract hierarchical features from raw data."},
{"doc_id": "doc_12", "title": "Gradient Descent",
"text": "Gradient descent minimizes a loss function by iteratively updating the model parameters."},
]
# query -> 相关文档 doc_id 集合（人工标注的 ground truth）
# 这些查询大多不与相关文档共享字面关键词，只在语义上相关——考的正是稠密检索的语义能力。
DEFAULT_LABELS: Dict[str, List[str]] = {
# kitten / feline 都不含字面 "cat"，稠密检索应凭语义召回，稀疏 BM25 则会漏
"a cat playing": ["doc_7", "doc_8"],
# "蒸馏"的两种写法，语义同一主题
"model distillation": ["doc_3", "doc_4"],
# 语义相关，字面不含 "neural network training"
"training neural networks": ["doc_11", "doc_12"],
"self attention in sequence models": ["doc_10"],
"web server resource not found": ["doc_6"],
}
DEFAULT_QUERY = "a cat playing"
DEFAULT_MODEL = "BAAI/bge-m3"
OFFLINE_HINT_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
# ---------------------------------------------------------------------------
# 稠密嵌入编码器：用 transformers 的 AutoModel 直接算句向量（mean / cls 池化 + L2 归一化）
# 这样既能加载书中默认的 BAAI/bge-m3（bge 系用 cls 池化），也能加载离线已缓存的
# sentence-transformers/all-MiniLM-L6-v2（mean 池化），无需依赖 FlagEmbedding。
# ---------------------------------------------------------------------------
class DenseEncoder:
def __init__(self, model_name: str, pooling: str = "auto",
device: str = "cpu", max_length: int = 512):
import torch
from transformers import AutoModel, AutoTokenizer
self.torch = torch
self.model_name = model_name
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name)
self.model.eval().to(device)
self.device = device
self.max_length = max_length
if pooling == "auto":
# bge / bge-m3 的稠密向量取 [CLS]；多数 sentence-transformers 模型用平均池化
pooling = "cls" if "bge" in model_name.lower() else "mean"
self.pooling = pooling
def encode(self, texts: List[str], batch_size: int = 16) -> np.ndarray:
vecs: List[np.ndarray] = []
for i in range(0, len(texts), batch_size):
batch = texts[i:i + batch_size]
enc = self.tokenizer(batch, padding=True, truncation=True,
max_length=self.max_length, return_tensors="pt").to(self.device)
with self.torch.no_grad():
out = self.model(**enc)
if self.pooling == "cls":
emb = out.last_hidden_state[:, 0]
else:
mask = enc["attention_mask"].unsqueeze(-1).float()
emb = (out.last_hidden_state * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
emb = self.torch.nn.functional.normalize(emb, p=2, dim=1)
vecs.append(emb.cpu().numpy().astype("float32"))
return np.vstack(vecs)
def load_encoder(model_name: str, pooling: str, device: str) -> Optional["DenseEncoder"]:
"""加载稠密编码器。离线且模型未缓存时给出清晰提示并返回 None（不影响参数解析验证）。"""
try:
import torch # noqa: F401
from transformers import AutoModel # noqa: F401
except Exception as e:
print("\n[稠密编码] 需要依赖 transformers 与 torch，当前环境缺失：", e)
print(" 安装：pip install torch transformers")
print(" （--compare-ann 使用合成向量，无需任何模型，可完全离线运行）")
return None
try:
print(f"正在加载嵌入模型 {model_name}（pooling={pooling}, device={device}）...")
t0 = time.time()
encoder = DenseEncoder(model_name, pooling=pooling, device=device)
print(f"模型加载完成，耗时 {time.time() - t0:.1f}s，池化方式 ={encoder.pooling}")
return encoder
except Exception as e:
print(f"\n[稠密编码] 无法加载模型 {model_name}：{e}")
print(f" 离线环境无法下载 {model_name} 权重（BGE-M3 约 2.3GB）。")
print(f" 可改用已缓存的小模型：--embedding-model {OFFLINE_HINT_MODEL}")
print(" 或先在联网环境预缓存目标模型；--compare-ann 则完全无需模型。")
return None
def load_corpus(path: Optional[str]) -> List[Dict]:
"""加载语料。支持 .json（文档数组）与 .jsonl（每行一个文档）。"""
if not path:
return DEFAULT_CORPUS
docs: List[Dict] = []
with open(path, "r", encoding="utf-8") as f:
if path.endswith(".jsonl"):
for line in f:
line = line.strip()
if line:
docs.append(json.loads(line))
else:
data = json.load(f)
docs = data["documents"] if isinstance(data, dict) else data
if not docs:
raise ValueError(f"语料文件为空：{path}")
return docs
def load_labels(path: Optional[str]) -> Dict[str, List[str]]:
"""加载评测标注：{query: [relevant_doc_id, ...]}。"""
if not path:
return DEFAULT_LABELS
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
# ---------------------------------------------------------------------------
# 稠密检索（精确暴力，用于单条查询与检索质量评测）
# ---------------------------------------------------------------------------
def dense_rank(query_vec: np.ndarray, doc_matrix: np.ndarray) -> List[int]:
"""向量已 L2 归一化，余弦相似度即点积；返回按相似度降序的文档下标。"""
sims = doc_matrix @ query_vec
return list(np.argsort(-sims)), sims
def run_search(encoder: "DenseEncoder", corpus: List[Dict], doc_matrix: np.ndarray,
query: str, top_k: int) -> List[Dict]:
"""执行单条稠密查询并打印结果，返回结构化结果供 --output 落盘。"""
q = encoder.encode([query])[0]
order, sims = dense_rank(q, doc_matrix)
print(f"\n查询: '{query}' (稠密检索, top-{top_k})")
print("-" * 60)
out = []
for rank, idx in enumerate(order[:top_k], 1):
d = corpus[idx]
title = d.get("title", "")
print(f" #{rank} {d.get('doc_id')} cos={float(sims[idx]):.4f} {title}")
print(f" 预览: {d['text'][:80]}...")
out.append({
"rank": rank,
"doc_id": d.get("doc_id"),
"score": float(sims[idx]),
"title": title,
})
return out
def _metrics_for_query(retrieved: List[str], relevant: Set[str], k: int) -> Dict:
"""单条查询的 recall@k / precision@k / 命中排名（用于 MRR）。"""
topk = retrieved[:k]
hits = [d for d in topk if d in relevant]
recall = len(set(hits)) / len(relevant) if relevant else 0.0
precision = len(hits) / len(topk) if topk else 0.0
rr = 0.0
for i, d in enumerate(retrieved, 1):
if d in relevant:
rr = 1.0 / i
break
return {"recall": recall, "precision": precision, "rr": rr,
"hits": hits, "retrieved": topk}
def run_eval(encoder: "DenseEncoder", corpus: List[Dict], doc_matrix: np.ndarray,
labels: Dict[str, List[str]], k: int) -> Dict:
"""在标注集上做稠密检索评测，打印每条查询指标 + 宏平均。"""
doc_ids = [d.get("doc_id") for d in corpus]
print(f"\n{'=' * 60}")
print(f"稠密检索质量评测 (recall@{k} / precision@{k} / MRR)")
print(f"{'=' * 60}")
per_query = {}
sum_recall = sum_prec = sum_rr = 0.0
q_vecs = encoder.encode(list(labels.keys()))
for (query, rel_list), qv in zip(labels.items(), q_vecs):
relevant = set(rel_list)
order, _ = dense_rank(qv, doc_matrix)
retrieved = [doc_ids[i] for i in order]
m = _metrics_for_query(retrieved, relevant, k)
per_query[query] = m
sum_recall += m["recall"]
sum_prec += m["precision"]
sum_rr += m["rr"]
flag = "" if m["recall"] > 0 else " <- 漏召回"
print(f"\n查询 '{query}' 相关文档={sorted(relevant)}")
print(f" 召回排序: {retrieved[:k]}")
print(f" recall@{k}={m['recall']:.2f} precision@{k}={m['precision']:.2f} RR={m['rr']:.2f}{flag}")
n = len(labels)
macro = {
"recall@k": sum_recall / n,
"precision@k": sum_prec / n,
"mrr": sum_rr / n,
"miss_rate@k": 1.0 - sum_recall / n,
}
print(f"\n{'-' * 60}")
print(f"宏平均 recall@{k}={macro['recall@k']:.3f} "
f"precision@{k}={macro['precision@k']:.3f} "
f"MRR={macro['mrr']:.3f} 漏召回率(1-recall@{k})={macro['miss_rate@k']:.3f}")
return {"k": k, "per_query": {q: {kk: vv for kk, vv in m.items() if kk != "retrieved"}
for q, m in per_query.items()},
"macro": macro}
# ---------------------------------------------------------------------------
# ANN 后端对比（实验 3-4 的重点）：复用服务端 indexing.py 里的 ANNOY / HNSW 实现，
# 在一批合成单位向量上对比二者相对"精确暴力检索"的召回率、建索引耗时与查询延迟。
# 用合成向量而非真实文本嵌入，是为了 (a) 完全离线、无需下载模型；(b) 语料足够大时
# ANN 的"近似"才会显现出与精确检索的差距，从而看清两类算法的取舍。
# ---------------------------------------------------------------------------
def _exact_topk(queries: np.ndarray, base: np.ndarray, k: int) -> List[Set[int]]:
"""精确暴力最近邻（余弦），作为 ANN 召回率的 ground truth。"""
sims = queries @ base.T
idx = np.argsort(-sims, axis=1)[:, :k]
return [set(row.tolist()) for row in idx]
def _sanity_ok(index, base: np.ndarray) -> bool:
"""自检：用库中已存在的向量查询，应能召回它自己。用于识别环境中损坏的索引后端。"""
probe = min(5, len(base))
for i in range(probe):
ids, _ = index.search(base[i], min(10, len(base)))
if f"v{i}" not in set(ids):
return False
return True
def compare_ann(base: np.ndarray, queries: np.ndarray, top_k: int, backends: List[str],
annoy_n_trees: int, hnsw_M: int, hnsw_ef_search: int,
hnsw_ef_construction: int) -> Dict:
dim = base.shape[1]
n = len(base)
exact_sets = _exact_topk(queries, base, top_k)
print(f"\n{'=' * 60}")
print(f"ANN 后端对比：{n} 条 {dim} 维向量，{len(queries)} 条查询，top-{top_k}")
print(f"指标：recall@{top_k} 相对精确暴力检索 / 建索引耗时 / 平均查询延迟")
print(f"{'=' * 60}")
report: Dict[str, Dict] = {}
for backend in backends:
if backend == "annoy":
index = AnnoyIndex(dimension=dim, n_trees=annoy_n_trees,
metric="angular", logger=None)
else:
index = HNSWIndex(dimension=dim, max_elements=n + 16,
ef_construction=hnsw_ef_construction, M=hnsw_M,
ef_search=hnsw_ef_search, space="cosine", logger=None)
t0 = time.time()
for i, v in enumerate(base):
index.add_item(f"v{i}", v)
if backend == "annoy":
index.rebuild_index()
build_time = time.time() - t0
healthy = _sanity_ok(index, base)
recalls: List[float] = []
qtimes: List[float] = []
for qi, q in enumerate(queries):
ts = time.time()
ids, _ = index.search(q, top_k)
qtimes.append(time.time() - ts)
got = {int(d[1:]) for d in ids}
recalls.append(len(got & exact_sets[qi]) / top_k)
mean_recall = float(np.mean(recalls))
mean_qms = float(np.mean(qtimes) * 1000)
params = (f"n_trees={annoy_n_trees}" if backend == "annoy"
else f"M={hnsw_M}, ef_search={hnsw_ef_search}, ef_construction={hnsw_ef_construction}")
report[backend] = {
"recall@k": mean_recall,
"build_time_s": build_time,
"mean_query_ms": mean_qms,
"params": params,
"healthy": healthy,
}
warn = "" if healthy else " [警告] 该后端连自身向量都召回不到，疑似当前环境下损坏，下列数字不可信"
print(f"\n[{backend.upper()}] {params}{warn}")
print(f" recall@{top_k} = {mean_recall:.3f}")
print(f" 建索引耗时 = {build_time * 1000:.1f} ms")
print(f" 平均查询延迟 = {mean_qms:.3f} ms")
if "annoy" in report and "hnsw" in report:
print(f"\n{'-' * 60}")
print("小结：HNSW 图结构通常召回率更高、支持增量插入，代价是更高内存与建索引开销；")
print(" ANNOY 树结构建索引快、内存省，但删除需重建，召回随 n_trees 调节。")
return report
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="cli.py",
description="稠密检索命令行工具（实验 3-4）：在小型语料上运行稠密嵌入检索并评测检索质量，"
"并对比 ANNOY / HNSW 两种 ANN 索引后端。",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""示例：
python cli.py # 默认演示（查询 "a cat playing"，需嵌入模型）
python cli.py -q "model distillation" -k 3 # 单条稠密查询
python cli.py --eval # 在标注集上算 recall/precision/MRR
python cli.py --embedding-model sentence-transformers/all-MiniLM-L6-v2 --eval # 离线小模型
python cli.py --compare-ann # ANNOY vs HNSW 召回率对比（合成向量，无需模型）
python cli.py --compare-ann --ann-base 5000 --annoy-n-trees 5 -k 10 -o ann.json
""",
)
parser.add_argument("-q", "--query", default=DEFAULT_QUERY,
help=f"查询字符串（默认: '{DEFAULT_QUERY}'）")
parser.add_argument("-c", "--corpus", default=None,
help="语料文件路径（.json 文档数组 或 .jsonl 每行一篇）；缺省用内置示例语料")
parser.add_argument("-k", "--top-k", type=int, default=5,
help="返回前 k 条结果（默认: 5）")
parser.add_argument("-o", "--output", default=None,
help="把结果/评测指标以 JSON 写入该文件")
parser.add_argument("--embedding-model", default=DEFAULT_MODEL,
help=f"稠密嵌入模型名（默认: {DEFAULT_MODEL}）；"
f"离线可用已缓存的 {OFFLINE_HINT_MODEL}")
parser.add_argument("--pooling", choices=["auto", "mean", "cls"], default="auto",
help="句向量池化方式：auto(bge*用cls，其余用mean) / mean / cls")
parser.add_argument("--device", default="cpu",
help="推理设备（cpu / cuda / mps，默认: cpu）")
parser.add_argument("--eval", action="store_true",
help="在标注集上评测 recall@k / precision@k / MRR，而非只跑单条查询")
parser.add_argument("--labels", default=None,
help="评测标注文件 {query: [相关doc_id,...]}；缺省用内置标注")
ann = parser.add_argument_group("ANN 后端对比（--compare-ann）")
ann.add_argument("--compare-ann", action="store_true",
help="对比 ANNOY 与 HNSW 的召回率/耗时（复用 indexing.py，用合成向量，无需模型）")
ann.add_argument("--backend", choices=["annoy", "hnsw", "both"], default="both",
help="参与对比的 ANN 后端（默认: both）")
ann.add_argument("--ann-base", type=int, default=3000,
help="合成底库向量数量（默认: 3000，越大 ANN 近似误差越明显）")
ann.add_argument("--ann-queries", type=int, default=100,
help="合成查询向量数量（默认: 100）")
ann.add_argument("--ann-dim", type=int, default=128,
help="合成向量维度（默认: 128）")
ann.add_argument("--annoy-n-trees", type=int, default=10,
help="ANNOY 树数量（默认: 10；越多越准越慢）")
ann.add_argument("--hnsw-M", type=int, default=16,
help="HNSW 每节点连接数 M（默认: 16；越大召回越高越占内存）")
ann.add_argument("--hnsw-ef-search", type=int, default=20,
help="HNSW 查询期动态候选表大小 ef_search（默认: 20）")
ann.add_argument("--hnsw-ef-construction", type=int, default=100,
help="HNSW 建索引期动态候选表大小 ef_construction（默认: 100）")
ann.add_argument("--seed", type=int, default=42,
help="合成向量随机种子（默认: 42）")
return parser
def main(argv: Optional[List[str]] = None) -> int:
args = build_parser().parse_args(argv)
payload: Dict = {"top_k": args.top_k}
# --- ANN 后端对比：合成向量，无需嵌入模型，完全离线 ---
if args.compare_ann:
rng = np.random.default_rng(args.seed)
base = rng.standard_normal((args.ann_base, args.ann_dim)).astype("float32")
base /= np.linalg.norm(base, axis=1, keepdims=True)
queries = rng.standard_normal((args.ann_queries, args.ann_dim)).astype("float32")
queries /= np.linalg.norm(queries, axis=1, keepdims=True)
backends = ["annoy", "hnsw"] if args.backend == "both" else [args.backend]
payload["compare_ann"] = compare_ann(
base, queries, args.top_k, backends,
annoy_n_trees=args.annoy_n_trees, hnsw_M=args.hnsw_M,
hnsw_ef_search=args.hnsw_ef_search, hnsw_ef_construction=args.hnsw_ef_construction)
if args.output:
with open(args.output, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
print(f"\n已写入结果：{args.output}")
return 0
# --- 稠密检索 / 评测：需要嵌入模型 ---
corpus = load_corpus(args.corpus)
print(f"已加载语料：{len(corpus)} 篇文档"
+ ("（内置示例）" if not args.corpus else f"（来自 {args.corpus}）"))
encoder = load_encoder(args.embedding_model, args.pooling, args.device)
if encoder is None:
return 0 # 已给出模型缺失提示，视为正常退出（参数解析已验证）
doc_matrix = encoder.encode([d["text"] for d in corpus])
payload["embedding_model"] = args.embedding_model
payload["query"] = args.query
if args.eval:
labels = load_labels(args.labels)
payload["eval"] = run_eval(encoder, corpus, doc_matrix, labels, args.top_k)
else:
payload["results"] = run_search(encoder, corpus, doc_matrix, args.query, args.top_k)
if args.output:
with open(args.output, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
print(f"\n已写入结果：{args.output}")
return 0
if __name__ == "__main__":
sys.exit(main())
config.py¶
"""Configuration for the dense embedding service."""
import os
from enum import Enum
from dataclasses import dataclass, field
from typing import Optional
class IndexType(Enum):
"""Supported index types."""
ANNOY = "annoy"
HNSW = "hnsw"
@dataclass
class ServiceConfig:
"""Service configuration."""
# Server settings
host: str = "0.0.0.0"
port: int = 4240 # Default port for dense embedding service
# Model settings
model_name: str = "BAAI/bge-m3"
use_fp16: bool = True
max_seq_length: int = 8192 # Increased to match HARD_LIMIT in chunking
# Index settings
index_type: IndexType = IndexType.HNSW
max_documents: int = 100000
# HNSW specific settings
hnsw_ef_construction: int = 200
hnsw_M: int = 16
hnsw_ef_search: int = 50
hnsw_space: str = "cosine"
# Annoy specific settings
annoy_n_trees: int = 50
annoy_metric: str = "angular"
# Logging settings
log_level: str = "INFO"
debug: bool = False
show_embeddings: bool = False
@classmethod
def from_env(cls):
"""Create config from environment variables."""
config = cls()
# Override with environment variables if present
if os.getenv("DENSE_PORT"):
config.port = int(os.getenv("DENSE_PORT"))
if os.getenv("DENSE_HOST"):
config.host = os.getenv("DENSE_HOST")
if os.getenv("DENSE_MODEL"):
config.model_name = os.getenv("DENSE_MODEL")
if os.getenv("DEBUG"):
config.debug = os.getenv("DEBUG").lower() == "true"
return config
document_store.py¶
"""In-memory document store for managing documents."""
from typing import Dict, Optional, List
from dataclasses import dataclass, field
from datetime import datetime
import uuid
from logger import VectorSearchLogger
@dataclass
class Document:
"""Document data class."""
id: str
text: str
metadata: Dict = field(default_factory=dict)
created_at: datetime = field(default_factory=datetime.now)
embedding: Optional[List[float]] = None
class DocumentStore:
"""In-memory document storage."""
def __init__(self, logger: Optional[VectorSearchLogger] = None):
"""
Initialize the document store.
Args:
logger: Logger instance for educational output
"""
self.documents: Dict[str, Document] = {}
self.logger = logger
if self.logger:
self.logger.logger.info("📦 Initialized in-memory document store")
def add_document(self, text: str, doc_id: Optional[str] = None,
metadata: Optional[Dict] = None) -> str:
"""
Add a document to the store.
Args:
text: Document text
doc_id: Optional document ID (will be generated if not provided)
metadata: Optional metadata dictionary
Returns:
Document ID
"""
# Generate ID if not provided
if doc_id is None:
doc_id = str(uuid.uuid4())
# Check if document already exists
if doc_id in self.documents:
if self.logger:
self.logger.logger.warning(f"Document {doc_id} already exists, updating...")
# Create document
doc = Document(
id=doc_id,
text=text,
metadata=metadata or {}
)
# Store document
self.documents[doc_id] = doc
if self.logger:
self.logger.logger.debug(f"📄 Stored document")
self.logger.logger.debug(f" - ID: {doc_id}")
self.logger.logger.debug(f" - Text length: {len(text)} chars")
self.logger.logger.debug(f" - Metadata keys: {list(metadata.keys()) if metadata else []}")
self.logger.logger.debug(f" - Total documents: {len(self.documents)}")
return doc_id
def get_document(self, doc_id: str) -> Optional[Document]:
"""
Retrieve a document by ID.
Args:
doc_id: Document ID
Returns:
Document or None if not found
"""
doc = self.documents.get(doc_id)
if self.logger:
if doc:
self.logger.logger.debug(f"✅ Retrieved document {doc_id}")
else:
self.logger.logger.warning(f"❌ Document {doc_id} not found")
return doc
def delete_document(self, doc_id: str) -> bool:
"""
Delete a document from the store.
Args:
doc_id: Document ID
Returns:
True if deleted, False if not found
"""
if doc_id in self.documents:
del self.documents[doc_id]
if self.logger:
self.logger.logger.debug(f"🗑️ Deleted document {doc_id}")
self.logger.logger.debug(f" Remaining documents: {len(self.documents)}")
return True
if self.logger:
self.logger.logger.warning(f"Document {doc_id} not found for deletion")
return False
def list_documents(self, limit: Optional[int] = None) -> List[Document]:
"""
List all documents in the store.
Args:
limit: Maximum number of documents to return
Returns:
List of documents
"""
docs = list(self.documents.values())
if limit:
docs = docs[:limit]
if self.logger:
self.logger.logger.debug(f"📋 Listing {len(docs)} documents")
return docs
def get_size(self) -> int:
"""Get the number of documents in the store."""
return len(self.documents)
def clear(self) -> None:
"""Clear all documents from the store."""
count = len(self.documents)
self.documents.clear()
if self.logger:
self.logger.logger.info(f"🧹 Cleared {count} documents from store")
def get_documents_by_ids(self, doc_ids: List[str]) -> List[Document]:
"""
Retrieve multiple documents by their IDs.
Args:
doc_ids: List of document IDs
Returns:
List of documents (only those found)
"""
docs = []
for doc_id in doc_ids:
doc = self.documents.get(doc_id)
if doc:
docs.append(doc)
if self.logger:
self.logger.logger.debug(f"Retrieved {len(docs)}/{len(doc_ids)} documents")
return docs
def update_document_embedding(self, doc_id: str, embedding: List[float]) -> bool:
"""
Update the embedding for a document.
Args:
doc_id: Document ID
embedding: Embedding vector
Returns:
True if updated, False if document not found
"""
if doc_id in self.documents:
self.documents[doc_id].embedding = embedding
if self.logger:
self.logger.logger.debug(f"Updated embedding for document {doc_id}")
return True
return False
embedding_service.py¶
"""Embedding service using BGE-M3 model."""
import time
import numpy as np
from typing import List, Dict, Optional
from FlagEmbedding import BGEM3FlagModel
from logger import VectorSearchLogger, log_execution_time
import logging
class EmbeddingService:
"""Service for generating embeddings using BGE-M3 model."""
def __init__(self, model_name: str = "BAAI/bge-m3", use_fp16: bool = True,
max_seq_length: int = 512, logger: Optional[VectorSearchLogger] = None):
"""
Initialize the embedding service with BGE-M3 model.
Args:
model_name: Name of the BGE-M3 model
use_fp16: Whether to use FP16 for inference
max_seq_length: Maximum sequence length
logger: Logger instance for educational output
"""
self.model_name = model_name
self.use_fp16 = use_fp16
self.max_seq_length = max_seq_length
self.logger = logger
self.std_logger = logging.getLogger("vector_search")
# Initialize the model
self._initialize_model()
def _initialize_model(self):
"""Initialize the BGE-M3 model."""
start_time = time.time()
if self.logger:
self.logger.logger.info(f"🚀 Initializing BGE-M3 model: {self.model_name}")
self.logger.logger.debug(f" - Using FP16: {self.use_fp16}")
self.logger.logger.debug(f" - Max sequence length: {self.max_seq_length}")
try:
self.model = BGEM3FlagModel(
self.model_name,
use_fp16=self.use_fp16
)
# Get embedding dimension by encoding a test sentence
test_embedding = self.model.encode(["test"])
if isinstance(test_embedding, dict):
self.embedding_dim = test_embedding['dense_vecs'].shape[1]
else:
self.embedding_dim = test_embedding.shape[1]
load_time = time.time() - start_time
if self.logger:
self.logger.logger.info(f"✅ Model loaded successfully in {load_time:.2f} seconds")
self.logger.logger.debug(f" - Embedding dimension: {self.embedding_dim}")
self.logger.logger.debug(f" - Model supports: dense, sparse, and multi-vector retrieval")
except Exception as e:
if self.logger:
self.logger.logger.error(f"Failed to load model: {e}")
raise
@log_execution_time()
def encode_text(self, text: str, return_sparse: bool = False,
return_colbert: bool = False) -> Dict[str, np.ndarray]:
"""
Encode a single text into embeddings.
Args:
text: Input text to encode
return_sparse: Whether to return sparse embeddings
return_colbert: Whether to return ColBERT embeddings
Returns:
Dictionary containing different types of embeddings
"""
start_time = time.time()
if self.logger:
self.logger.logger.debug(f"📝 Encoding text (length: {len(text)} chars)")
self.logger.logger.debug(f" Text preview: {text[:100]}..." if len(text) > 100 else f" Text: {text}")
# Encode the text
embeddings = self.model.encode(
[text],
return_dense=True,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert
)
# Extract dense embeddings
dense_vec = embeddings['dense_vecs'][0]
result = {
'dense': dense_vec,
'dimension': len(dense_vec)
}
# Add sparse embeddings if requested
if return_sparse and 'lexical_weights' in embeddings:
result['sparse'] = embeddings['lexical_weights'][0]
if self.logger:
num_tokens = len(result['sparse'])
self.logger.logger.debug(f" Sparse embedding: {num_tokens} non-zero tokens")
# Add ColBERT embeddings if requested
if return_colbert and 'colbert_vecs' in embeddings:
result['colbert'] = embeddings['colbert_vecs'][0]
if self.logger:
colbert_shape = result['colbert'].shape
self.logger.logger.debug(f" ColBERT embedding shape: {colbert_shape}")
encoding_time = time.time() - start_time
if self.logger:
self.logger.logger.debug(f"✅ Encoding completed in {encoding_time:.4f} seconds")
self.logger.log_embedding_vector(dense_vec, sample_size=10)
return result
@log_execution_time()
def encode_batch(self, texts: List[str], return_sparse: bool = False,
return_colbert: bool = False) -> Dict[str, np.ndarray]:
"""
Encode multiple texts into embeddings.
Args:
texts: List of input texts to encode
return_sparse: Whether to return sparse embeddings
return_colbert: Whether to return ColBERT embeddings
Returns:
Dictionary containing different types of embeddings for all texts
"""
start_time = time.time()
if self.logger:
self.logger.logger.info(f"📚 Batch encoding {len(texts)} texts")
total_chars = sum(len(t) for t in texts)
self.logger.logger.debug(f" Total characters: {total_chars}")
self.logger.logger.debug(f" Average text length: {total_chars/len(texts):.1f} chars")
# Encode all texts
embeddings = self.model.encode(
texts,
return_dense=True,
return_sparse=return_sparse,
return_colbert_vecs=return_colbert
)
result = {
'dense': embeddings['dense_vecs'],
'dimension': embeddings['dense_vecs'].shape[1],
'num_texts': len(texts)
}
# Add sparse embeddings if requested
if return_sparse and 'lexical_weights' in embeddings:
result['sparse'] = embeddings['lexical_weights']
# Add ColBERT embeddings if requested
if return_colbert and 'colbert_vecs' in embeddings:
result['colbert'] = embeddings['colbert_vecs']
encoding_time = time.time() - start_time
if self.logger:
self.logger.logger.info(f"✅ Batch encoding completed in {encoding_time:.4f} seconds")
self.logger.logger.debug(f" Average time per text: {encoding_time/len(texts):.4f} seconds")
return result
def get_embedding_dimension(self) -> int:
"""Get the dimension of the embeddings."""
return self.embedding_dim
def compute_similarity(self, vec1: np.ndarray, vec2: np.ndarray,
metric: str = "cosine") -> float:
"""
Compute similarity between two vectors.
Args:
vec1: First vector
vec2: Second vector
metric: Similarity metric ('cosine', 'euclidean', 'dot')
Returns:
Similarity score
"""
if metric == "cosine":
# Cosine similarity
dot_product = np.dot(vec1, vec2)
norm1 = np.linalg.norm(vec1)
norm2 = np.linalg.norm(vec2)
similarity = dot_product / (norm1 * norm2)
elif metric == "euclidean":
# Euclidean distance (negative for similarity)
similarity = -np.linalg.norm(vec1 - vec2)
elif metric == "dot":
# Dot product
similarity = np.dot(vec1, vec2)
else:
raise ValueError(f"Unknown metric: {metric}")
if self.logger:
self.logger.logger.debug(f" Similarity ({metric}): {similarity:.6f}")
return float(similarity)
indexing.py¶
"""Vector index implementations using ANNOY and HNSW."""
from abc import ABC, abstractmethod
from typing import List, Tuple, Dict, Optional
import numpy as np
import annoy
import hnswlib
import time
from logger import VectorSearchLogger
class VectorIndex(ABC):
"""Abstract base class for vector indexes."""
@abstractmethod
def add_item(self, doc_id: str, vector: np.ndarray) -> None:
"""Add an item to the index."""
pass
@abstractmethod
def delete_item(self, doc_id: str) -> bool:
"""Delete an item from the index."""
pass
@abstractmethod
def search(self, query_vector: np.ndarray, top_k: int) -> Tuple[List[str], List[float]]:
"""Search for top-k similar items."""
pass
@abstractmethod
def get_size(self) -> int:
"""Get the current number of items in the index."""
pass
@abstractmethod
def rebuild_index(self) -> None:
"""Rebuild the index if necessary."""
pass
class AnnoyIndex(VectorIndex):
"""ANNOY-based vector index implementation."""
def __init__(self, dimension: int, n_trees: int = 50, metric: str = "angular",
logger: Optional[VectorSearchLogger] = None):
"""
Initialize ANNOY index.
Args:
dimension: Dimension of vectors
n_trees: Number of trees for ANNOY (affects precision/speed tradeoff)
metric: Distance metric ('angular', 'euclidean', 'manhattan', 'hamming', 'dot')
logger: Logger instance for educational output
"""
self.dimension = dimension
self.n_trees = n_trees
self.metric = metric
self.logger = logger
# Create index
self.index = annoy.AnnoyIndex(dimension, metric)
# Mapping between internal indices and document IDs
self.id_to_index: Dict[str, int] = {}
self.index_to_id: Dict[int, str] = {}
self.vectors_cache: Dict[int, np.ndarray] = {}
self.next_index = 0
self.is_built = False
if self.logger:
self.logger.logger.info(f"📚 Initialized ANNOY index")
self.logger.logger.debug(f" - Dimension: {dimension}")
self.logger.logger.debug(f" - Number of trees: {n_trees}")
self.logger.logger.debug(f" - Metric: {metric}")
def add_item(self, doc_id: str, vector: np.ndarray) -> None:
"""Add an item to the ANNOY index."""
start_time = time.time()
if doc_id in self.id_to_index:
if self.logger:
self.logger.logger.warning(f"Document {doc_id} already exists in index, updating...")
# Remove old entry
old_index = self.id_to_index[doc_id]
del self.index_to_id[old_index]
del self.vectors_cache[old_index]
# Add to index
current_index = self.next_index
self.index.add_item(current_index, vector.tolist())
# Update mappings
self.id_to_index[doc_id] = current_index
self.index_to_id[current_index] = doc_id
self.vectors_cache[current_index] = vector.copy()
self.next_index += 1
# Mark index as needing rebuild
self.is_built = False
if self.logger:
time_taken = time.time() - start_time
self.logger.logger.debug(f"✅ Added document to ANNOY index in {time_taken:.4f}s")
self.logger.logger.debug(f" - Document ID: {doc_id}")
self.logger.logger.debug(f" - Internal index: {current_index}")
self.logger.logger.debug(f" - Index needs rebuild: True")
def delete_item(self, doc_id: str) -> bool:
"""
Delete an item from the index.
Note: ANNOY doesn't support deletion, so we need to rebuild without the item.
"""
if doc_id not in self.id_to_index:
if self.logger:
self.logger.logger.warning(f"Document {doc_id} not found in index")
return False
if self.logger:
self.logger.logger.info(f"🗑️ Deleting from ANNOY index (requires rebuild)")
# Remove from mappings
old_index = self.id_to_index[doc_id]
del self.id_to_index[doc_id]
del self.index_to_id[old_index]
del self.vectors_cache[old_index]
# Rebuild index without the deleted item
self._rebuild_without_deleted()
if self.logger:
self.logger.logger.debug(f"✅ Document {doc_id} deleted and index rebuilt")
return True
def _rebuild_without_deleted(self):
"""Rebuild the index without deleted items."""
start_time = time.time()
# Create new index
new_index = annoy.AnnoyIndex(self.dimension, self.metric)
# Create new mappings
new_id_to_index = {}
new_index_to_id = {}
new_vectors_cache = {}
# Add all remaining items to new index
new_idx = 0
for old_idx, doc_id in self.index_to_id.items():
if old_idx in self.vectors_cache:
vector = self.vectors_cache[old_idx]
new_index.add_item(new_idx, vector.tolist())
new_id_to_index[doc_id] = new_idx
new_index_to_id[new_idx] = doc_id
new_vectors_cache[new_idx] = vector
new_idx += 1
# Build the new index
new_index.build(self.n_trees)
# Replace old index with new one
self.index = new_index
self.id_to_index = new_id_to_index
self.index_to_id = new_index_to_id
self.vectors_cache = new_vectors_cache
self.next_index = new_idx
self.is_built = True
if self.logger:
time_taken = time.time() - start_time
self.logger.logger.debug(f" Rebuild completed in {time_taken:.4f}s")
self.logger.logger.debug(f" New index size: {len(self.id_to_index)} documents")
def search(self, query_vector: np.ndarray, top_k: int) -> Tuple[List[str], List[float]]:
"""Search for top-k similar items in the ANNOY index."""
# Build index if needed
if not self.is_built:
self.rebuild_index()
start_time = time.time()
# Ensure we don't request more items than we have
actual_k = min(top_k, len(self.index_to_id))
if actual_k == 0:
if self.logger:
self.logger.logger.warning("Index is empty, returning no results")
return [], []
# Search in index
indices, distances = self.index.get_nns_by_vector(
query_vector.tolist(), actual_k, include_distances=True
)
# Convert indices to document IDs
doc_ids = [self.index_to_id[idx] for idx in indices if idx in self.index_to_id]
valid_distances = distances[:len(doc_ids)]
if self.logger:
time_taken = time.time() - start_time
self.logger.logger.debug(f"⚡ ANNOY search completed in {time_taken:.4f}s")
self.logger.logger.debug(f" Retrieved {len(doc_ids)} results")
return doc_ids, valid_distances
def get_size(self) -> int:
"""Get the current number of items in the index."""
return len(self.id_to_index)
def rebuild_index(self) -> None:
"""Build/rebuild the ANNOY index."""
if self.is_built and self.logger:
self.logger.logger.debug("Index already built, skipping rebuild")
return
start_time = time.time()
if self.logger:
self.logger.logger.info(f"🏗️ Building ANNOY index with {self.n_trees} trees")
self.index.build(self.n_trees)
self.is_built = True
if self.logger:
time_taken = time.time() - start_time
self.logger.logger.debug(f"✅ Index built in {time_taken:.4f}s")
class HNSWIndex(VectorIndex):
"""HNSW-based vector index implementation."""
def __init__(self, dimension: int, max_elements: int = 100000,
ef_construction: int = 200, M: int = 32, ef_search: int = 100,
space: str = "cosine", logger: Optional[VectorSearchLogger] = None):
"""
Initialize HNSW index.
Args:
dimension: Dimension of vectors
max_elements: Maximum number of elements
ef_construction: Size of the dynamic list (affects build time/accuracy)
M: Number of bi-directional links (affects memory/accuracy)
ef_search: Size of the dynamic list for search (affects search time/accuracy)
space: Distance metric ('l2', 'ip', 'cosine')
logger: Logger instance for educational output
"""
self.dimension = dimension
self.max_elements = max_elements
self.ef_construction = ef_construction
self.M = M
self.ef_search = ef_search
self.space = space
self.logger = logger
# Create index
self.index = hnswlib.Index(space=space, dim=dimension)
self.index.init_index(max_elements=max_elements, ef_construction=ef_construction, M=M)
self.index.set_ef(ef_search)
# Mapping between document IDs and internal labels
self.id_to_label: Dict[str, int] = {}
self.label_to_id: Dict[int, str] = {}
self.available_labels: List[int] = []
self.next_label = 0
if self.logger:
self.logger.logger.info(f"📚 Initialized HNSW index")
self.logger.logger.debug(f" - Dimension: {dimension}")
self.logger.logger.debug(f" - Max elements: {max_elements}")
self.logger.logger.debug(f" - ef_construction: {ef_construction}")
self.logger.logger.debug(f" - M: {M}")
self.logger.logger.debug(f" - ef_search: {ef_search}")
self.logger.logger.debug(f" - Space: {space}")
def add_item(self, doc_id: str, vector: np.ndarray) -> None:
"""Add an item to the HNSW index."""
start_time = time.time()
# Check if document already exists
if doc_id in self.id_to_label:
if self.logger:
self.logger.logger.warning(f"Document {doc_id} already exists, updating...")
# Remove old entry first
self.delete_item(doc_id)
# Get a label for this document
if self.available_labels:
label = self.available_labels.pop()
else:
label = self.next_label
self.next_label += 1
# Add to index
self.index.add_items(vector.reshape(1, -1), np.array([label]))
# Update mappings
self.id_to_label[doc_id] = label
self.label_to_id[label] = doc_id
if self.logger:
time_taken = time.time() - start_time
self.logger.logger.debug(f"✅ Added document to HNSW index in {time_taken:.4f}s")
self.logger.logger.debug(f" - Document ID: {doc_id}")
self.logger.logger.debug(f" - Internal label: {label}")
self.logger.logger.debug(f" - Current index size: {self.index.get_current_count()}")
def delete_item(self, doc_id: str) -> bool:
"""Delete an item from the HNSW index."""
if doc_id not in self.id_to_label:
if self.logger:
self.logger.logger.warning(f"Document {doc_id} not found in index")
return False
start_time = time.time()
# Get label and mark for deletion
label = self.id_to_label[doc_id]
try:
# Mark as deleted in HNSW (soft delete)
self.index.mark_deleted(label)
# Update mappings
del self.id_to_label[doc_id]
del self.label_to_id[label]
# Add label back to available labels for reuse
self.available_labels.append(label)
if self.logger:
time_taken = time.time() - start_time
self.logger.logger.debug(f"✅ Deleted document from HNSW index in {time_taken:.4f}s")
self.logger.logger.debug(f" - Document ID: {doc_id}")
self.logger.logger.debug(f" - Internal label: {label} (marked for reuse)")
return True
except Exception as e:
if self.logger:
self.logger.logger.error(f"Error deleting document: {e}")
return False
def search(self, query_vector: np.ndarray, top_k: int) -> Tuple[List[str], List[float]]:
"""Search for top-k similar items in the HNSW index."""
start_time = time.time()
# Ensure we don't request more items than we have
actual_k = min(top_k, len(self.label_to_id))
if actual_k == 0:
if self.logger:
self.logger.logger.warning("Index is empty, returning no results")
return [], []
# Search in index
labels, distances = self.index.knn_query(query_vector.reshape(1, -1), k=actual_k)
# Convert labels to document IDs
doc_ids = []
valid_distances = []
for label, distance in zip(labels[0], distances[0]):
if label in self.label_to_id:
doc_ids.append(self.label_to_id[label])
valid_distances.append(float(distance))
if self.logger:
time_taken = time.time() - start_time
self.logger.logger.debug(f"⚡ HNSW search completed in {time_taken:.4f}s")
self.logger.logger.debug(f" Retrieved {len(doc_ids)} results")
self.logger.logger.debug(f" Search ef parameter: {self.ef_search}")
return doc_ids, valid_distances
def get_size(self) -> int:
"""Get the current number of items in the index."""
return len(self.id_to_label)
def rebuild_index(self) -> None:
"""HNSW doesn't require explicit rebuild."""
if self.logger:
self.logger.logger.debug("HNSW index doesn't require explicit rebuild")
pass
logger.py¶
"""Educational logging configuration with extensive debug information."""
import logging
import sys
import time
from typing import Optional
import colorlog
from functools import wraps
def setup_logger(name: str = "vector_search", level: str = "DEBUG") -> logging.Logger:
"""
Set up a colorful and informative logger for educational purposes.
Args:
name: Logger name
level: Logging level (DEBUG, INFO, WARNING, ERROR)
Returns:
Configured logger instance
"""
# Create logger
logger = logging.getLogger(name)
logger.setLevel(getattr(logging, level))
# Clear existing handlers
logger.handlers = []
# Create console handler with colors
console_handler = colorlog.StreamHandler(sys.stdout)
console_handler.setLevel(getattr(logging, level))
# Create detailed formatter for educational purposes
log_format = (
"%(log_color)s%(asctime)s - %(name)s - [%(levelname)s] - "
"%(filename)s:%(lineno)d - %(funcName)s() - %(message)s%(reset)s"
)
formatter = colorlog.ColoredFormatter(
log_format,
datefmt="%Y-%m-%d %H:%M:%S",
reset=True,
log_colors={
'DEBUG': 'cyan',
'INFO': 'green',
'WARNING': 'yellow',
'ERROR': 'red',
'CRITICAL': 'red,bg_white',
}
)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)
return logger
def log_execution_time(logger: Optional[logging.Logger] = None):
"""
Decorator to log function execution time for educational purposes.
Args:
logger: Logger instance to use
"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
nonlocal logger
if logger is None:
logger = logging.getLogger("vector_search")
logger.debug(f"Starting execution of {func.__name__}")
start_time = time.time()
try:
result = func(*args, **kwargs)
execution_time = time.time() - start_time
logger.info(
f"✅ {func.__name__} completed successfully in {execution_time:.4f} seconds"
)
return result
except Exception as e:
execution_time = time.time() - start_time
logger.error(
f"❌ {func.__name__} failed after {execution_time:.4f} seconds: {str(e)}"
)
raise
return wrapper
return decorator
class VectorSearchLogger:
"""Educational logger for vector search operations with detailed debugging."""
def __init__(self, logger: logging.Logger, show_embeddings: bool = False):
self.logger = logger
self.show_embeddings = show_embeddings
def log_indexing_start(self, doc_id: str, text: str):
"""Log the start of document indexing."""
self.logger.debug("=" * 80)
self.logger.info(f"📝 Starting INDEXING operation")
self.logger.debug(f"Document ID: {doc_id}")
self.logger.debug(f"Text length: {len(text)} characters")
self.logger.debug(f"Text preview: {text[:100]}..." if len(text) > 100 else f"Text: {text}")
def log_embedding_generation(self, text: str, embedding_shape: tuple, time_taken: float):
"""Log embedding generation details."""
self.logger.debug(f"🧮 Generating embeddings using BGE-M3 model")
self.logger.debug(f"Input text length: {len(text)} characters")
self.logger.debug(f"Embedding shape: {embedding_shape}")
self.logger.debug(f"Embedding generation time: {time_taken:.4f} seconds")
def log_embedding_vector(self, embedding, sample_size: int = 10):
"""Log embedding vector details for educational purposes."""
if self.show_embeddings:
self.logger.debug(f"Embedding vector (first {sample_size} dimensions): {embedding[:sample_size]}")
self.logger.debug(f"Embedding statistics - Min: {embedding.min():.6f}, Max: {embedding.max():.6f}, Mean: {embedding.mean():.6f}")
def log_index_update(self, index_type: str, doc_id: str, current_size: int):
"""Log index update operations."""
self.logger.info(f"📊 Updating {index_type.upper()} index")
self.logger.debug(f"Adding document {doc_id} to index")
self.logger.debug(f"Current index size: {current_size} documents")
def log_search_start(self, query: str, top_k: int):
"""Log the start of search operation."""
self.logger.debug("=" * 80)
self.logger.info(f"🔍 Starting SEARCH operation")
self.logger.debug(f"Query: {query}")
self.logger.debug(f"Retrieving top {top_k} results")
def log_search_results(self, results: list, distances: list, time_taken: float):
"""Log search results with detailed information."""
self.logger.info(f"✨ Search completed in {time_taken:.4f} seconds")
self.logger.debug(f"Found {len(results)} matching documents")
for i, (doc_id, distance) in enumerate(zip(results, distances), 1):
self.logger.debug(f" Rank {i}: Document {doc_id} (distance: {distance:.6f})")
def log_deletion(self, doc_id: str):
"""Log document deletion."""
self.logger.debug("=" * 80)
self.logger.info(f"🗑️ Starting DELETE operation")
self.logger.debug(f"Deleting document: {doc_id}")
def log_error(self, operation: str, error: Exception):
"""Log errors with context."""
self.logger.error(f"❌ Error during {operation}: {type(error).__name__}: {str(error)}")
self.logger.debug(f"Full error details:", exc_info=True)
def log_index_build(self, index_type: str, num_documents: int, parameters: dict):
"""Log index building process."""
self.logger.info(f"🏗️ Building {index_type.upper()} index")
self.logger.debug(f"Number of documents: {num_documents}")
self.logger.debug(f"Index parameters: {parameters}")
main.py¶
"""Main FastAPI application for vector similarity search service."""
import time
import argparse
from typing import List, Optional, Dict, Any
from contextlib import asynccontextmanager
import uvicorn
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import numpy as np
from config import ServiceConfig, IndexType
from logger import setup_logger, VectorSearchLogger
from embedding_service import EmbeddingService
from indexing import AnnoyIndex, HNSWIndex, VectorIndex
from document_store import DocumentStore
# Request/Response models
class IndexRequest(BaseModel):
"""Request model for indexing documents."""
text: str = Field(..., description="Text content 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 searching documents."""
query: str = Field(..., description="Search query text")
top_k: int = Field(default=10, ge=1, le=100, description="Number of results to return")
return_documents: bool = Field(default=True, description="Whether to return full documents")
class DeleteRequest(BaseModel):
"""Request model for deleting documents."""
doc_id: str = Field(..., description="Document ID to delete")
class SearchResult(BaseModel):
"""Search result model."""
doc_id: str
score: float
text: Optional[str] = None
metadata: Optional[Dict[str, Any]] = None
rank: int
class IndexResponse(BaseModel):
"""Response model for indexing operations."""
success: bool
doc_id: str
message: str
index_size: int
class DeleteResponse(BaseModel):
"""Response model for deletion operations."""
success: bool
message: str
index_size: int
class SearchResponse(BaseModel):
"""Response model for search operations."""
success: bool
query: str
results: List[SearchResult]
total_results: int
search_time_ms: float
class StatsResponse(BaseModel):
"""Response model for service statistics."""
index_type: str
index_size: int
document_count: int
embedding_dimension: int
model_name: str
# Global instances
config: ServiceConfig = None
logger = None
vec_logger: VectorSearchLogger = None
embedding_service: EmbeddingService = None
vector_index: VectorIndex = None
document_store: DocumentStore = None
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Manage application lifecycle."""
# Startup
global config, logger, vec_logger, embedding_service, vector_index, document_store
logger.info("=" * 80)
logger.info("🚀 Starting Vector Similarity Search Service")
logger.info("=" * 80)
# Initialize embedding service
logger.info("Initializing BGE-M3 embedding service...")
embedding_service = EmbeddingService(
model_name=config.model_name,
use_fp16=config.use_fp16,
max_seq_length=config.max_seq_length,
logger=vec_logger
)
# Initialize vector index based on configuration
embedding_dim = embedding_service.get_embedding_dimension()
logger.info(f"Initializing {config.index_type.value.upper()} vector index...")
if config.index_type == IndexType.ANNOY:
vector_index = AnnoyIndex(
dimension=embedding_dim,
n_trees=config.annoy_n_trees,
metric=config.annoy_metric,
logger=vec_logger
)
else: # HNSW
vector_index = HNSWIndex(
dimension=embedding_dim,
max_elements=config.max_documents,
ef_construction=config.hnsw_ef_construction,
M=config.hnsw_M,
ef_search=config.hnsw_ef_search,
space=config.hnsw_space,
logger=vec_logger
)
# Initialize document store
logger.info("Initializing document store...")
document_store = DocumentStore(logger=vec_logger)
logger.info("=" * 80)
logger.info("✅ Service initialized successfully!")
logger.info(f"📍 API available at http://{config.host}:{config.port}")
logger.info(f"📚 Docs available at http://{config.host}:{config.port}/docs")
logger.info("=" * 80)
yield
# Shutdown
logger.info("Shutting down service...")
# Create FastAPI app
app = FastAPI(
title="Vector Similarity Search Service",
description="Educational service for vector similarity search using BGE-M3 embeddings with ANNOY/HNSW indexing",
version="1.0.0",
lifespan=lifespan
)
# Add CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
@app.get("/", response_model=Dict[str, str])
async def root():
"""Root endpoint."""
return {
"service": "Vector Similarity Search",
"status": "running",
"index_type": config.index_type.value,
"model": config.model_name
}
@app.post("/index", response_model=IndexResponse)
async def index_document(request: IndexRequest):
"""
Index a new document.
This endpoint:
1. Generates embeddings using BGE-M3
2. Adds the document to the document store
3. Adds the embedding to the vector index
"""
try:
vec_logger.log_indexing_start(request.doc_id or "auto-generated", request.text)
# Generate embedding
start_time = time.time()
embedding_result = embedding_service.encode_text(request.text)
embedding = embedding_result['dense']
embedding_time = time.time() - start_time
vec_logger.log_embedding_generation(
request.text,
embedding.shape,
embedding_time
)
# Store document
doc_id = document_store.add_document(
text=request.text,
doc_id=request.doc_id,
metadata=request.metadata
)
# Update document with embedding
document_store.update_document_embedding(doc_id, embedding.tolist())
# Add to vector index
vector_index.add_item(doc_id, embedding)
vec_logger.log_index_update(
config.index_type.value,
doc_id,
vector_index.get_size()
)
# Rebuild index if necessary (for ANNOY)
if config.index_type == IndexType.ANNOY:
vector_index.rebuild_index()
return IndexResponse(
success=True,
doc_id=doc_id,
message=f"Document indexed successfully using {config.index_type.value.upper()}",
index_size=vector_index.get_size()
)
except Exception as e:
vec_logger.log_error("indexing", e)
raise HTTPException(status_code=500, detail=str(e))
@app.post("/search", response_model=SearchResponse)
async def search_documents(request: SearchRequest):
"""
Search for similar documents.
This endpoint:
1. Generates query embedding using BGE-M3
2. Searches the vector index for similar documents
3. Returns ranked results with scores
"""
try:
vec_logger.log_search_start(request.query, request.top_k)
# Generate query embedding
start_time = time.time()
embedding_result = embedding_service.encode_text(request.query)
query_embedding = embedding_result['dense']
embedding_time = time.time() - start_time
logger.debug(f"Query embedding generated in {embedding_time:.4f}s")
vec_logger.log_embedding_vector(query_embedding, sample_size=10)
# Search in index
search_start = time.time()
doc_ids, distances = vector_index.search(query_embedding, request.top_k)
search_time = time.time() - search_start
vec_logger.log_search_results(doc_ids, distances, search_time)
# Prepare results
results = []
if request.return_documents:
documents = document_store.get_documents_by_ids(doc_ids)
doc_map = {doc.id: doc for doc in documents}
for rank, (doc_id, distance) in enumerate(zip(doc_ids, distances), 1):
doc = doc_map.get(doc_id)
if doc:
results.append(SearchResult(
doc_id=doc_id,
score=float(1.0 / (1.0 + distance)), # Convert distance to similarity score
text=doc.text,
metadata=doc.metadata,
rank=rank
))
else:
for rank, (doc_id, distance) in enumerate(zip(doc_ids, distances), 1):
results.append(SearchResult(
doc_id=doc_id,
score=float(1.0 / (1.0 + distance)),
rank=rank
))
total_time_ms = (time.time() - start_time) * 1000
return SearchResponse(
success=True,
query=request.query,
results=results,
total_results=len(results),
search_time_ms=total_time_ms
)
except Exception as e:
vec_logger.log_error("search", e)
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/index", response_model=DeleteResponse)
async def delete_document(request: DeleteRequest):
"""
Delete a document from the index.
This endpoint:
1. Removes the document from the document store
2. Removes the embedding from the vector index
"""
try:
vec_logger.log_deletion(request.doc_id)
# Delete from document store
doc_deleted = document_store.delete_document(request.doc_id)
if not doc_deleted:
return DeleteResponse(
success=False,
message=f"Document {request.doc_id} not found",
index_size=vector_index.get_size()
)
# Delete from vector index
index_deleted = vector_index.delete_item(request.doc_id)
if index_deleted:
return DeleteResponse(
success=True,
message=f"Document {request.doc_id} deleted successfully",
index_size=vector_index.get_size()
)
else:
return DeleteResponse(
success=False,
message=f"Document {request.doc_id} deleted from store but not from index",
index_size=vector_index.get_size()
)
except Exception as e:
vec_logger.log_error("deletion", e)
raise HTTPException(status_code=500, detail=str(e))
@app.get("/stats", response_model=StatsResponse)
async def get_stats():
"""Get service statistics."""
return StatsResponse(
index_type=config.index_type.value,
index_size=vector_index.get_size(),
document_count=document_store.get_size(),
embedding_dimension=embedding_service.get_embedding_dimension(),
model_name=config.model_name
)
@app.get("/documents", response_model=List[Dict[str, Any]])
async def list_documents(limit: int = Query(default=10, ge=1, le=100)):
"""List documents in the store."""
docs = document_store.list_documents(limit=limit)
return [
{
"id": doc.id,
"text": doc.text[:200] + "..." if len(doc.text) > 200 else doc.text,
"metadata": doc.metadata,
"created_at": doc.created_at.isoformat()
}
for doc in docs
]
def main():
"""Main entry point."""
global config, logger, vec_logger
# Parse command line arguments
parser = argparse.ArgumentParser(description="Vector Similarity Search Service")
parser.add_argument(
"--index-type",
type=str,
choices=["annoy", "hnsw"],
default="hnsw",
help="Type of index to use (default: hnsw)"
)
parser.add_argument(
"--host",
type=str,
default="0.0.0.0",
help="Host to bind to (default: 0.0.0.0)"
)
parser.add_argument(
"--port",
type=int,
default=4240,
help="Port to bind to (default: 4240)"
)
parser.add_argument(
"--debug",
action="store_true",
help="Enable debug mode"
)
parser.add_argument(
"--show-embeddings",
action="store_true",
help="Show embedding vectors in logs"
)
args = parser.parse_args()
# Create configuration
config = ServiceConfig(
index_type=IndexType(args.index_type),
host=args.host,
port=args.port,
debug=args.debug,
show_embeddings=args.show_embeddings
)
# Setup logging
logger = setup_logger("vector_search", config.log_level)
vec_logger = VectorSearchLogger(logger, config.show_embeddings)
# Run the service
uvicorn.run(
app,
host=config.host,
port=config.port,
log_level=config.log_level.lower(),
reload=False
)
if __name__ == "__main__":
main()
quick_demo.py¶
#!/usr/bin/env python3
"""Quick demo script to showcase the vector similarity search service."""
import time
import sys
def print_section(title):
"""Print a formatted section header."""
print("\n" + "=" * 60)
print(f" {title}")
print("=" * 60)
def main():
"""Run a quick demo of the service."""
print_section("Vector Similarity Search - Quick Demo")
print("""
This educational service demonstrates vector similarity search
using BGE-M3 embeddings with ANNOY/HNSW indexing.
EDUCATIONAL CONCEPTS DEMONSTRATED:
1. Text → Vector embedding generation
2. Approximate nearest neighbor search
3. Cosine similarity for semantic matching
4. Trade-offs between index types (ANNOY vs HNSW)
""")
print("\n📚 STEP 1: Start the service")
print("-" * 40)
print("\nOption A - Using HNSW (high precision):")
print(" python main.py --index-type hnsw --debug")
print("\nOption B - Using ANNOY (fast, memory-efficient):")
print(" python main.py --index-type annoy --debug")
print("\nOption C - Using the startup script:")
print(" ./start_service.sh hnsw 8000 true")
print("\n📝 STEP 2: Index some documents")
print("-" * 40)
print("""
Example using curl:
curl -X POST http://localhost:8000/index \\
-H "Content-Type: application/json" \\
-d '{
"text": "Machine learning is a subset of AI that enables systems to learn from data.",
"metadata": {"category": "AI", "level": "beginner"}
}'
""")
print("\n🔍 STEP 3: Search for similar documents")
print("-" * 40)
print("""
Example search:
curl -X POST http://localhost:8000/search \\
-H "Content-Type: application/json" \\
-d '{
"query": "What is deep learning?",
"top_k": 5
}'
""")
print("\n🎯 STEP 4: Run the test client")
print("-" * 40)
print("""
The test client will:
- Index 10 sample documents about AI, programming, and DevOps
- Perform 5 different similarity searches
- Demonstrate document deletion
- Show performance metrics
Run it with:
python test_client.py
For performance testing (100 documents):
python test_client.py --performance
""")
print("\n📊 KEY LEARNING POINTS")
print("-" * 40)
print("""
1. EMBEDDINGS: BGE-M3 converts text → 1024-dimensional vectors
- Semantic meaning is captured in vector space
- Similar texts have similar vectors
2. INDEXING: Two algorithms for efficient similarity search
- ANNOY: Tree-based, fast but approximate
- HNSW: Graph-based, slower but more accurate
3. SIMILARITY: Cosine distance measures semantic similarity
- Score close to 1.0 = very similar
- Score close to 0.0 = not similar
4. TRADE-OFFS:
- Speed vs Accuracy (ANNOY vs HNSW)
- Memory vs Performance (index parameters)
- Build time vs Search time
""")
print("\n🔗 USEFUL ENDPOINTS")
print("-" * 40)
print("""
- API Documentation: http://localhost:8000/docs
- Service Status: http://localhost:8000/
- Statistics: http://localhost:8000/stats
- List Documents: http://localhost:8000/documents
""")
print("\n💡 EXPERIMENT IDEAS")
print("-" * 40)
print("""
1. Compare ANNOY vs HNSW accuracy on same queries
2. Measure indexing time for different document sizes
3. Test multilingual search (BGE-M3 supports 100+ languages)
4. Analyze how different parameters affect performance
5. Try searching with synonyms and paraphrases
""")
print_section("Ready to Start!")
print("\nNext steps:")
print("1. Start the service: python main.py --debug")
print("2. Run the demo: python test_client.py")
print("3. Explore the API: http://localhost:8000/docs")
print()
if __name__ == "__main__":
main()
test_client.py¶
"""Test client for the vector similarity search service."""
import requests
import json
import time
from typing import List, Dict, Any
class VectorSearchClient:
"""Client for testing the vector search service."""
def __init__(self, base_url: str = "http://localhost:8000"):
"""Initialize the client."""
self.base_url = base_url
def index_document(self, text: str, doc_id: str = None, metadata: Dict = None) -> Dict:
"""Index a document."""
response = requests.post(
f"{self.base_url}/index",
json={
"text": text,
"doc_id": doc_id,
"metadata": metadata or {}
}
)
return response.json()
def search(self, query: str, top_k: int = 5, return_documents: bool = True) -> Dict:
"""Search for similar documents."""
response = requests.post(
f"{self.base_url}/search",
json={
"query": query,
"top_k": top_k,
"return_documents": return_documents
}
)
return response.json()
def delete_document(self, doc_id: str) -> Dict:
"""Delete a document."""
response = requests.delete(
f"{self.base_url}/index",
json={"doc_id": doc_id}
)
return response.json()
def get_stats(self) -> Dict:
"""Get service statistics."""
response = requests.get(f"{self.base_url}/stats")
return response.json()
def list_documents(self, limit: int = 10) -> List[Dict]:
"""List documents in the store."""
response = requests.get(f"{self.base_url}/documents", params={"limit": limit})
return response.json()
def run_demo():
"""Run a comprehensive demo of the vector search service."""
print("=" * 80)
print("🚀 Vector Similarity Search Service - Demo Client")
print("=" * 80)
# Initialize client
client = VectorSearchClient()
# Check service status
print("\n📊 Checking service status...")
try:
response = requests.get("http://localhost:8000/")
status = response.json()
print(f"✅ Service is running")
print(f" - Index type: {status['index_type']}")
print(f" - Model: {status['model']}")
except Exception as e:
print(f"❌ Service is not running: {e}")
print("Please start the service first with: python main.py")
return
# Sample documents
documents = [
{
"text": "Machine learning is a subset of artificial intelligence that enables systems to learn and improve from experience without being explicitly programmed.",
"metadata": {"category": "AI", "topic": "machine_learning"}
},
{
"text": "Deep learning is a type of machine learning based on artificial neural networks with multiple layers that progressively extract higher-level features from raw input.",
"metadata": {"category": "AI", "topic": "deep_learning"}
},
{
"text": "Natural language processing (NLP) is a branch of AI that helps computers understand, interpret and manipulate human language.",
"metadata": {"category": "AI", "topic": "nlp"}
},
{
"text": "Computer vision enables machines to interpret and understand visual information from the world, similar to how humans use their eyes and brains.",
"metadata": {"category": "AI", "topic": "computer_vision"}
},
{
"text": "Reinforcement learning is an area of machine learning where an agent learns to make decisions by taking actions in an environment to maximize cumulative reward.",
"metadata": {"category": "AI", "topic": "reinforcement_learning"}
},
{
"text": "Python is a high-level programming language known for its simplicity and readability, widely used in data science and web development.",
"metadata": {"category": "Programming", "topic": "python"}
},
{
"text": "JavaScript is a versatile programming language primarily used for creating interactive web applications and running code in browsers.",
"metadata": {"category": "Programming", "topic": "javascript"}
},
{
"text": "Docker is a platform that uses containerization to package applications with their dependencies, ensuring consistency across different environments.",
"metadata": {"category": "DevOps", "topic": "containerization"}
},
{
"text": "Kubernetes is an open-source container orchestration platform that automates the deployment, scaling, and management of containerized applications.",
"metadata": {"category": "DevOps", "topic": "orchestration"}
},
{
"text": "The transformer architecture revolutionized NLP by introducing self-attention mechanisms that process sequences in parallel rather than sequentially.",
"metadata": {"category": "AI", "topic": "transformers"}
}
]
# Index documents
print("\n📝 Indexing documents...")
doc_ids = []
for i, doc in enumerate(documents, 1):
print(f" [{i}/{len(documents)}] Indexing: {doc['text'][:50]}...")
result = client.index_document(
text=doc["text"],
metadata=doc["metadata"]
)
doc_ids.append(result["doc_id"])
time.sleep(0.1) # Small delay for demonstration
print(f"\n✅ Indexed {len(documents)} documents")
# Show statistics
print("\n📊 Current statistics:")
stats = client.get_stats()
for key, value in stats.items():
print(f" - {key}: {value}")
# Perform searches
queries = [
"What is deep learning and neural networks?",
"How to deploy applications with containers?",
"Programming languages for web development",
"Learning from environment and rewards",
"Understanding human language with computers"
]
print("\n🔍 Performing searches...")
for query in queries:
print(f"\n Query: '{query}'")
print(" " + "-" * 60)
result = client.search(query, top_k=3)
if result["success"]:
print(f" Found {result['total_results']} results in {result['search_time_ms']:.2f}ms:")
for res in result["results"]:
print(f"\n Rank {res['rank']}: (Score: {res['score']:.4f})")
print(f" Text: {res['text'][:100]}...")
if res.get("metadata"):
print(f" Metadata: {res['metadata']}")
else:
print(f" ❌ Search failed")
time.sleep(0.5) # Delay for readability
# Test deletion
print("\n🗑️ Testing document deletion...")
if doc_ids:
doc_to_delete = doc_ids[0]
print(f" Deleting document: {doc_to_delete}")
delete_result = client.delete_document(doc_to_delete)
print(f" Result: {delete_result['message']}")
print(f" New index size: {delete_result['index_size']}")
# List remaining documents
print("\n📋 Listing documents (first 5)...")
docs = client.list_documents(limit=5)
for i, doc in enumerate(docs, 1):
print(f" {i}. ID: {doc['id'][:8]}... | Text: {doc['text'][:50]}...")
print("\n" + "=" * 80)
print("✅ Demo completed successfully!")
print("=" * 80)
def run_performance_test():
"""Run a performance test."""
print("\n" + "=" * 80)
print("⚡ Performance Test")
print("=" * 80)
client = VectorSearchClient()
# Generate test documents
num_docs = 100
print(f"\n📝 Indexing {num_docs} documents for performance testing...")
start_time = time.time()
for i in range(num_docs):
text = f"This is test document number {i}. It contains various information about topic {i % 10}. " \
f"The content is randomly generated for testing purposes. Keywords: test, document, {i}, performance."
client.index_document(text)
if (i + 1) % 20 == 0:
print(f" Indexed {i + 1}/{num_docs} documents...")
index_time = time.time() - start_time
print(f"\n✅ Indexing completed in {index_time:.2f} seconds")
print(f" Average: {index_time/num_docs*1000:.2f}ms per document")
# Test search performance
print(f"\n🔍 Testing search performance with 20 queries...")
search_times = []
for i in range(20):
query = f"Find information about topic {i % 10} and document testing"
start_time = time.time()
result = client.search(query, top_k=10, return_documents=False)
search_time = time.time() - start_time
search_times.append(search_time)
avg_search_time = sum(search_times) / len(search_times)
min_search_time = min(search_times)
max_search_time = max(search_times)
print(f"\n📊 Search Performance Results:")
print(f" - Average search time: {avg_search_time*1000:.2f}ms")
print(f" - Min search time: {min_search_time*1000:.2f}ms")
print(f" - Max search time: {max_search_time*1000:.2f}ms")
# Final stats
stats = client.get_stats()
print(f"\n📊 Final Statistics:")
print(f" - Total documents: {stats['document_count']}")
print(f" - Index size: {stats['index_size']}")
print(f" - Index type: {stats['index_type']}")
if __name__ == "__main__":
import sys
if len(sys.argv) > 1 and sys.argv[1] == "--performance":
run_performance_test()
else:
run_demo()
start_service.sh¶
#!/bin/bash
# Vector Similarity Search Service Startup Script
echo "========================================"
echo "Vector Similarity Search Service"
echo "========================================"
# Default values
INDEX_TYPE=${1:-hnsw}
PORT=${2:-8000}
DEBUG=${3:-true}
echo ""
echo "Configuration:"
echo " Index Type: $INDEX_TYPE"
echo " Port: $PORT"
echo " Debug: $DEBUG"
echo ""
# Check if virtual environment exists
if [ -d "venv" ]; then
echo "Activating virtual environment..."
source venv/bin/activate
fi
# Install dependencies if needed
echo "Checking dependencies..."
pip list | grep -q "FlagEmbedding" || {
echo "Installing dependencies..."
pip install -r requirements.txt
}
# Start the service
echo ""
echo "Starting service..."
echo "========================================"
if [ "$DEBUG" = "true" ]; then
python main.py --index-type $INDEX_TYPE --port $PORT --debug
else
python main.py --index-type $INDEX_TYPE --port $PORT
fi