search.py

import argparse
import faiss
import h5py
import numpy as np
import os
from pathlib import Path
import time
from tqdm import tqdm

from deglib import Metric, Mt19937
from deglib.builder import EvenRegularGraphBuilder, LID
from deglib.graph import SizeBoundedGraph

from datasets import DATASETS, prepare, get_fn

def store_results(dst, algo, dataset, task, D, I, buildtime, querytime, params):
    os.makedirs(Path(dst).parent, exist_ok=True)
    f = h5py.File(dst, 'w')
    f.attrs['algo'] = algo
    f.attrs['dataset'] = dataset
    f.attrs['task'] = task
    f.attrs['buildtime'] = buildtime
    f.attrs['querytime'] = querytime
    f.attrs['params'] = params
    f.create_dataset('knns', I.shape, dtype=I.dtype)[:] = I
    f.create_dataset('dists', D.shape, dtype=D.dtype)[:] = D
    f.close()

def load_data(dataset, task):
    """Load data and queries for a given dataset and task."""
    prepare(dataset, task)

    fn, _ = get_fn(dataset, task)
    with h5py.File(fn) as f:
        queries = np.array(DATASETS[dataset][task]['queries'](f)).astype(np.float32)

    return queries

def run_faissIVF(dataset, task, k, threads=8):
    faiss.omp_set_num_threads(threads)
    print(f'\nRunning {task} on {dataset} with faiss IVF')

    data = load_data(dataset, task)

    n, d = data.shape
    if task == 'task2':
        k = k + 1 # need to search for one more NN since we cannot remove self-loop

    # setup IVF index
    nlist = 1024 # number of clusters/centroids to build the IVF from
    index_identifier = f"IVF{nlist},SQfp16"
    index = faiss.index_factory(d, index_identifier)

    # build index
    print(f"Training {index_identifier} index on {data.shape}")
    start = time.time()
    index.train(data)
    index.add(data)
    elapsed_build = time.time() - start
    print(f"Done training in {elapsed_build}s.")
    assert index.is_trained

    # search index
    for nprobe in [1, 2, 5, 10]:
        print(f"Starting search on {data.shape} with nprobe={nprobe}")
        start = time.time()
        index.nprobe = nprobe
        D, I = index.search(data, k)
        elapsed_search = time.time() - start + (elapsed_build if task == 'task2' else 0)
        print(f"Done searching in {elapsed_search}s. Got {I.shape} indices ({I.dtype}) and {D.shape} distances {D.dtype}.")

        I = I + 1 # FAISS is 0-indexed, groundtruth is 1-indexed
        identifier = f"index=({index_identifier}),query=(nprobe={nprobe})"
        store_results(os.path.join("results/", dataset, task, f"{identifier}.h5"), "faissIVF", dataset, task, D, I, elapsed_build, elapsed_search, identifier)


def run_faissHNSW(dataset, task, k, threads=8):
    faiss.omp_set_num_threads(threads)
    print(f'\nRunning {task} on {dataset} with faiss HNSW')

    data = load_data(dataset, task)

    n, d = data.shape
    if task == 'task2':
        k = k + 1 # need to search for one more NN since we cannot remove self-loop

    # setup HNSW graph index
    index_M = 16  # M parameter in HNSW
    index_metric = faiss.METRIC_INNER_PRODUCT
    index = faiss.IndexHNSWFlat(d, index_M, index_metric)
    index.hnsw.efConstruction = 30  # Set efConstruction for the graph building phase
    metric_short = 'L2' if index_metric == faiss.METRIC_L2 else 'IP'
    algo = f"HNSW{index_M},{metric_short},efc={index.hnsw.efConstruction}"

    # build graph
    print(f"Add data to {algo} index on {data.shape}")
    start = time.time()
    index.add(data)
    elapsed_build = time.time() - start
    print(f"Done building in {elapsed_build}s.")

    # perform search
    for ef_search in [10, 50, 100]:
        print(f"Starting search on {data.shape} with efSearch={ef_search}")
        start = time.time()
        index.hnsw.efSearch = ef_search
        D, I = index.search(data, k)
        elapsed_search = time.time() - start + (elapsed_build if task == 'task2' else 0)
        print(f"Done searching in {elapsed_search}s. Got {I.shape} indices ({I.dtype}) and {D.shape} distances {D.dtype}.")

        I = I + 1  # FAISS is 0-indexed, groundtruth is 1-indexed
        identifier = f"index=({algo}),query=(efSearch={ef_search})"
        store_results(os.path.join("results/", dataset, task, f"{identifier}.h5"), algo, dataset, task, D, I, elapsed_build, elapsed_search, identifier)


def run_deg(dataset, task, k, threads=8):
    print(f'\nRunning {task} on {dataset} with DEG')

    data = load_data(dataset, task)

    n, d = data.shape
    if task == 'task2':
        k = k + 1 # need to search for one more NN since we cannot remove self-loop

    # setup DEG index
    edges_per_vertex = 16
    metric: Metric = Metric.L2_Uint8
    lid = LID.Low
    extend_k = 16
    extend_eps = 0.1
    graph = SizeBoundedGraph.create_empty(n, data.shape[1], edges_per_vertex, metric)
    builder = EvenRegularGraphBuilder(
        graph, Mt19937(), lid=lid, extend_k=extend_k, extend_eps=extend_eps, improve_k=0, improve_eps=0,
        max_path_length=0, swap_tries=0, additional_swap_tries=0
    )
    builder.set_thread_count(threads)
    metric_short = 'IP' if metric == Metric.InnerProduct else metric.name
    algo = f"DEG{edges_per_vertex},{metric_short},LID={lid.name},extK={extend_k},extEps={extend_eps}"

    print(f"Prepare and add data to {algo} index on {data.shape}")
    start = time.time()

    # normalize the data and queries into uint8
    if metric == Metric.L2_Uint8:
        data_max = np.max(data, axis=0, keepdims=True)
        data_min = np.min(data, axis=0, keepdims=True)
        denom = data_max - data_min
        denom[denom == 0] = 1  # Prevent division by zero
        data = ((data - data_min) / denom * 255).astype(np.uint8)

    # Add entries in batches
    labels = np.arange(n, dtype=np.uint32)
    batch_size = 10000
    num_batches = (n + batch_size - 1) // batch_size
    for i in tqdm(range(num_batches), desc="Adding entries", unit="batch"):
        start_idx = i * batch_size
        end_idx = min((i + 1) * batch_size, n)
        builder.add_entry(labels[start_idx:end_idx], data[start_idx:end_idx])
        builder.build()

    elapsed_build = time.time() - start
    print(f"Done building in {elapsed_build}s.")

    # perform search
    for eps_search in [0.01, 0.02, 0.05, 0.1]:
        print(f"Starting search on {data.shape} with epsSearch={eps_search}")
        start = time.time()
        I, D = graph.search(data, threads=threads, thread_batch_size=1024, eps=eps_search, k=k)
        elapsed_search = time.time() - start + (elapsed_build if task == 'task2' else 0)
        print(f"Done searching in {elapsed_search}s. Got {I.shape} indices ({I.dtype}) and {D.shape} distances {D.dtype}.")

        I = I + 1  # deglib is 0-indexed, groundtruth is 1-indexed
        identifier = f"index=({algo}),query=(epsSearch={eps_search})"
        store_results(os.path.join("results/", dataset, task, f"{identifier}.h5"), algo, dataset, task, D, I, elapsed_build, elapsed_search, identifier)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--task",
        choices=['task2'],
        default='task2'
    )

    parser.add_argument(
        '--dataset',
        choices=DATASETS.keys(),
        default='ccnews-small'
    )

    parser.add_argument(
        '--algo',
        choices=['ivf', 'hnsw', 'deg', 'all'],
        default='deg',
        help='Algorithm to use: ivf, hnsw, deg or all'
    )

    parser.add_argument(
        '--threads',
        type=int,
        default=8,
        help='Number of threads to use for parallel processing'
    )

    args = parser.parse_args()

    if args.algo == 'ivf' or args.algo == 'all':
        run_faissIVF(args.dataset, args.task, DATASETS[args.dataset][args.task]['k'], args.threads)

    if args.algo == 'hnsw' or args.algo == 'all':
        run_faissHNSW(args.dataset, args.task, DATASETS[args.dataset][args.task]['k'], args.threads)

    if args.algo == 'deg' or args.algo == 'all':
        run_deg(args.dataset, args.task, DATASETS[args.dataset][args.task]['k'], args.threads)