Function main

sample/python/ivf_rabitq_indexing.py:17–48

(args=None)

def main(args=None) -> None:
    # 1. Load data
    data = read_fvecs(args.data_file)
    n, dim = data.shape
    print(f"Data loaded")
    print(f"\tN: {n}")
    print(f"\tDIM: {dim}")

    # 2. Cluster with FAISS
    centroids, cluster_ids = cluster_data(data, args.num_clusters, args.metric, args.num_threads)
    print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}")

    # 3. Build IVF index
    print(f"\nBuilding IVF index: bits={args.total_bits}, metric={args.metric}, "
          f"num_threads={args.num_threads}, faster_quant={args.faster_quant}")

    idx = IvfIndex(
        dim=dim,
        max_elements=n,
        num_clusters=args.num_clusters,
        nbits=args.total_bits,
        metric=args.metric,
    )

    t0 = time()
    idx.build(data, centroids, cluster_ids, num_threads=args.num_threads, fast_quantization=args.faster_quant)
    elapsed_min = (time() - t0) / 60

    print("IVF constructed")
    idx.save(args.index_file)
    print(f"Indexing time: {elapsed_min:.4f} min")
    print(f"Index saved → {args.index_file}")


if __name__ == "__main__":