()
| 237 | |
| 238 | |
| 239 | def main(): |
| 240 | parser = argparse.ArgumentParser( |
| 241 | description="Convert SIFT1M fvecs/ivecs binary files to JSON format." |
| 242 | ) |
| 243 | parser.add_argument( |
| 244 | "--data-dir", |
| 245 | default=os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data"), |
| 246 | help="Directory containing fvecs/ivecs files and for JSON output (default: data/)", |
| 247 | ) |
| 248 | parser.add_argument( |
| 249 | "--shard-size", |
| 250 | type=int, |
| 251 | default=100000, |
| 252 | help="Maximum vectors per shard file for base vectors (default: 100000)", |
| 253 | ) |
| 254 | args = parser.parse_args() |
| 255 | |
| 256 | data_dir = args.data_dir |
| 257 | shard_size = args.shard_size |
| 258 | |
| 259 | if shard_size <= 0: |
| 260 | print("Error: --shard-size must be a positive integer.", file=sys.stderr) |
| 261 | sys.exit(1) |
| 262 | |
| 263 | # Input file paths |
| 264 | base_fvecs = os.path.join(data_dir, "sift_base.fvecs") |
| 265 | query_fvecs = os.path.join(data_dir, "sift_query.fvecs") |
| 266 | gt_ivecs = os.path.join(data_dir, "sift_groundtruth.ivecs") |
| 267 | |
| 268 | # --- Convert base vectors --- |
| 269 | print(f"Parsing base vectors from {base_fvecs}...") |
| 270 | try: |
| 271 | base_vectors = read_fvecs(base_fvecs) |
| 272 | except (FileNotFoundError, ValueError) as e: |
| 273 | print(f"Error: {e}", file=sys.stderr) |
| 274 | sys.exit(1) |
| 275 | print(f" Loaded {len(base_vectors)} vectors, dim={len(base_vectors[0])}") |
| 276 | |
| 277 | print("Converting base vectors to JSON...") |
| 278 | base_files = convert_base_vectors(base_vectors, data_dir, shard_size) |
| 279 | |
| 280 | # --- Convert query vectors --- |
| 281 | print(f"\nParsing query vectors from {query_fvecs}...") |
| 282 | try: |
| 283 | query_vectors = read_fvecs(query_fvecs) |
| 284 | except (FileNotFoundError, ValueError) as e: |
| 285 | print(f"Error: {e}", file=sys.stderr) |
| 286 | sys.exit(1) |
| 287 | print(f" Loaded {len(query_vectors)} vectors, dim={len(query_vectors[0])}") |
| 288 | |
| 289 | print("Converting query vectors to JSON...") |
| 290 | query_file = convert_query_vectors(query_vectors, data_dir) |
| 291 | |
| 292 | # --- Convert ground truth --- |
| 293 | print(f"\nParsing ground truth from {gt_ivecs}...") |
| 294 | try: |
| 295 | gt_vectors = read_ivecs(gt_ivecs) |
| 296 | except (FileNotFoundError, ValueError) as e: |
no test coverage detected