Function main

knn_utils/build_index.py:204–258

()

    return index_name

def main():
    parser = argparse.ArgumentParser(description="Build a FAISS index from an Arrow datastore")
    
    # Only required parameter is the dstore_path
    parser.add_argument(
        "--dstore_path", 
        type=str, 
        required=True,
        help="Path to the Arrow file containing the datastore"
    )
    
    # Optional parameters
    parser.add_argument(
        "--num_keys_to_add_at_a_time", 
        type=int, 
        default=1_000_000,
        help="Number of keys to add at a time"
    )
    parser.add_argument(
        "--ncentroids", 
        type=int, 
        default=4096,
        help="Number of centroids for IVFPQ"
    )
    parser.add_argument(
        "--seed", 
        type=int, 
        default=42,
        help="Random seed"
    )
    parser.add_argument(
        "--code_size", 
        type=int, 
        default=32,
        help="Code size for PQ"
    )
    parser.add_argument(
        "--probe", 
        type=int, 
        default=8,
        help="Number of probes for query"
    )
    
    args = parser.parse_args()
    
    index_path = build_index(
        dstore_path=args.dstore_path,
        num_keys_to_add_at_a_time=args.num_keys_to_add_at_a_time,
        ncentroids=args.ncentroids,
        seed=args.seed,
        code_size=args.code_size,
        probe=args.probe
    )
    
    logger.info(f"Index saved to: {index_path}")

if __name__ == "__main__":
    main()