ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/dense-embedding/benchmark.py

@@ -0,0 +1,303 @@
+#!/usr/bin/env python3
+"""Real-embedding ANNOY vs HNSW benchmark for Experiment 3-4."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import statistics
+import subprocess
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Sequence
+
+import numpy as np
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE.parent))
+from experiment_utils import sha256_file, write_campaign_evidence
+
+from indexing import AnnoyIndex, HNSWIndex
+
+
+TOPICS = [
+    ("vector search", "Approximate nearest-neighbor indexes accelerate semantic vector retrieval."),
+    ("database transactions", "Database transactions use atomicity, consistency, isolation and durability."),
+    ("photosynthesis", "Green plants turn sunlight and carbon dioxide into chemical energy."),
+    ("quantum entanglement", "Entangled particles exhibit correlated quantum measurement outcomes."),
+    ("contract law", "A valid contract generally requires offer acceptance and consideration."),
+    ("neural networks", "Deep neural networks learn layered nonlinear representations from data."),
+    ("cybersecurity", "Zero trust security continuously verifies identity and device posture."),
+    ("volcanoes", "Volcanoes form when magma rises through fractures in the planetary crust."),
+    ("water cycle", "Evaporation condensation precipitation and runoff form the water cycle."),
+    ("operating systems", "An operating system schedules processes and manages memory and devices."),
+    ("HTTP errors", "HTTP status 403 means a server understood but refused a request."),
+    ("machine translation", "Multilingual models translate meaning between natural languages."),
+    ("financial risk", "Portfolio diversification reduces exposure to idiosyncratic financial risk."),
+    ("medical imaging", "Radiology systems analyze X-rays CT scans and magnetic resonance images."),
+    ("supply chains", "Supply chain planning coordinates inventory logistics demand and suppliers."),
+    ("climate science", "Climate models simulate long-term interactions among atmosphere ocean and land."),
+    ("CPU instructions", "SIMD instructions apply one operation to several packed numeric values."),
+    ("compiler design", "A compiler parses source code optimizes intermediate form and emits machine code."),
+    ("graph theory", "Graph algorithms traverse vertices and edges to discover paths and communities."),
+    ("astronomy", "Astronomers infer stellar properties from spectra luminosity and orbital motion."),
+]
+
+
+class TransformerEncoder:
+    def __init__(self, model_name: str, device: str):
+        import torch
+        from transformers import AutoModel, AutoTokenizer
+
+        self.torch = torch
+        self.model_name = model_name
+        self.device = device
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
+        self.model = AutoModel.from_pretrained(model_name).to(device).eval()
+
+    def encode(self, texts: Sequence[str], query: bool = False, batch_size: int = 16) -> np.ndarray:
+        prefix = "Instruct: Retrieve semantically relevant passages.\nQuery:" if query else ""
+        values = [prefix + text for text in texts]
+        vectors = []
+        for start in range(0, len(values), batch_size):
+            batch = values[start : start + batch_size]
+            tokens = self.tokenizer(
+                batch, padding=True, truncation=True, max_length=192, return_tensors="pt"
+            ).to(self.device)
+            with self.torch.no_grad():
+                output = self.model(**tokens).last_hidden_state[:, -1].float()
+            output = self.torch.nn.functional.normalize(output, p=2, dim=1)
+            vectors.append(output.cpu().numpy())
+        return np.concatenate(vectors, axis=0).astype("float32")
+
+
+def build_corpus(n_docs: int) -> tuple[List[str], List[str]]:
+    docs, ids = [], []
+    variants = (
+        "A concise technical overview.",
+        "This passage explains the central mechanism and its practical use.",
+        "An engineering handbook entry with definitions and examples.",
+        "A research summary intended for a multilingual knowledge base.",
+        "Operational notes emphasizing trade-offs, reliability, and performance.",
+    )
+    for i in range(n_docs):
+        topic, sentence = TOPICS[i % len(TOPICS)]
+        variant = variants[(i // len(TOPICS)) % len(variants)]
+        docs.append(f"Topic: {topic}. {sentence} {variant} Document revision {i:04d}.")
+        ids.append(f"doc_{i:04d}")
+    return ids, docs
+
+
+def percentiles(values: List[float]) -> Dict[str, float]:
+    return {
+        "mean": statistics.mean(values),
+        "p50": float(np.percentile(values, 50)),
+        "p95": float(np.percentile(values, 95)),
+    }
+
+
+def exact_neighbors(matrix: np.ndarray, queries: np.ndarray, k: int) -> List[List[int]]:
+    scores = queries @ matrix.T
+    return [np.argsort(-row)[:k].tolist() for row in scores]
+
+
+def measure_index(name: str, index: Any, ids: List[str], vectors: np.ndarray,
+                  query_vectors: np.ndarray, truth: List[List[int]], k: int,
+                  repeats: int) -> Dict[str, Any]:
+    build_start = time.perf_counter()
+    for doc_id, vector in zip(ids, vectors):
+        index.add_item(doc_id, vector)
+    index.rebuild_index()
+    build_ms = (time.perf_counter() - build_start) * 1000
+
+    id_to_pos = {doc_id: pos for pos, doc_id in enumerate(ids)}
+    recalls, latencies = [], []
+    rankings = []
+    for q_idx, query in enumerate(query_vectors):
+        first = None
+        for _ in range(repeats):
+            started = time.perf_counter()
+            found, distances = index.search(query, k)
+            latencies.append((time.perf_counter() - started) * 1000)
+            if first is None:
+                first = found
+        found_pos = {id_to_pos[x] for x in first if x in id_to_pos}
+        recalls.append(len(found_pos & set(truth[q_idx])) / k)
+        rankings.append({"query_index": q_idx, "doc_ids": first})
+
+    with tempfile.NamedTemporaryFile(suffix=f".{name}") as handle:
+        if name == "annoy":
+            index.index.save(handle.name)
+        else:
+            index.index.save_index(handle.name)
+        serialized_bytes = os.path.getsize(handle.name)
+    return {
+        "build_ms": round(build_ms, 3),
+        "recall_at_k": statistics.mean(recalls),
+        "query_latency_ms": percentiles(latencies),
+        "serialized_bytes": serialized_bytes,
+        "rankings": rankings,
+    }
+
+
+def local_annoy_healthy() -> bool:
+    probe = AnnoyIndex(3, n_trees=5)
+    for i in range(5):
+        probe.add_item(str(i), np.array([i + 1, i + 2, i + 3], dtype="float32"))
+    probe.rebuild_index()
+    found, _ = probe.search(np.array([2, 3, 4], dtype="float32"), 3)
+    return len(found) == 3
+
+
+def measure_annoy_docker(ids: List[str], vectors: np.ndarray, queries: np.ndarray,
+                         initial_truth: List[List[int]], full_truth: List[List[int]],
+                         initial_n: int, k: int, repeats: int) -> tuple[Dict[str, Any], Dict[str, Any]]:
+    """Run the real Spotify ANNOY library in Linux when the macOS ARM extension
+    returns only item zero (a reproducible host-wheel defect in this environment)."""
+    with tempfile.TemporaryDirectory() as raw_dir:
+        work = Path(raw_dir)
+        np.savez(
+            work / "input.npz",
+            ids=np.asarray(ids, dtype="U32"), vectors=vectors, queries=queries,
+            initial_truth=np.asarray(initial_truth, dtype="int64"),
+            full_truth=np.asarray(full_truth, dtype="int64"),
+            parameters=np.asarray([initial_n, k, repeats], dtype="int64"),
+        )
+        command = [
+            "docker", "run", "--rm",
+            "-v", f"{work}:/work",
+            "-v", f"{HERE / 'docker_annoy_runner.py'}:/runner.py:ro",
+            "python:3.11-slim",
+            "sh", "-lc",
+            "apt-get update -qq && apt-get install -y -qq g++ >/dev/null && "
+            "pip install -q numpy annoy && python /runner.py /work/input.npz /work/output.json",
+        ]
+        started = time.perf_counter()
+        proc = subprocess.run(command, text=True, capture_output=True, timeout=600)
+        wall_ms = (time.perf_counter() - started) * 1000
+        if proc.returncode != 0:
+            raise RuntimeError(f"Docker ANNOY runner failed: {proc.stderr[-2000:]}")
+        result = json.loads((work / "output.json").read_text(encoding="utf-8"))
+    inspect = subprocess.check_output(
+        ["docker", "image", "inspect", "python:3.11-slim", "--format", "{{json .RepoDigests}}"],
+        text=True,
+    ).strip()
+    runtime = {
+        "kind": "docker-linux-aarch64",
+        "reason": "host macOS ARM ANNOY extension failed health check (returned fewer than k items)",
+        "base_image": "python:3.11-slim",
+        "base_image_repo_digests": json.loads(inspect),
+        "container_setup_and_run_wall_ms": round(wall_ms, 3),
+    }
+    return result, runtime
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Experiment 3-4 ANNOY/HNSW real embedding benchmark")
+    parser.add_argument("--model", default="Qwen/Qwen3-Embedding-0.6B")
+    parser.add_argument("--device", default="cpu")
+    parser.add_argument("--docs", type=int, default=300)
+    parser.add_argument("--top-k", type=int, default=10)
+    parser.add_argument("--repeats", type=int, default=5)
+    parser.add_argument("--seed", type=int, default=37)
+    args = parser.parse_args()
+    np.random.seed(args.seed)
+
+    ids, docs = build_corpus(args.docs)
+    queries = [f"Find technical information about {topic}." for topic, _ in TOPICS]
+    encoder = TransformerEncoder(args.model, args.device)
+    embed_start = time.perf_counter()
+    vectors = encoder.encode(docs)
+    query_vectors = encoder.encode(queries, query=True)
+    embedding_ms = (time.perf_counter() - embed_start) * 1000
+    truth = exact_neighbors(vectors, query_vectors, args.top_k)
+    dim = vectors.shape[1]
+
+    initial_n = int(args.docs * 0.8)
+    initial_truth = exact_neighbors(vectors[:initial_n], query_vectors, args.top_k)
+    backends = {
+        "hnsw": HNSWIndex(dim, max_elements=args.docs + 10, ef_construction=200, M=16, ef_search=100),
+    }
+    results = {}
+    if local_annoy_healthy():
+        backends["annoy"] = AnnoyIndex(dim, n_trees=50, metric="angular")
+        annoy_runtime = {"kind": "host", "health_check": "passed"}
+    else:
+        results["annoy"], annoy_runtime = measure_annoy_docker(
+            ids, vectors, query_vectors, initial_truth, truth,
+            initial_n, args.top_k, args.repeats,
+        )
+        print(f"annoy (Docker): recall@{args.top_k}={results['annoy']['recall_at_k']:.3f}, "
+              f"build={results['annoy']['build_ms']:.1f}ms")
+    for name, backend in backends.items():
+        initial = measure_index(
+            name, backend, ids[:initial_n], vectors[:initial_n], query_vectors,
+            initial_truth, args.top_k, args.repeats,
+        )
+        update_start = time.perf_counter()
+        for doc_id, vector in zip(ids[initial_n:], vectors[initial_n:]):
+            backend.add_item(doc_id, vector)
+        requires_rebuild = name == "annoy"
+        if requires_rebuild:
+            backend.rebuild_index()
+        update_ms = (time.perf_counter() - update_start) * 1000
+        id_to_pos = {doc_id: pos for pos, doc_id in enumerate(ids)}
+        update_recalls = []
+        for q_idx, query in enumerate(query_vectors):
+            found, _ = backend.search(query, args.top_k)
+            update_recalls.append(len({id_to_pos[x] for x in found} & set(truth[q_idx])) / args.top_k)
+        initial["incremental_update"] = {
+            "items_added": args.docs - initial_n,
+            "latency_ms": round(update_ms, 3),
+            "requires_full_rebuild": requires_rebuild,
+            "recall_at_k_after_update": statistics.mean(update_recalls),
+        }
+        results[name] = initial
+        print(f"{name}: recall@{args.top_k}={initial['recall_at_k']:.3f}, build={initial['build_ms']:.1f}ms")
+
+    cache_ref = Path.home() / ".cache" / "huggingface" / "hub" / f"models--{args.model.replace('/', '--')}" / "refs" / "main"
+    model_revision = cache_ref.read_text(encoding="utf-8").strip() if cache_ref.exists() else None
+    full = args.docs >= 300 and all(results[name]["recall_at_k"] >= 0.8 for name in results)
+    evidence = {
+        "status": "passed" if full else "partial",
+        "configuration": {
+            "embedding_model": args.model,
+            "model_revision": model_revision,
+            "device": args.device,
+            "seed": args.seed,
+            "dimension": dim,
+            "documents": args.docs,
+            "queries": len(queries),
+            "top_k": args.top_k,
+            "annoy_runtime": annoy_runtime,
+        },
+        "acceptance": {
+            "real_embedding_model": True,
+            "same_vectors_and_queries": True,
+            "exact_search_ground_truth": True,
+            "recall_latency_build_size_measured": True,
+            "incremental_behavior_measured": True,
+            "both_backends_recall_at_least_0_8": all(results[name]["recall_at_k"] >= 0.8 for name in results),
+            "passed": full,
+        },
+        "summary": {
+            "embedding_latency_ms": round(embedding_ms, 3),
+            "annoy": {k: v for k, v in results["annoy"].items() if k != "rankings"},
+            "hnsw": {k: v for k, v in results["hnsw"].items() if k != "rankings"},
+        },
+        "corpus": {"doc_ids": ids, "texts": docs, "queries": queries},
+        "results": results,
+    }
+    manifest = write_campaign_evidence(
+        HERE, "3-4", evidence,
+        input_paths=[HERE / "indexing.py", HERE / "benchmark.py", HERE / "docker_annoy_runner.py"]
+    )
+    print(json.dumps(manifest["summary"], ensure_ascii=False, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())