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

chapter3/dense-embedding/docker_annoy_runner.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""Linux-isolated ANNOY measurement used when the host ARM wheel is broken."""
+
+import json
+import os
+import statistics
+import sys
+import tempfile
+import time
+
+import numpy as np
+from annoy import AnnoyIndex
+
+
+def latency_stats(values):
+    return {
+        "mean": statistics.mean(values),
+        "p50": float(np.percentile(values, 50)),
+        "p95": float(np.percentile(values, 95)),
+    }
+
+
+def main():
+    input_path, output_path = sys.argv[1:3]
+    data = np.load(input_path, allow_pickle=False)
+    ids = [str(x) for x in data["ids"]]
+    vectors = data["vectors"].astype("float32")
+    queries = data["queries"].astype("float32")
+    initial_truth = data["initial_truth"]
+    full_truth = data["full_truth"]
+    initial_n, k, repeats = (int(x) for x in data["parameters"])
+    dimension = vectors.shape[1]
+
+    index = AnnoyIndex(dimension, "angular")
+    started = time.perf_counter()
+    for i, vector in enumerate(vectors[:initial_n]):
+        index.add_item(i, vector.tolist())
+    index.build(50)
+    build_ms = (time.perf_counter() - started) * 1000
+
+    recalls, latencies, rankings = [], [], []
+    for q_idx, query in enumerate(queries):
+        first = None
+        for _ in range(repeats):
+            started = time.perf_counter()
+            found = index.get_nns_by_vector(query.tolist(), k, -1, False)
+            latencies.append((time.perf_counter() - started) * 1000)
+            if first is None:
+                first = found
+        recalls.append(len(set(first) & set(initial_truth[q_idx].tolist())) / k)
+        rankings.append({"query_index": q_idx, "doc_ids": [ids[i] for i in first]})
+
+    with tempfile.NamedTemporaryFile() as handle:
+        index.save(handle.name)
+        serialized_bytes = os.path.getsize(handle.name)
+
+    # ANNOY cannot mutate a built index: full update means rebuilding a fresh tree.
+    started = time.perf_counter()
+    updated = AnnoyIndex(dimension, "angular")
+    for i, vector in enumerate(vectors):
+        updated.add_item(i, vector.tolist())
+    updated.build(50)
+    update_ms = (time.perf_counter() - started) * 1000
+    update_recalls = []
+    for q_idx, query in enumerate(queries):
+        found = updated.get_nns_by_vector(query.tolist(), k, -1, False)
+        update_recalls.append(len(set(found) & set(full_truth[q_idx].tolist())) / k)
+
+    payload = {
+        "build_ms": round(build_ms, 3),
+        "recall_at_k": statistics.mean(recalls),
+        "query_latency_ms": latency_stats(latencies),
+        "serialized_bytes": serialized_bytes,
+        "rankings": rankings,
+        "incremental_update": {
+            "items_added": len(ids) - initial_n,
+            "latency_ms": round(update_ms, 3),
+            "requires_full_rebuild": True,
+            "recall_at_k_after_update": statistics.mean(update_recalls),
+        },
+    }
+    with open(output_path, "w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2)
+
+
+if __name__ == "__main__":
+    main()