Learned Metric Index (LMI) is a machine learning based data structure for fast look-up of approximate nearest neighbor queries for complex data
from li.Logger import Logger
from li.utils import pairwise_cosine
import time
import numpy as np


class Baseline(Logger):
    def __init__(self):
        pass

    def search(self, queries, data, k=10):
        s = time.time()
        anns = pairwise_cosine(data, queries).T
        nns = anns.argsort()[:, :k] + 1
        dists = np.sort(anns)[:, :k]
        return dists, nns, time.time() - s

    def build(self, data):
        s = time.time()
        self.logger.info("No build method implemented for baseline.")
        return time.time() - s

search/li/Baseline.py
LearnedMetricIndex/LearnedMetricIndex
from dataclasses import dataclass, field
from typing import Any, List, Union

from li.clustering import ClusteringAlgorithm
from li.model import ModelParameters


@dataclass
class BuildConfiguration:
    """
    BuildConfiguration specifies parameters of models on each level of the hierarchy.

    Expectes the arguments to be either a single value or a list of values of the same size
    as n_categories. If a single value is provided, it is expanded to the size of n_categories.
    The n_categories argument must be a list of integers specifying the number of categories
    at each level of the hierarchy.

    Example:
    ```
    config1 = BuildConfiguration(
        clustering.algorithms["faiss_kmeans"],
        40,
        "MLP-3",
        0.001,
        [5, 10],
    )

    config2 = BuildConfiguration(
        [clustering.algorithms["faiss_kmeans"]],
        [40],
        ["MLP-3"],
        [0.001],
        [5, 10],
    )
    ```

    In both cases, the arguments are expanded to the following (identical) configuration:
    ```
    config3 = BuildConfiguration(
        [clustering.algorithms["faiss_kmeans"], clustering.algorithms["faiss_kmeans"]],
        [40, 40],
        ["MLP-3", "MLP-3"],
        [0.001, 0.001],
        [5, 10],
    )
    ```
    """

    clustering_algorithms: List[ClusteringAlgorithm]
    epochs: List[int]
    model_types: List[str]
    lrs: List[float]
    n_categories: List[int]

    level_configurations: List[ModelParameters] = field(init=False)
    n_levels: int = field(init=False)

    def __init__(
        self,
        clustering_algorithms: Union[List[ClusteringAlgorithm], ClusteringAlgorithm],
        epochs: Union[List[int], int],
        model_types: Union[List[str], str],
        lrs: Union[List[float], float],
        n_categories: List[int],
    ):
        BuildConfiguration._validate(
            clustering_algorithms, epochs, model_types, lrs, n_categories
        )

        # Expand the arguments to the size of n_categories
        self.clustering_algorithms = BuildConfiguration._expand(
            clustering_algorithms, len(n_categories)
        )
        self.epochs = BuildConfiguration._expand(epochs, len(n_categories))
        self.model_types = BuildConfiguration._expand(model_types, len(n_categories))
        self.lrs = BuildConfiguration._expand(lrs, len(n_categories))
        self.n_categories = n_categories

        # Fields that are populated mainly for convenience and readability
        self.level_configurations = [
            ModelParameters(
                clustering_algorithm=self.clustering_algorithms[i],
                model_type=self.model_types[i],
                epochs=self.epochs[i],
                lr=self.lrs[i],
                n_categories=self.n_categories[i],
            )
            for i in range(len(self.n_categories))
        ]
        self.n_levels = len(self.n_categories)

    @staticmethod
    def _validate(
        clustering_algorithms: Union[List[ClusteringAlgorithm], ClusteringAlgorithm],
        epochs: Union[List[int], int],
        model_types: Union[List[str], str],
        lrs: Union[List[float], float],
        n_categories: List[int],
    ) -> None:
        """
        Validate the arguments to BuildConfiguration and raise an AssertionError if they are invalid.
        """
        assert len(n_categories) > 0, "n_categories must specify at least one level"

        arguments = [clustering_algorithms, epochs, model_types, lrs]

        arguments_are_lists = all([isinstance(arg, list) for arg in arguments])
        arguments_are_scalars = all(
            [
                callable(clustering_algorithms),
                isinstance(epochs, int),
                isinstance(model_types, str),
                isinstance(lrs, float),
            ]
        )

        assert (
            arguments_are_lists or arguments_are_scalars
        ), "clustering_algorithms, epochs, model_types, and lrs must be lists or single values"

        for arg in arguments:
            if isinstance(arg, list):
                assert len(arg) == 1 or len(arg) == len(n_categories), (
                    "clustering_algorithms, epochs, model_types, and lrs must "
                    "be lists of size 1 or the same size as n_categories"
                )

    @staticmethod
    def _expand(arg: Union[List[Any], Any], n_categories: int) -> List[Any]:
        """
        Expects arg to be either a single value, a list with a single value or a list
        of the same size as n_categories. If single value inputs are provided,
        they are expanded to the size of n_categories and returned as lists.
        """
        if isinstance(arg, list):
            if len(arg) == 1:
                return [arg[0]] * n_categories
            else:
                return arg
        else:
            return [arg] * n_categories

search/li/BuildConfiguration.py
LearnedMetricIndex/LearnedMetricIndex
import time
from collections import defaultdict
from logging import DEBUG, INFO
from typing import Dict, List, Tuple

import faiss
import numpy as np
import numpy.typing as npt
import pandas as pd
import torch
import torch.utils.data
from li.Logger import Logger
from li.model import NeuralNetwork, data_X_to_torch
from li.PriorityQueue import EMPTY_VALUE, PriorityQueue
from li.utils import filter_path_idxs, log_runtime
from tqdm import tqdm

torch.manual_seed(2023)
np.random.seed(2023)


class LearnedIndex(Logger):
    def __init__(
        self,
        root_model: NeuralNetwork,
        internal_models: Dict[Tuple, NeuralNetwork],
        bucket_paths: List[Tuple],
    ):
        self.root_model = root_model
        """The rood model of the index."""

        self.internal_models = internal_models
        """
        Dictionary mapping the path to the internal model.
        A path is padded with `EMPTY_VALUE` to the right to match the length of the longest path.
        """

        self.bucket_paths = bucket_paths
        """List of paths to the buckets."""

    def search(
        self,
        data_navigation: pd.DataFrame,
        queries_navigation: npt.NDArray[np.float32],
        data_search: pd.DataFrame,
        queries_search: npt.NDArray[np.float32],
        data_prediction: npt.NDArray[np.int64],
        n_categories: List[int],
        n_buckets: int = 1,
        k: int = 10,
    ) -> Tuple[npt.NDArray, npt.NDArray[np.uint32], Dict[str, float]]:
        """Searches for `k` nearest neighbors for each query in `queries`.

        Implementation details:
        - The search is done in two steps:
            1. The order in which the queries visit the buckets is precomputed.
            2. The queries are then searched in the `n_buckets` most similar buckets.

        Parameters
        ----------
        data_navigation : pd.DataFrame
            Data used for navigation.
        queries_navigation : npt.NDArray[np.float32]
            Queries used for navigation.
        data_search : pd.DataFrame
            Data used for the sequential search.
        queries_search : npt.NDArray[np.float32]
            Queries used for the sequential search.
        data_prediction : npt.NDArray[np.int64]
            Predicted paths for each data point.
        n_categories : List[int]
            Number of categories for each level of the index.
        n_buckets : int, optional
            Number of most similar buckets to search in, by default 1
        k : int, optional
            Number of nearest neighbors to search for, by default 10

        Returns
        -------
        Tuple[npt.NDArray, npt.NDArray[np.uint32], Dict[str, float]]
            Array of shape (queries_search.shape[0], k) with distances to nearest neighbors for each query,
            array of shape (queries_search.shape[0], k) with nearest neighbors for each query,
            dictionary with measured times.
        """
        measured_time = defaultdict(float)

        s = time.time()
        n_levels = len(n_categories)

        anns_final = None
        dists_final = None

        self.logger.debug("Precomputing bucket order")
        bucket_order, measured_time["inference"] = self._precompute_bucket_order(
            queries_navigation=queries_navigation,
            n_buckets=n_buckets,
            n_categories=n_categories,
        )

        # Add bucket location to each object as searching is done sequentially per bucket
        for level_to_search in range(1, n_levels + 1):
            data_navigation[f"category_L{level_to_search}"] = data_prediction[
                :, (level_to_search - 1)
            ]

        # Search in the `n_buckets` most similar buckets
        for bucket_order_idx in range(n_buckets):
            self.logger.debug(
                f"Searching in bucket {bucket_order_idx + 1} out of {n_buckets}"
            )
            (dists, anns, t_all, t_seq_search, t_sort) = self._search_single_bucket(
                data_navigation=data_navigation,
                data_search=data_search,
                queries_search=queries_search,
                bucket_path=bucket_order[:, bucket_order_idx, :],
                n_levels=n_levels,
            )

            measured_time["search_within_buckets"] += t_all
            measured_time["seq_search"] += t_seq_search
            measured_time["sort"] += t_sort

            self.logger.debug("Sorting the results")
            t = time.time()
            if anns_final is None:
                anns_final = anns
                dists_final = dists
            else:
                # stacks the results from the previous sorted anns and dists
                # *_final arrays now have shape (queries.shape[0], k*2)
                anns_final = np.hstack((anns_final, anns))
                dists_final = np.hstack((dists_final, dists))
                # gets the sorted indices of the stacked dists
                idx_sorted = dists_final.argsort(kind="stable", axis=1)[:, :k]
                # indexes the final arrays with the sorted indices
                # *_final arrays now have shape (queries.shape[0], k)
                idx = np.ogrid[tuple(map(slice, dists_final.shape))]
                idx[1] = idx_sorted
                dists_final = dists_final[tuple(idx)]
                anns_final = anns_final[tuple(idx)]

                assert (
                    anns_final.shape
                    == dists_final.shape
                    == (queries_search.shape[0], k)
                )
            self.logger.debug(f"Sorted the results in: {time.time() - t}")

        assert dists_final is not None
        assert anns_final is not None

        # Cleanup the placement of the objects
        for level_idx_to_search in range(n_levels):
            column_name = f"category_L{level_idx_to_search + 1}"

            if column_name in data_navigation.columns:
                data_navigation.drop(column_name, axis=1, inplace=True)

        measured_time["search"] = time.time() - s

        return dists_final, anns_final, measured_time

    @log_runtime(INFO, "Precomputed bucket order time: {}")
    def _precompute_bucket_order(
        self,
        queries_navigation: npt.NDArray[np.float32],
        n_buckets: int,
        n_categories: List[int],
    ) -> Tuple[npt.NDArray[np.int32], float]:
        """
        Precomputes the order in which the queries visit the buckets.

        Implementation details:
        - When visiting an internal node, the paths to the nodes/buckets under that node
        are added to the priority queue.
        - When visiting a bucket, the path to the bucket is stored in `bucket_order`.
        - The priority queue is then sorted by the probability of the next node/bucket to visit.
        - The computation is done until `n_buckets` buckets are visited for each query.

        Parameters
        ----------
        queries_navigation : np.ndarray
            Queries used for navigation.
        n_buckets : int
            Number of most similar buckets to precompute the order for.
        n_categories : List[int]
            Number of categories for each level of the index.

        Returns
        -------
        Tuple[npt.NDArray[np.int32], float]
            Array of shape (queries_navigation.shape[0], n_buckets, len(n_categories))
            with the order in which the queries visit the buckets,
            total inference time.
        """

        n_queries = queries_navigation.shape[0]
        n_levels = len(n_categories)
        assert self.root_model is not None, "Model is not trained, call `build` first."

        total_inference_t = 0.0

        s = time.time()
        pred_l1_prob, pred_l1_paths = self.root_model.predict_proba(
            data_X_to_torch(queries_navigation)
        )
        total_inference_t += time.time() - s

        if n_levels == 1:
            bucket_order = np.full(
                (n_queries, n_buckets, n_levels), fill_value=EMPTY_VALUE, dtype=np.int32
            )
            bucket_order[:, :n_buckets, 0] = pred_l1_paths[:, :n_buckets]
            return bucket_order, total_inference_t

        queue_length_upper_bound = int(np.prod(n_categories))
        pq = PriorityQueue(n_queries, queue_length_upper_bound, n_levels)

        # Populates the priority queue with the first level of the index
        # * Relies on the fact that the pred_l1_categories and pred_l1_probs are sorted,
        # * therefore the priority queue does not need to be sorted after this for loop
        for l1_idx in reversed(range(n_categories[0])):
            l1_paths = np.full(
                (n_queries, n_levels), fill_value=EMPTY_VALUE, dtype=np.int32
            )
            l1_paths[:, 0] = pred_l1_paths[:, l1_idx]
            pq.add(np.arange(n_queries), l1_paths, pred_l1_prob[:, l1_idx])

        bucket_order = np.full(
            (n_queries, n_buckets, n_levels), fill_value=EMPTY_VALUE, dtype=np.int32
        )
        bucket_order_length = np.zeros(n_queries, dtype=np.int32)

        while not np.all(bucket_order_length == n_buckets):
            query_idxs = np.where(bucket_order_length < n_buckets)[0]
            path_to_visit = pq.pop(query_idxs)

            inference_t = self._visit_internal_nodes(
                queries_navigation, query_idxs, pq, path_to_visit, n_levels
            )
            self._visit_buckets(
                query_idxs,
                path_to_visit,
                bucket_order,
                bucket_order_length,
            )

            total_inference_t += inference_t

            pq.sort()

        return bucket_order, total_inference_t

    def _visit_internal_nodes(
        self,
        queries_navigation: npt.NDArray[np.float32],
        all_query_idxs: npt.NDArray[np.int32],
        pq: PriorityQueue,
        path_to_visit: npt.NDArray[np.int32],
        n_levels: int,
    ) -> float:
        """
        Visits the internal nodes specified by `paths`.
        Paths to the buckets under a specific internal node are then added into `pq`.
        Done for each possible path to an internal node separately.
        """
        total_inference_t = 0.0

        for path in self.internal_models.keys():
            path_idxs = filter_path_idxs(path_to_visit, path)
            query_idxs = all_query_idxs[path_idxs]
            if query_idxs.shape[0] == 0:
                continue

            model = self.internal_models[path]

            s = time.time()
            probabilities, categories = model.predict_proba(
                data_X_to_torch(queries_navigation[query_idxs])
            )
            total_inference_t += time.time() - s

            level = len(path) - path.count(EMPTY_VALUE)
            n_model_categories = categories.shape[1]

            for child_idx in range(n_model_categories):
                child_paths = np.full(
                    (query_idxs.shape[0], n_levels),
                    fill_value=EMPTY_VALUE,
                    dtype=np.int32,
                )
                child_paths[:] = np.array(path)
                child_paths[:, level] = categories[:, child_idx]

                pq.add(
                    query_idxs,
                    child_paths,
                    probabilities[:, child_idx],
                )

        return total_inference_t

    def _visit_buckets(
        self,
        all_query_idxs: npt.NDArray[np.int32],
        path_to_visit: npt.NDArray[np.int32],
        bucket_order: npt.NDArray[np.int32],
        bucket_order_length: npt.NDArray[np.int32],
    ) -> None:
        """
        Visits the buckets specified by `paths`.
        The path to the bucket relevant to each query is then stored in `bucket_order`.
        Done for each possible bucket separately.
        """
        for path in self.bucket_paths:
            path_idxs = filter_path_idxs(path_to_visit, path)
            query_idxs = all_query_idxs[path_idxs]

            if query_idxs.shape[0] == 0:
                continue

            bucket_order[query_idxs, bucket_order_length[query_idxs], :] = np.array(
                path
            )
            bucket_order_length[query_idxs] += 1

    @log_runtime(DEBUG, "Searched the bucket in: {}")
    def _search_single_bucket(
        self,
        data_navigation: pd.DataFrame,
        data_search: pd.DataFrame,
        queries_search: npt.NDArray[np.float32],
        bucket_path: npt.NDArray[np.int32],
        k: int = 10,
        n_levels: int = 1,
    ) -> Tuple[npt.NDArray, npt.NDArray[np.uint32], float, float, float]:
        s_all = time.time()

        n_queries = queries_search.shape[0]
        nns = np.zeros((n_queries, k), dtype=np.uint32)
        dists = np.full((n_queries, k), fill_value=float("inf"), dtype=float)

        possible_bucket_paths = [
            f"category_L{level_to_search}" for level_to_search in range(1, n_levels + 1)
        ]

        t_seq_search = 0.0
        t_sort = 0.0

        for path, g in tqdm(data_navigation.groupby(possible_bucket_paths)):
            bucket_obj_indexes = g.index

            relevant_query_idxs = filter_path_idxs(bucket_path, path)

            if bucket_obj_indexes.shape[0] != 0 and relevant_query_idxs.shape[0] != 0:
                queries_for_this_bucket = queries_search[relevant_query_idxs]
                data_in_this_bucket = data_search.loc[bucket_obj_indexes].to_numpy()

                s = time.time()
                similarity, indices = faiss.knn(
                    queries_for_this_bucket,
                    data_in_this_bucket,
                    k,
                    metric=faiss.METRIC_INNER_PRODUCT,
                )
                t_seq_search += time.time() - s

                distances = 1 - similarity

                nns[relevant_query_idxs] = bucket_obj_indexes.to_numpy()[indices]
                dists[relevant_query_idxs] = distances

        return dists, nns, time.time() - s_all, t_seq_search, t_sort

search/li/LearnedIndex.py
LearnedMetricIndex/LearnedMetricIndex
import time
from itertools import product, takewhile
from logging import DEBUG
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import numpy.typing as npt
import pandas as pd
import torch
import torch.utils.data
from li.BuildConfiguration import BuildConfiguration
from li.clustering import ClusteringAlgorithm
from li.LearnedIndex import LearnedIndex
from li.Logger import Logger
from li.model import LIDataset, ModelParameters, NeuralNetwork, data_X_to_torch
from li.PriorityQueue import EMPTY_VALUE
from li.utils import filter_path_idxs, log_runtime
from tqdm import tqdm


class LearnedIndexBuilder(Logger):
    def __init__(self, data: pd.DataFrame, config: BuildConfiguration):
        # Attributes used during the build
        self.data = data
        """Data to build the index on."""

        self.config = config
        """Model configurations for each level of the index."""

        # Attributes passed to the index after the build
        self.root_model: Optional[NeuralNetwork] = None
        """The root model of the index."""

        self.internal_models: Dict[Tuple, NeuralNetwork] = {}
        """
        Dictionary mapping the path to the internal model.
        A path is padded with `EMPTY_VALUE` to the right to match the length of the longest path.
        """

        self.bucket_paths: List[Tuple] = []
        """List of paths to the buckets."""

    def build(
        self,
    ) -> Tuple[LearnedIndex, npt.NDArray[np.int64], int, float, float]:
        """
        Builds the index.

        Parameters
        ----------
        data : pd.DataFrame
            Data to build the index on.
        config : BuildConfiguration
            Configuration for the training.

        Returns
        -------
        Tuple[npt.NDArray[np.int64], int, float, float]
            An array of shape (data.shape[0], len(config.n_levels)) with predicted paths for each data point,
            number of buckets, time it took to build the index, time it took to cluster the data.
        """
        s = time.time()

        n_levels = self.config.n_levels

        # Where should the training data be placed with respect to each level
        data_prediction: npt.NDArray[np.int64] = np.full(
            (self.data.shape[0], n_levels), fill_value=EMPTY_VALUE, dtype=np.int64
        )

        self.logger.debug("Training the root model.")
        self.root_model, root_cluster_t = self._train_model(
            self.data,
            self.config.level_configurations[0],
        )
        data_prediction[:, 0] = self.root_model.predict(data_X_to_torch(self.data))

        if n_levels == 1:
            for bucket_index in range(len(np.unique(data_prediction[:, 0]))):
                self.bucket_paths.append((bucket_index,))

            return (
                self._create_index(),
                data_prediction,
                len(self.bucket_paths),
                time.time() - s,
                root_cluster_t,
            )

        self.logger.debug(f"Training {self.config.n_categories[:-1]} internal models.")
        s_internal = time.time()
        internal_cluster_t = self._train_internal_models(
            self.data,
            data_prediction,
            self.config,
        )
        self.logger.debug(
            f"Trained {self.config.n_categories[:-1]} internal models in {time.time()-s_internal:.2f}s."
        )

        return (
            self._create_index(),
            data_prediction,
            len(self.bucket_paths),
            time.time() - s,
            root_cluster_t + internal_cluster_t,
        )

    def _create_index(self) -> LearnedIndex:
        """Creates the index from the trained models."""
        assert self.root_model is not None, "The root model is not trained."

        return LearnedIndex(
            self.root_model,
            self.internal_models,
            self.bucket_paths,
        )

    @log_runtime(DEBUG, "Trained the model in: {}")
    def _train_model(
        self,
        data: pd.DataFrame,
        model_parameters: ModelParameters,
    ) -> Tuple[NeuralNetwork, float]:
        """
        Trains a single model.
        The model is trained until it predicts the correct number of categories.
        The same number of epochs is used for each training iteration.

        Parameters
        ----------
        data : pd.DataFrame
            Data to train the model on.
        clustering_algorithm : ClusteringAlgorithm
            Clustering algorithm to use.
        model_type : str
            Type of the model.
        epochs : int
            The minimal number of epochs to train the model for.
        lr : float
            Learning rate for the model.
        n_categories : int
            Number of categories to predict.

        Returns
        -------
        Tuple[NeuralNetwork, float]
            Trained model, time it took to cluster the data.

        Raises
        ------
        RuntimeError
            If the model does not converge after 1000 iterations
            (after training `epochs` epochs 1000 times).
        """
        clustering_algorithm, model_type, epochs, lr, n_categories = model_parameters

        _, labels, cluster_t = self._cluster(data, clustering_algorithm, n_categories)
        n_clusters = len(np.unique(labels))

        if n_clusters != n_categories:
            self.logger.debug(
                "Clustering algorithm did not return %d clusters, got %d.",
                n_categories,
                n_clusters,
            )
            self.logger.debug("Setting n_categories to %d.", n_clusters)
            n_categories = n_clusters

        train_loader = torch.utils.data.DataLoader(
            dataset=LIDataset(data, labels),
            batch_size=256,
            sampler=torch.utils.data.SubsetRandomSampler(data.index.values.tolist()),
        )
        torch_data = data_X_to_torch(data)

        model = NeuralNetwork(
            input_dim=data.shape[1],
            output_dim=n_categories,
            lr=lr,
            model_type=model_type,
        )
        is_trained = False

        iters = 0
        while not is_trained:
            model.train_batch(train_loader, epochs=epochs, logger=self.logger)
            predictions = model.predict(torch_data)
            iters += 1

            if iters > 1_000:
                raise RuntimeError("The model did not converge after 1000 iterations.")

            is_trained = len(np.unique(predictions)) == n_categories

        if iters > 1:
            self.logger.debug(
                f"Trained for {iters * epochs} epochs instead of {epochs}."
            )

        return model, cluster_t

    def _train_internal_models(
        self,
        data: pd.DataFrame,
        data_prediction: npt.NDArray[np.int64],
        config: BuildConfiguration,
    ) -> float:
        """
        Trains the internal models.

        ! The `data_prediction` array is modified in-place.

        Parameters
        ----------
        data : pd.DataFrame
            Data to train the models on.
        data_prediction : npt.NDArray[np.int64]
            Predicted paths for each data point.
        config : BuildConfiguration
            Configuration for the training.

        Returns
        -------
        float
            Time it took to cluster the data.
        """
        assert (
            self.root_model is not None
        ), "The root model is not trained, call `_train_root_model` first."

        overall_cluster_t = 0.0

        for level in range(1, config.n_levels):
            internal_node_paths = self._generate_internal_node_paths(
                level, config.n_levels, config.n_categories
            )
            self.logger.debug(f"Training level {level}.")

            for path in tqdm(internal_node_paths):
                self.logger.debug(f"Training model on path {path}.")

                data_idxs = filter_path_idxs(data_prediction, path)
                assert (
                    data_idxs.shape[0] != 0
                ), "There are no data points associated with the given path."

                # +1 as the data is indexed from 1
                training_data = data.loc[data_idxs + 1]

                # The subset needs to be reindexed; otherwise, the object accesses are invalid.
                original_pd_indices = training_data.index.values
                training_data = training_data.set_index(
                    pd.Index(range(1, training_data.shape[0] + 1))
                )

                model, cluster_t = self._train_model(
                    training_data,
                    config.level_configurations[level],
                )
                self.internal_models[path] = model

                overall_cluster_t += cluster_t

                # Restore back to the original indices
                training_data = training_data.set_index(
                    pd.Index(original_pd_indices.tolist())
                )

                predictions = model.predict(data_X_to_torch(training_data))

                # original_pd_indices-1 as data is indexed from 1
                # level as we are predicting the next level but the indexing is 0-based
                data_prediction[original_pd_indices - 1, level] = predictions

                if level == config.n_levels - 1:
                    for bucket_index in range(len(np.unique(predictions))):
                        self.bucket_paths.append(path[:-1] + (bucket_index,))

        return overall_cluster_t

    def _cluster(
        self,
        data: pd.DataFrame,
        clustering_algorithm: ClusteringAlgorithm,
        n_clusters: int,
    ) -> Tuple[Optional[Any], npt.NDArray[np.int32], float]:
        s = time.time()

        if data.shape[0] < 2:
            return None, np.array([0] * data.shape[0]), time.time() - s

        if data.shape[0] < n_clusters:
            n_clusters = data.shape[0] // 5
            if n_clusters < 2:
                n_clusters = 2

        clustering_object, labels = clustering_algorithm(
            np.array(data),
            n_clusters,
            None,
        )

        return clustering_object, labels, time.time() - s

    def _serialize_path(self, path: Tuple) -> str:
        """
        Serializes the path to a string.

        Example:
        >>> self._serialize_path((1, 2, -1, -1))
        "1.2"
        """
        valid_path = takewhile(lambda x: x != EMPTY_VALUE, path)

        return ".".join(map(str, valid_path))

    def _deserialize_path(self, path: str, n_levels: int) -> Tuple:
        """
        Deserializes the path from a string.

        Example:
        >>> self._deserialize_path("1.2", 4)
        (1, 2, -1, -1)
        """
        levels = path.split(".")

        return tuple(list(map(int, levels)) + [EMPTY_VALUE] * (n_levels - len(levels)))

    def _generate_internal_node_paths(
        self, level: int, n_levels: int, n_categories: List[int]
    ) -> List[Tuple]:
        """Generates all possible paths to internal nodes at the given `level`.

        Parameters
        ----------
        level : int
            Desired level of the internal nodes.
        n_levels : int
            Total number of levels in the index.
        n_categories : List[int]
            Number of categories for each level of the index.

        Returns
        -------
        List[Tuple]
            List of all possible paths to internal nodes at the given `level`.
        """
        path_combinations = [range(n_categories[lvl]) for lvl in range(level)]
        padding = [[EMPTY_VALUE]] * (n_levels - level)

        return list(product(*path_combinations, *padding))

search/li/LearnedIndexBuilder.py
LearnedMetricIndex/LearnedMetricIndex
import logging


def get_logger_config() -> str:
    return "[%(asctime)s][%(levelname)-5.5s][%(name)-.20s] %(message)s"


def remove_logger_handlers():
    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)


class Logger:
    @property
    def logger(self, level=logging.INFO):
        component = "{}.{}".format(type(self).__module__, type(self).__name__)
        logging.basicConfig(level=level, format=get_logger_config())
        return logging.getLogger(component)

search/li/Logger.py
LearnedMetricIndex/LearnedMetricIndex
import numpy as np
import numpy.typing as npt

EMPTY_VALUE = -1


class PriorityQueue:
    """
    A priority queue storing probabilities and paths for the next nodes to visit for each query.

    The priority queue is realized by three numpy arrays:
    - `probability`: stores the probability of the next node/bucket to visit
    - `path`: stores the path to the next node/bucket
    - `length`: stores the current length of the queue for each query
    - `should_sort`: stores whether the queue associated with this query should be sorted
    """

    def __init__(self, n_queries: int, queue_length_upper_bound: int, n_levels: int):
        self.probability: npt.NDArray[np.float32] = np.full(
            (n_queries, queue_length_upper_bound),
            fill_value=EMPTY_VALUE,
            dtype=np.float32,
        )
        self.path: npt.NDArray[np.int32] = np.full(
            (n_queries, queue_length_upper_bound, n_levels),
            fill_value=EMPTY_VALUE,
            dtype=np.int32,
        )
        self.length: npt.NDArray[np.int32] = np.zeros(n_queries, dtype=np.int32)

        self.should_sort: npt.NDArray[np.bool_] = np.full(
            n_queries, fill_value=np.False_, dtype=np.bool_
        )
        self.n_levels = n_levels

    def add(
        self,
        indices: npt.NDArray[np.int32],
        path: npt.NDArray[np.int32],
        probabilities: npt.NDArray[np.float32],
    ) -> None:
        """
        Adds a new node/bucket path to visit to the priority queue
        but only for queries specified by `indices`.
        """
        self.probability[indices, self.length[indices]] = probabilities
        self.path[indices, self.length[indices], :] = path
        self.should_sort[indices] = np.True_

        self.length[indices] += 1

    def pop(self, indices: npt.NDArray[np.int32]) -> npt.NDArray[np.int32]:
        """Pops the next node/bucket path to visit for each query."""
        self.length[indices] -= 1

        return self.path[indices, self.length[indices], :]

    def sort(self) -> None:
        """
        Sorts the queues by the probability.
        A particular queue is sorted only if `should_sort` is `True`.

        Implementation details:
        Firstly, we obtain the indexes of the sorted probabilities.
        Then, we use these indexes to sort the probabilities and paths.
        Sorting of paths is done for each level separately.
        The whole process is repeated for each queue length separately.
        """
        for queue_length in np.unique(self.length):
            if queue_length in {0, 1}:
                continue

            idxs_to_sort = np.where(
                np.logical_and(
                    self.length == queue_length,
                    self.should_sort == np.True_,
                )
            )[0]

            sorted_idxs = self.probability[idxs_to_sort, :queue_length].argsort()

            self.probability[idxs_to_sort, :queue_length] = np.take_along_axis(
                self.probability[idxs_to_sort, :queue_length],
                sorted_idxs,
                axis=1,
            )
            for level_idx in range(self.n_levels):
                self.path[idxs_to_sort, :queue_length, level_idx] = np.take_along_axis(
                    self.path[idxs_to_sort, :queue_length, level_idx],
                    sorted_idxs,
                    axis=1,
                )

            self.should_sort[idxs_to_sort] = np.False_

search/li/PriorityQueue.py
LearnedMetricIndex/LearnedMetricIndex
from typing import Any, Callable, Dict, Optional, Tuple

import numpy as np
import numpy.typing as npt

from .faiss_kmeans import cluster as faiss_kmeans
from .scikit_kmeans import cluster as scikit_kmeans

ClusteringAlgorithm = Callable[
    [npt.NDArray[np.float32], int, Optional[Dict[str, Any]]],
    Tuple[Any, npt.NDArray[np.int32]],
]

algorithms: Dict[str, ClusteringAlgorithm] = {
    "faiss_kmeans": faiss_kmeans,
    "scikit_kmeans": scikit_kmeans,
}

search/li/clustering/__init__.py
LearnedMetricIndex/LearnedMetricIndex
from typing import Any, Dict, Optional, Tuple

import numpy as np
import numpy.typing as npt
from faiss import Kmeans


def cluster(
    data: npt.NDArray[np.float32],
    n_clusters: int,
    parameters: Optional[Dict[str, Any]],
) -> Tuple[Kmeans, npt.NDArray[np.int32]]:
    if parameters is None:
        parameters = {"verbose": False, "seed": 2023}

    _, d = data.shape

    kmeans = Kmeans(d=d, k=n_clusters, **parameters)
    kmeans.train(data)

    labels = kmeans.index.search(data, 1)[1].T[0]  # type: ignore

    return kmeans, labels.astype(np.int32)

search/li/clustering/faiss_kmeans.py
LearnedMetricIndex/LearnedMetricIndex
from typing import Any, Dict, Optional, Tuple

import numpy as np
import numpy.typing as npt
from sklearn.cluster import KMeans


def cluster(
    data: npt.NDArray[np.float32],
    n_clusters: int,
    parameters: Optional[Dict[str, Any]],
) -> Tuple[KMeans, npt.NDArray[np.int32]]:
    if parameters is None:
        parameters = {
            "verbose": 0,
            "random_state": 2023,
            # The same default values as in faiss.Kmeans
            "init": "random",
            "max_iter": 25,
            "n_init": 1,
            #
        }

    kmeans = KMeans(n_clusters=n_clusters, **parameters)
    kmeans.fit(data)

    labels = kmeans.labels_

    return kmeans, labels

search/li/clustering/scikit_kmeans.py
LearnedMetricIndex/LearnedMetricIndex
from dataclasses import astuple, dataclass
from typing import Callable, Dict, Optional, Tuple

import numpy as np
import torch
import torch.nn.functional as nnf
import torch.utils.data
from li.clustering import ClusteringAlgorithm
from li.Logger import Logger
from torch import nn
from torch.nn import Linear, ReLU, Sequential

torch.manual_seed(2023)
np.random.seed(2023)


@dataclass(frozen=True)
class ModelParameters:
    clustering_algorithm: ClusteringAlgorithm
    model_type: str
    epochs: int
    lr: float
    n_categories: int

    def __iter__(self):
        return iter(astuple(self))


supported_models: Dict[str, Callable[[int, int], Sequential]] = {
    "MLP": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 128),
        ReLU(),
        Linear(128, output_dim),
    ),
    "MLP-2": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 64),
        ReLU(),
        Linear(64, output_dim),
    ),
    "MLP-3": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 256),
        ReLU(),
        Linear(256, output_dim),
    ),
    "MLP-4": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 512),
        ReLU(),
        Linear(512, output_dim),
    ),
    "MLP-5": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 256),
        ReLU(),
        Linear(256, 128),
        ReLU(),
        Linear(128, output_dim),
    ),
    "MLP-6": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 32),
        ReLU(),
        Linear(32, output_dim),
    ),
    "MLP-7": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 16),
        ReLU(),
        Linear(16, output_dim),
    ),
    "MLP-8": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 8),
        ReLU(),
        Linear(8, output_dim),
    ),
    "MLP-9": lambda input_dim, output_dim: Sequential(
        Linear(input_dim, 8),
        ReLU(),
        Linear(8, 16),
        ReLU(),
        Linear(16, output_dim),
    ),
}


def init_layers(model_type: Optional[str], input_dim: int, output_dim: int):
    if model_type not in supported_models:
        raise ValueError(f"Model type {model_type} not supported.")

    return supported_models[model_type](input_dim, output_dim)


class Model(nn.Module):
    def __init__(
        self, input_dim=768, output_dim=1000, model_type: Optional[str] = None
    ):
        super().__init__()
        self.layers = init_layers(model_type, input_dim, output_dim)
        self.n_output_neurons = output_dim

    def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
        outputs = self.layers(x)
        return outputs


def data_X_to_torch(data) -> torch.FloatTensor:
    """Creates torch training data."""
    data_X = torch.from_numpy(np.array(data).astype(np.float32))
    return data_X


def data_to_torch(data, labels) -> Tuple[torch.FloatTensor, torch.LongTensor]:
    """Creates torch training data and labels."""
    data_X = data_X_to_torch(data)
    data_y = torch.as_tensor(torch.from_numpy(labels), dtype=torch.long)
    return data_X, data_y


def get_device() -> torch.device:
    """Gets the `device` to be used by torch.
    This arugment is needed to operate with the PyTorch model instance.

    Returns
    ------
    torch.device
        Device
    """
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda:0" if use_cuda else "cpu")
    torch.backends.cudnn.benchmark = True
    return device


class NeuralNetwork(Logger):
    """The neural network class corresponding to every inner node.

    Parameters
    ----------
    input_dim : int
        The input dimension.
    output_dim : int
        The output dimension.
    loss : torch.nn, optional
        The loss function, the default is torch.nn.CrossEntropyLoss.
    lr : float, optional
        The learning rate, the default is 0.001.
    model_type : str, optional
        The model type, the default is 'MLP'.
    class_weight : torch.FloatTensor, optional
        The class weights, the default is None.
    """

    def __init__(
        self,
        input_dim,
        output_dim,
        loss=torch.nn.CrossEntropyLoss,
        lr=0.1,
        model_type="MLP",
        class_weight=None,
    ):
        self.device = get_device()
        self.model = Model(input_dim, output_dim, model_type=model_type).to(self.device)
        if not isinstance(class_weight, type(None)):
            self.loss = loss(weight=class_weight.to(self.device))
        else:
            self.loss = loss()
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=lr)

    def train(
        self,
        data_X: torch.FloatTensor,
        data_y: torch.LongTensor,
        epochs=500,
        logger=None,
    ):
        step = epochs // 10
        losses = []
        if logger:
            logger.debug(f"Epochs: {epochs}, step: {step}")
        for ep in range(epochs):
            pred_y = self.model(data_X.to(self.device))
            curr_loss = self.loss(pred_y, data_y.to(self.device))
            if ep % step == 0 and ep != 0:
                if logger:
                    logger.debug(f"Epoch {ep} | Loss {curr_loss.item()}")
            losses.append(curr_loss.item())

            self.model.zero_grad()
            curr_loss.backward()

            self.optimizer.step()
        return losses

    def train_batch(self, dataset, epochs=5, logger=None):
        step = epochs // 10
        step = step if step > 0 else 1
        losses = []
        if logger:
            logger.debug(f"Epochs: {epochs}, step: {step}")
        for ep in range(epochs):
            for data_X, data_y in iter(dataset):
                pred_y = self.model(data_X.to(self.device))
                curr_loss = self.loss(pred_y, data_y.to(self.device))

            if ep % step == 0 and ep != 0:
                if logger:
                    logger.debug(f"Epoch {ep} | Loss {curr_loss.item():.5f}")
            losses.append(curr_loss.item())

            self.model.zero_grad()
            curr_loss.backward()

            self.optimizer.step()
        return losses

    def predict(self, data_X: torch.FloatTensor):
        """Collects predictions for multiple data points (used in structure building)."""
        self.model = self.model.to(self.device)
        self.model.eval()

        all_outputs = torch.tensor([], device=self.device)
        with torch.no_grad():
            outputs = self.model(data_X.to(self.device))
            all_outputs = torch.cat((all_outputs, outputs), 0)

        _, y_pred = torch.max(all_outputs, 1)
        return y_pred.cpu().numpy()

    def predict_proba(self, data_X: torch.FloatTensor):
        """Collects predictions for a single data point (used in query predictions)."""
        self.model = self.model.to(self.device)
        self.model.eval()

        with torch.no_grad():
            outputs = self.model(data_X.to(self.device))

        if outputs.dim() == 1:
            dim = 0
        else:
            dim = 1
        prob = nnf.softmax(outputs, dim=dim)
        probs, classes = prob.topk(prob.shape[1])

        return probs.cpu().numpy(), classes.cpu().numpy()


class LIDataset(torch.utils.data.Dataset):
    def __init__(self, dataset_x, dataset_y):
        self.dataset_x, self.dataset_y = data_to_torch(dataset_x, dataset_y)

    def __len__(self):
        return self.dataset_x.shape[0]

    def __getitem__(self, idx):
        return self.dataset_x[idx - 1], self.dataset_y[idx - 1]

search/li/model.py
LearnedMetricIndex/LearnedMetricIndex
import functools
import time
from typing import Any, List, Tuple, Union

import numpy as np
import numpy.typing as npt
from sklearn.metrics.pairwise import cosine_similarity


def pairwise_cosine(x, y):
    return 1 - cosine_similarity(x, y)


def save_as_pickle(filename: str, obj):
    """
    Saves an object as a pickle file.
    Expects that the directory with the file exists.

    Parameters
    ----------
    filename : str
        Path to the file to write to.
    obj: object
        The object to save.
    """
    import pickle

    with open(filename, "wb") as f:
        pickle.dump(obj, f)


def log_runtime(level: int, message: str):
    """
    A decorator that prints the runtime of the decorated method.
    The decorated method must be a method of a class that has a logger property
    (e.g., a class that inherits from Logger).
    The message must contain a placeholder for the runtime.
    """

    def decorator_log_runtime(func):
        @functools.wraps(func)
        def wrapper_log_runtime(self, *args, **kwargs):
            start = time.time()
            result = func(self, *args, **kwargs)
            stop = time.time()

            self.logger.log(level, message.format(stop - start))

            return result

        return wrapper_log_runtime

    return decorator_log_runtime


def serialize(lst: List[Any]) -> str:
    """Serializes a list of objects into a string."""
    return ",".join(map(str, lst))


def filter_path_idxs(
    paths: npt.NDArray[Union[np.int32, np.int64]], path: Tuple
) -> npt.NDArray[np.int32]:
    """Returns the indexes of `paths` that match the given `path`."""
    return np.where(np.all(paths == np.array(path), axis=1))[0]

search/li/utils.py
LearnedMetricIndex/LearnedMetricIndex
import argparse
import logging
import os
import time
from pathlib import Path
from typing import Any, Dict, List
from urllib.request import urlretrieve

import h5py
import numpy as np
import numpy.typing as npt
import pandas as pd
from li.Baseline import Baseline
from li.BuildConfiguration import BuildConfiguration
from li.clustering import algorithms
from li.LearnedIndexBuilder import LearnedIndexBuilder
from li.utils import save_as_pickle, serialize
from sklearn import preprocessing

np.random.seed(2023)

logging.basicConfig(
    level=logging.INFO,
    format="[%(asctime)s][%(levelname)-5.5s][%(name)-.20s] %(message)s",
)
LOG = logging.getLogger(__name__)

MODELS_DIR_NAME = "models"


def download(src, dst):
    if not os.path.exists(dst):
        os.makedirs(Path(dst).parent, exist_ok=True)
        LOG.info("downloading %s -> %s..." % (src, dst))
        urlretrieve(src, dst)


def prepare(kind, size):
    url = "https://sisap-23-challenge.s3.amazonaws.com/SISAP23-Challenge"
    task = {
        "query": f"{url}/public-queries-10k-{kind}.h5",
        "dataset": f"{url}/laion2B-en-{kind}-n={size}.h5",
    }

    for version, url in task.items():
        target_path = os.path.join("data", kind, size, f"{version}.h5")
        download(url, target_path)
        assert os.path.exists(target_path), f"Failed to download {url}"


def store_results(dst, algo, kind, dists, anns, buildtime, querytime, params, size):
    os.makedirs(Path(dst).parent, exist_ok=True)
    LOG.info(f"Storing results in {dst}")
    f = h5py.File(dst, "w")
    f.attrs["algo"] = algo
    f.attrs["data"] = kind
    f.attrs["buildtime"] = buildtime
    f.attrs["querytime"] = querytime
    f.attrs["size"] = size
    f.attrs["params"] = params
    f.create_dataset("knns", anns.shape, dtype=anns.dtype)[:] = anns
    f.create_dataset("dists", dists.shape, dtype=dists.dtype)[:] = dists
    f.close()


def format_identifier(
    bucket: int,
    kind: str,
    config: BuildConfiguration,
    clustering_algorithms: List[str],
    short_identifier: str,
    size: str,
):
    return (
        f"{short_identifier}"
        f"-{kind}"
        f"-{size}"
        f"-ep={serialize(config.epochs)}"
        f"-lr={serialize(config.lrs)}"
        f"-cat={serialize(config.n_categories)}"
        f"-model={serialize(config.model_types)}"
        f"-buck={bucket}"
        f"-clustering_algorithm={serialize(clustering_algorithms)}"
        f"-{os.environ['PBS_JOBID']}"
    )


def format_models_filename(
    kind: str,
    config: BuildConfiguration,
    clustering_algorithms: List[str],
    preprocess: bool,
    size: str,
):
    return (
        f"./{MODELS_DIR_NAME}/{kind}"
        f"-{size}"
        f"-ep={serialize(config.epochs)}"
        f"-lr={serialize(config.lrs)}"
        f"-cat={serialize(config.n_categories)}"
        f"-model={serialize(config.model_types)}"
        f"-prep={preprocess}"
        f"-clustering_algorithm={serialize(clustering_algorithms)}"
        f"-{os.environ['PBS_JOBID']}"
    )


def run(
    kind: str,
    key: str,
    size: str,
    k: int,
    index_type: str,
    n_buckets_perc: List[int],
    n_categories: List[int],
    epochs: List[int],
    model_types: List[str],
    lr: List[float],
    preprocess: bool,
    save: bool,
    clustering_algorithms: List[str],
):
    assert index_type in {
        "baseline",
        "learned-index",
    }, f"Unknown index type: {index_type}"

    LOG.info(
        f"Running with: kind={kind}, key={key}, size={size}, n_buckets_perc={n_buckets_perc},"
        f" n_categories={n_categories}, clustering_algorithms={clustering_algorithms},"
        f" epochs={epochs}, lr={lr}, model_types={model_types}, preprocess={preprocess}, save={save}"
    )

    prepare(kind, size)

    data: npt.NDArray[np.float32] = np.array(
        h5py.File(os.path.join("data", kind, size, "dataset.h5"), "r")[key]
    )
    queries: npt.NDArray[np.float32] = np.array(
        h5py.File(os.path.join("data", kind, size, "query.h5"), "r")[key]
    )
    if preprocess:
        data = preprocessing.normalize(data)  # type: ignore
        queries = preprocessing.normalize(queries)  # type: ignore

    n, d = data.shape
    LOG.info(f"Loaded downloaded data, shape: n={n}, d={d}")
    LOG.info(f"Loaded downloaded queries, shape: queries={queries.shape}")

    if index_type == "baseline":
        baseline = Baseline()
        build_t = baseline.build(data)
        LOG.info(f"Build time: {build_t}")
        dists, nns, search_t = baseline.search(queries=queries, data=data, k=k)
    elif index_type == "learned-index":
        evaluate_learned_index(
            data,
            clustering_algorithms,
            epochs,
            model_types,
            lr,
            k,
            kind,
            n_buckets_perc,
            n_categories,
            preprocess,
            queries,
            save,
            size,
        )


def evaluate_learned_index(
    data: npt.NDArray[np.float32],
    clustering_algorithms: List[str],
    epochs: List[int],
    model_type: List[str],
    lr: List[float],
    k: int,
    kind: str,
    n_buckets_perc: List[int],
    n_categories: List[int],
    preprocess: bool,
    queries: npt.NDArray[np.float32],
    save: bool,
    size: str,
):
    s = time.time()
    # ---- data to pd.DataFrame ---- #
    data_pd = pd.DataFrame(data)
    data_pd.index += 1
    kind_search = "clip768v2"
    key_search = "emb"
    if kind != kind_search:
        LOG.info("Loading data to be used in search")
        prepare(kind_search, size)
        # ---- data_search to pd.DataFrame ---- #
        data_search = pd.DataFrame(
            np.array(
                h5py.File(os.path.join("data", kind_search, size, "dataset.h5"), "r")[
                    key_search
                ]
            )
        )
        data_search.index += 1
        queries_search = np.array(
            h5py.File(os.path.join("data", kind_search, size, "query.h5"), "r")[
                key_search
            ]
        )
        n, d = data_search.shape
        LOG.info(f"Loaded downloaded data, shape: n={n}, d={d}")
        LOG.info(f"Loaded downloaded queries, shape: queries={queries_search.shape}")
    else:
        data_search = data_pd
        queries_search = queries
    # ---- prepare index configuration ---- #
    config = BuildConfiguration(
        [algorithms[algo] for algo in clustering_algorithms],
        epochs,
        model_type,
        lr,
        n_categories,
    )
    # ---- instantiate the index builder ---- #
    builder = LearnedIndexBuilder(data_pd, config)
    # ---- build the index ---- #
    li, data_prediction, n_buckets_in_index, build_t, cluster_t = builder.build()
    LOG.info(f"Total number of buckets in the index: {n_buckets_in_index}")
    LOG.info(f"Cluster time: {cluster_t}")
    LOG.info(f"Pure build time: {build_t}")
    LOG.info(f"Overall build time: {time.time() - s}")

    if save:
        if not os.path.isdir(MODELS_DIR_NAME):
            os.mkdir(MODELS_DIR_NAME)
        save_filename = format_models_filename(
            kind, config, clustering_algorithms, preprocess, size
        )
        LOG.info(f"Saving as {save_filename}")
        save_as_pickle(f"{save_filename}.pkl", li)

    n_buckets = [int((p / 100) * n_buckets_in_index) for p in n_buckets_perc]
    n_buckets = list(set([b for b in n_buckets if b > 0]))
    LOG.info(f"Number of buckets to search in: {n_buckets}")

    for bucket in n_buckets:
        LOG.info(f"Searching with {bucket} buckets")

        dists, nns, measured_time = li.search(
            data_navigation=data_pd,
            queries_navigation=queries,
            data_search=data_search,
            queries_search=queries_search,
            data_prediction=data_prediction,
            n_categories=n_categories,
            n_buckets=bucket,
            k=k,
        )

        LOG.info(f"Inference time: {measured_time['inference']}")
        LOG.info(f"Search time: {measured_time['search']}")
        LOG.info(
            f"Search within buckets time: {measured_time['search_within_buckets']}"
        )
        LOG.info(f"Sequential search time: {measured_time['seq_search']}")
        LOG.info(f"Sort time: {measured_time['sort']}")

        short_identifier = "learned-index"
        identifier = format_identifier(
            bucket, kind, config, clustering_algorithms, short_identifier, size
        )
        store_results(
            os.path.join("result/", kind, size, f"{identifier}.h5"),
            short_identifier.capitalize(),
            kind,
            dists,
            nns,
            build_t,
            measured_time["search"],
            identifier,
            size,
        )


def expand(array: List[Any], size: int):
    """Expands an array of size one to the given size."""
    assert len(array) == 1
    return [array[0]] * size


def validate_and_expand_per_level_arguments(args: Dict[str, Any]):
    """
    Validates that the arguments that are per-level are either of size 1
    and expanded to the size of n_categories or of the same size as n_categories.
    """
    arguments = ["clustering_algorithm", "model_type", "epochs", "lr", "n_categories"]

    for arg in arguments:
        if len(args[arg]) == 1:
            args[arg] = expand(args[arg], len(args["n_categories"]))
        else:
            assert len(args[arg]) == len(args["n_categories"])


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", default="pca96v2")
    parser.add_argument("--emb", default="pca96")
    parser.add_argument(
        "--size", default="100K", choices=["100K", "300K", "10M", "30M", "100M"]
    )
    parser.add_argument("--k", default=10, type=int)
    parser.add_argument("--n-categories", nargs="+", default=[10, 10], type=int)
    parser.add_argument("--epochs", nargs="+", default=[100], type=int)
    parser.add_argument("--model-type", nargs="+", default=["MLP"])
    parser.add_argument("--lr", nargs="+", default=[0.01], type=float)
    parser.add_argument("-b", "--n-buckets", nargs="+", default=[2, 3, 4], type=int)
    parser.add_argument("-bp", "--buckets-perc", nargs="+", default=[10], type=int)
    parser.add_argument("--preprocess", default=True, type=bool)
    parser.add_argument("--save", default=True, type=bool)
    parser.add_argument(
        "--clustering-algorithm",
        nargs="+",
        default=["faiss_kmeans"],
        choices=algorithms.keys(),
    )
    args = parser.parse_args()

    validate_and_expand_per_level_arguments(vars(args))

    if args.save and "PBS_JOBID" not in os.environ:
        os.environ["PBS_JOBID"] = "unknown"

    run(
        args.dataset,
        args.emb,
        args.size,
        args.k,
        "learned-index",
        args.buckets_perc,
        args.n_categories,
        args.epochs,
        args.model_type,
        args.lr,
        args.preprocess,
        args.save,
        args.clustering_algorithm,
    )

search/search.py
LearnedMetricIndex/LearnedMetricIndex
from glob import glob
from os.path import basename, splitext
from setuptools import find_packages, setup

setup(
    name="li",
    version="0.1",
    packages=find_packages(where="search"),
    package_dir={"": "search"},
    py_modules=[splitext(basename(path))[0] for path in glob("search/*.py")],
)

setup.py
LearnedMetricIndex/LearnedMetricIndex