olifog.com
Embeddable hybrid graph-vector database in OCaml
6.9 kB
218 lines
1#!/usr/bin/env python3
2# /// script
3# requires-python = ">=3.10"
4# dependencies = [
5# "h5py",
6# "numpy",
7# ]
8# ///
9"""Download and convert standard ANN benchmark datasets.
10
11Downloads HDF5 files from ann-benchmarks.com and converts them to a simple
12binary format that the OCaml benchmark suite can load directly.
13
14Binary format (.fbin):
15 - 4 bytes: int32_le n (number of vectors)
16 - 4 bytes: int32_le dim (dimension)
17 - n * dim * 4 bytes: float32_le row-major data
18
19Ground truth format (.ibin):
20 - 4 bytes: int32_le n_queries
21 - 4 bytes: int32_le k
22 - n_queries * k * 4 bytes: int32_le row-major indices
23
24Usage:
25 uv run scripts/download_datasets.py [--dataset sift-128] [--output datasets/]
26 uv run scripts/download_datasets.py --list
27 uv run scripts/download_datasets.py --all
28"""
29
30import argparse
31import os
32import struct
33import sys
34import urllib.request
35
36import h5py
37import numpy as np
38
39DATASETS = {
40 "sift-128": {
41 "url": "http://ann-benchmarks.com/sift-128-euclidean.hdf5",
42 "metric": "euclidean",
43 "dim": 128,
44 "description": "SIFT 1M (128d, euclidean) — the universal ANN baseline",
45 },
46 "glove-100": {
47 "url": "http://ann-benchmarks.com/glove-100-angular.hdf5",
48 "metric": "angular",
49 "dim": 100,
50 "description": "GloVe 1.2M (100d, angular/cosine) — word embeddings",
51 },
52 "gist-960": {
53 "url": "http://ann-benchmarks.com/gist-960-euclidean.hdf5",
54 "metric": "euclidean",
55 "dim": 960,
56 "description": "GIST 1M (960d, euclidean) — high-dimensional stress test",
57 },
58 "fashion-mnist-784": {
59 "url": "http://ann-benchmarks.com/fashion-mnist-784-euclidean.hdf5",
60 "metric": "euclidean",
61 "dim": 784,
62 "description": "Fashion-MNIST 60K (784d, euclidean) — small & fast",
63 },
64 "nytimes-256": {
65 "url": "http://ann-benchmarks.com/nytimes-256-angular.hdf5",
66 "metric": "angular",
67 "dim": 256,
68 "description": "NYTimes 290K (256d, angular/cosine) — document embeddings",
69 },
70}
71
72
73def download_file(url: str, dest: str) -> None:
74 """Download a file with progress reporting."""
75 if os.path.exists(dest):
76 print(f" Already downloaded: {dest}")
77 return
78
79 print(f" Downloading {url}...")
80 tmp = dest + ".tmp"
81
82 req = urllib.request.Request(url, headers={
83 "User-Agent": "Mozilla/5.0 (gvecdb benchmark downloader)"
84 })
85 with urllib.request.urlopen(req) as resp:
86 total_size = int(resp.headers.get("Content-Length", 0))
87 downloaded = 0
88 with open(tmp, "wb") as f:
89 while True:
90 chunk = resp.read(1024 * 1024)
91 if not chunk:
92 break
93 f.write(chunk)
94 downloaded += len(chunk)
95 if total_size > 0:
96 pct = min(100, downloaded * 100 // total_size)
97 mb = downloaded / (1024 * 1024)
98 total_mb = total_size / (1024 * 1024)
99 print(f"\r {mb:.1f}/{total_mb:.1f} MB ({pct}%)", end="", flush=True)
100 print()
101 os.rename(tmp, dest)
102
103
104def write_fbin(path: str, data: np.ndarray) -> None:
105 """Write float32 vectors in binary format."""
106 n, dim = data.shape
107 data = data.astype(np.float32)
108 with open(path, "wb") as f:
109 f.write(struct.pack("<ii", n, dim))
110 f.write(data.tobytes())
111 print(f" Wrote {path}: {n} vectors, {dim}d ({os.path.getsize(path) / 1e6:.1f} MB)")
112
113
114def write_ibin(path: str, data: np.ndarray) -> None:
115 """Write int32 ground truth in binary format."""
116 n, k = data.shape
117 data = data.astype(np.int32)
118 with open(path, "wb") as f:
119 f.write(struct.pack("<ii", n, k))
120 f.write(data.tobytes())
121 print(f" Wrote {path}: {n} queries, k={k}")
122
123
124def convert_dataset(name: str, info: dict, output_dir: str) -> None:
125 """Download HDF5 and convert to binary format."""
126 dataset_dir = os.path.join(output_dir, name)
127 os.makedirs(dataset_dir, exist_ok=True)
128
129 hdf5_path = os.path.join(dataset_dir, f"{name}.hdf5")
130 download_file(info["url"], hdf5_path)
131
132 print(f" Converting {name}...")
133 with h5py.File(hdf5_path, "r") as f:
134 train = np.array(f["train"])
135 test = np.array(f["test"])
136 neighbors = np.array(f["neighbors"])
137
138 print(f" Base vectors: {train.shape}")
139 print(f" Query vectors: {test.shape}")
140 print(f" Ground truth: {neighbors.shape}")
141
142 write_fbin(os.path.join(dataset_dir, "base.fbin"), train)
143 write_fbin(os.path.join(dataset_dir, "queries.fbin"), test)
144
145 # Ground truth: take top-100 (or whatever's available)
146 k = min(100, neighbors.shape[1])
147 write_ibin(
148 os.path.join(dataset_dir, "groundtruth.ibin"), neighbors[:, :k]
149 )
150
151 # Write metadata
152 meta_path = os.path.join(dataset_dir, "metadata.txt")
153 with open(meta_path, "w") as f:
154 f.write(f"name: {name}\n")
155 f.write(f"metric: {info['metric']}\n")
156 f.write(f"dim: {info['dim']}\n")
157 f.write(f"base_vectors: {train.shape[0]}\n")
158 f.write(f"query_vectors: {test.shape[0]}\n")
159 f.write(f"ground_truth_k: {k}\n")
160
161 print(f" Done: {dataset_dir}/")
162
163
164def main():
165 parser = argparse.ArgumentParser(
166 description="Download standard ANN benchmark datasets"
167 )
168 parser.add_argument(
169 "--dataset",
170 type=str,
171 help=f"Dataset name. Available: {', '.join(DATASETS.keys())}",
172 )
173 parser.add_argument(
174 "--output", type=str, default="datasets", help="Output directory"
175 )
176 parser.add_argument(
177 "--list", action="store_true", help="List available datasets"
178 )
179 parser.add_argument(
180 "--all", action="store_true", help="Download all datasets"
181 )
182 args = parser.parse_args()
183
184 if args.list:
185 print("Available datasets:")
186 for name, info in DATASETS.items():
187 print(f" {name:25s} {info['description']}")
188 return
189
190 if args.all:
191 for name, info in DATASETS.items():
192 print(f"\n=== {name} ===")
193 try:
194 convert_dataset(name, info, args.output)
195 except Exception as e:
196 print(f" ERROR: {e}", file=sys.stderr)
197 return
198
199 if not args.dataset:
200 # Default: download sift-128 and glove-100
201 for name in ["sift-128", "glove-100"]:
202 print(f"\n=== {name} ===")
203 convert_dataset(name, DATASETS[name], args.output)
204 return
205
206 if args.dataset not in DATASETS:
207 print(
208 f"Unknown dataset: {args.dataset}. Available: {', '.join(DATASETS.keys())}",
209 file=sys.stderr,
210 )
211 sys.exit(1)
212
213 print(f"\n=== {args.dataset} ===")
214 convert_dataset(args.dataset, DATASETS[args.dataset], args.output)
215
216
217if __name__ == "__main__":
218 main()