quantization.py

import logging
import time
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union

import numpy as np
from torch import Tensor

logger = logging.getLogger(__name__)


if TYPE_CHECKING:
    import faiss
    import usearch


def semantic_search_faiss(
    query_embeddings: np.ndarray,
    corpus_embeddings: Optional[np.ndarray] = None,
    corpus_index: Optional["faiss.Index"] = None,
    corpus_precision: Literal["float32", "uint8", "ubinary"] = "float32",
    top_k: int = 10,
    ranges: Optional[np.ndarray] = None,
    calibration_embeddings: Optional[np.ndarray] = None,
    rescore: bool = True,
    rescore_multiplier: int = 2,
    exact: bool = True,
    output_index: bool = False,
) -> Tuple[List[List[Dict[str, Union[int, float]]]], float, "faiss.Index"]:
    """
    Performs semantic search using the FAISS library.

    Rescoring will be performed if:
    1. `rescore` is True
    2. The query embeddings are not quantized
    3. The corpus is quantized, i.e. the corpus precision is not float32
    Only if these conditions are true, will we search for `top_k * rescore_multiplier` samples and then rescore to only
    keep `top_k`.

    Args:
        query_embeddings: Embeddings of the query sentences. Ideally not
            quantized to allow for rescoring.
        corpus_embeddings: Embeddings of the corpus sentences. Either
            `corpus_embeddings` or `corpus_index` should be used, not
            both. The embeddings can be quantized to "int8" or "binary"
            for more efficient search.
        corpus_index: FAISS index for the corpus sentences. Either
            `corpus_embeddings` or `corpus_index` should be used, not
            both.
        corpus_precision: Precision of the corpus embeddings. The
            options are "float32", "int8", or "binary". Default is
            "float32".
        top_k: Number of top results to retrieve. Default is 10.
        ranges: Ranges for quantization of embeddings. This is only used
            for int8 quantization, where the ranges refers to the
            minimum and maximum values for each dimension. So, it's a 2D
            array with shape (2, embedding_dim). Default is None, which
            means that the ranges will be calculated from the
            calibration embeddings.
        calibration_embeddings: Embeddings used for calibration during
            quantization. This is only used for int8 quantization, where
            the calibration embeddings can be used to compute ranges,
            i.e. the minimum and maximum values for each dimension.
            Default is None, which means that the ranges will be
            calculated from the query embeddings. This is not
            recommended.
        rescore: Whether to perform rescoring. Note that rescoring still
            will only be used if the query embeddings are not quantized
            and the corpus is quantized, i.e. the corpus precision is
            not "float32". Default is True.
        rescore_multiplier: Oversampling factor for rescoring. The code
            will now search `top_k * rescore_multiplier` samples and
            then rescore to only keep `top_k`. Default is 2.
        exact: Whether to use exact search or approximate search.
            Default is True.
        output_index: Whether to output the FAISS index used for the
            search. Default is False.

    Returns:
        A tuple containing a list of search results and the time taken
        for the search. If `output_index` is True, the tuple will also
        contain the FAISS index used for the search.

    Raises:
        ValueError: If both `corpus_embeddings` and `corpus_index` are
            provided or if neither is provided.

    The list of search results is in the format: [[{"corpus_id": int, "score": float}, ...], ...]
    The time taken for the search is a float value.
    """
    import faiss

    if corpus_embeddings is not None and corpus_index is not None:
        raise ValueError("Only corpus_embeddings or corpus_index should be used, not both.")
    if corpus_embeddings is None and corpus_index is None:
        raise ValueError("Either corpus_embeddings or corpus_index should be used.")

    # If corpus_index is not provided, create a new index
    if corpus_index is None:
        if corpus_precision in ("float32", "uint8"):
            if exact:
                corpus_index = faiss.IndexFlatIP(corpus_embeddings.shape[1])
            else:
                corpus_index = faiss.IndexHNSWFlat(corpus_embeddings.shape[1], 16)

        elif corpus_precision == "ubinary":
            if exact:
                corpus_index = faiss.IndexBinaryFlat(corpus_embeddings.shape[1] * 8)
            else:
                corpus_index = faiss.IndexBinaryHNSW(corpus_embeddings.shape[1] * 8, 16)

        corpus_index.add(corpus_embeddings)

    # If rescoring is enabled and the query embeddings are in float32, we need to quantize them
    # to the same precision as the corpus embeddings. Also update the top_k value to account for the
    # rescore_multiplier
    rescore_embeddings = None
    k = top_k
    if query_embeddings.dtype not in (np.uint8, np.int8):
        if rescore:
            if corpus_precision != "float32":
                rescore_embeddings = query_embeddings
                k *= rescore_multiplier
            else:
                logger.warning(
                    "Rescoring is enabled but the corpus is not quantized. Either pass `rescore=False` or "
                    'quantize the corpus embeddings with `quantize_embeddings(embeddings, precision="...") `'
                    'and pass `corpus_precision="..."` to `semantic_search_faiss`.'
                )

        query_embeddings = quantize_embeddings(
            query_embeddings,
            precision=corpus_precision,
            ranges=ranges,
            calibration_embeddings=calibration_embeddings,
        )
    elif rescore:
        logger.warning(
            "Rescoring is enabled but the query embeddings are quantized. Either pass `rescore=False` or don't quantize the query embeddings."
        )

    # Perform the search using the usearch index
    start_t = time.time()
    scores, indices = corpus_index.search(query_embeddings, k)

    # If rescoring is enabled, we need to rescore the results using the rescore_embeddings
    if rescore_embeddings is not None:
        top_k_embeddings = np.array(
            [[corpus_index.reconstruct(idx.item()) for idx in query_indices] for query_indices in indices]
        )
        # If the corpus precision is binary, we need to unpack the bits
        if corpus_precision == "ubinary":
            top_k_embeddings = np.unpackbits(top_k_embeddings, axis=-1).astype(int)
        else:
            top_k_embeddings = top_k_embeddings.astype(int)

        # rescore_embeddings: [num_queries, embedding_dim]
        # top_k_embeddings: [num_queries, top_k, embedding_dim]
        # updated_scores: [num_queries, top_k]
        # We use einsum to calculate the dot product between the query and the top_k embeddings, equivalent to looping
        # over the queries and calculating 'rescore_embeddings[i] @ top_k_embeddings[i].T'
        rescored_scores = np.einsum("ij,ikj->ik", rescore_embeddings, top_k_embeddings)
        rescored_indices = np.argsort(-rescored_scores)[:, :top_k]
        indices = indices[np.arange(len(query_embeddings))[:, None], rescored_indices]
        scores = rescored_scores[np.arange(len(query_embeddings))[:, None], rescored_indices]

    delta_t = time.time() - start_t

    outputs = (
        [
            [
                {"corpus_id": int(neighbor), "score": float(score)}
                for score, neighbor in zip(scores[query_id], indices[query_id])
            ]
            for query_id in range(len(query_embeddings))
        ],
        delta_t,
    )
    if output_index:
        outputs = (*outputs, corpus_index)
    return outputs


def semantic_search_usearch(
    query_embeddings: np.ndarray,
    corpus_embeddings: Optional[np.ndarray] = None,
    corpus_index: Optional["usearch.index.Index"] = None,
    corpus_precision: Literal["float32", "int8", "binary"] = "float32",
    top_k: int = 10,
    ranges: Optional[np.ndarray] = None,
    calibration_embeddings: Optional[np.ndarray] = None,
    rescore: bool = True,
    rescore_multiplier: int = 2,
    exact: bool = True,
    output_index: bool = False,
) -> Tuple[List[List[Dict[str, Union[int, float]]]], float, "usearch.index.Index"]:
    """
    Performs semantic search using the usearch library.

    Rescoring will be performed if:
    1. `rescore` is True
    2. The query embeddings are not quantized
    3. The corpus is quantized, i.e. the corpus precision is not float32
    Only if these conditions are true, will we search for `top_k * rescore_multiplier` samples and then rescore to only
    keep `top_k`.

    Args:
        query_embeddings: Embeddings of the query sentences. Ideally not
            quantized to allow for rescoring.
        corpus_embeddings: Embeddings of the corpus sentences. Either
            `corpus_embeddings` or `corpus_index` should be used, not
            both. The embeddings can be quantized to "int8" or "binary"
            for more efficient search.
        corpus_index: usearch index for the corpus sentences. Either
            `corpus_embeddings` or `corpus_index` should be used, not
            both.
        corpus_precision: Precision of the corpus embeddings. The
            options are "float32", "int8", or "binary". Default is
            "float32".
        top_k: Number of top results to retrieve. Default is 10.
        ranges: Ranges for quantization of embeddings. This is only used
            for int8 quantization, where the ranges refers to the
            minimum and maximum values for each dimension. So, it's a 2D
            array with shape (2, embedding_dim). Default is None, which
            means that the ranges will be calculated from the
            calibration embeddings.
        calibration_embeddings: Embeddings used for calibration during
            quantization. This is only used for int8 quantization, where
            the calibration embeddings can be used to compute ranges,
            i.e. the minimum and maximum values for each dimension.
            Default is None, which means that the ranges will be
            calculated from the query embeddings. This is not
            recommended.
        rescore: Whether to perform rescoring. Note that rescoring still
            will only be used if the query embeddings are not quantized
            and the corpus is quantized, i.e. the corpus precision is
            not "float32". Default is True.
        rescore_multiplier: Oversampling factor for rescoring. The code
            will now search `top_k * rescore_multiplier` samples and
            then rescore to only keep `top_k`. Default is 2.
        exact: Whether to use exact search or approximate search.
            Default is True.
        output_index: Whether to output the usearch index used for the
            search. Default is False.

    Returns:
        A tuple containing a list of search results and the time taken
        for the search. If `output_index` is True, the tuple will also
        contain the usearch index used for the search.

    Raises:
        ValueError: If both `corpus_embeddings` and `corpus_index` are
            provided or if neither is provided.

    The list of search results is in the format: [[{"corpus_id": int, "score": float}, ...], ...]
    The time taken for the search is a float value.
    """
    from usearch.compiled import ScalarKind
    from usearch.index import Index

    if corpus_embeddings is not None and corpus_index is not None:
        raise ValueError("Only corpus_embeddings or corpus_index should be used, not both.")
    if corpus_embeddings is None and corpus_index is None:
        raise ValueError("Either corpus_embeddings or corpus_index should be used.")
    if corpus_precision not in ["float32", "int8", "binary"]:
        raise ValueError('corpus_precision must be "float32", "int8", or "binary" for usearch')

    # If corpus_index is not provided, create a new index
    if corpus_index is None:
        if corpus_precision == "float32":
            corpus_index = Index(
                ndim=corpus_embeddings.shape[1],
                metric="cos",
                dtype="f32",
            )
        elif corpus_precision == "int8":
            corpus_index = Index(
                ndim=corpus_embeddings.shape[1],
                metric="ip",
                dtype="i8",
            )
        elif corpus_precision == "binary":
            corpus_index = Index(
                ndim=corpus_embeddings.shape[1],
                metric="hamming",
                dtype="b1",
            )
        corpus_index.add(np.arange(len(corpus_embeddings)), corpus_embeddings)

    # If rescoring is enabled and the query embeddings are in float32, we need to quantize them
    # to the same precision as the corpus embeddings. Also update the top_k value to account for the
    # rescore_multiplier
    rescore_embeddings = None
    k = top_k
    if query_embeddings.dtype not in (np.uint8, np.int8):
        if rescore:
            if corpus_index.dtype != ScalarKind.F32:
                rescore_embeddings = query_embeddings
                k *= rescore_multiplier
            else:
                logger.warning(
                    "Rescoring is enabled but the corpus is not quantized. Either pass `rescore=False` or "
                    'quantize the corpus embeddings with `quantize_embeddings(embeddings, precision="...") `'
                    'and pass `corpus_precision="..."` to `semantic_search_usearch`.'
                )

        query_embeddings = quantize_embeddings(
            query_embeddings,
            precision=corpus_precision,
            ranges=ranges,
            calibration_embeddings=calibration_embeddings,
        )
    elif rescore:
        logger.warning(
            "Rescoring is enabled but the query embeddings are quantized. Either pass `rescore=False` or don't quantize the query embeddings."
        )

    # Perform the search using the usearch index
    start_t = time.time()
    matches = corpus_index.search(query_embeddings, count=k, exact=exact)
    scores = matches.distances
    indices = matches.keys

    if scores.ndim < 2:
        scores = np.atleast_2d(scores)
    if indices.ndim < 2:
        indices = np.atleast_2d(indices)

    # If rescoring is enabled, we need to rescore the results using the rescore_embeddings
    if rescore_embeddings is not None:
        top_k_embeddings = np.array([corpus_index.get(query_indices) for query_indices in indices])
        # If the corpus precision is binary, we need to unpack the bits
        if corpus_precision == "binary":
            top_k_embeddings = np.unpackbits(top_k_embeddings.astype(np.uint8), axis=-1)
        top_k_embeddings = top_k_embeddings.astype(int)

        # rescore_embeddings: [num_queries, embedding_dim]
        # top_k_embeddings: [num_queries, top_k, embedding_dim]
        # updated_scores: [num_queries, top_k]
        # We use einsum to calculate the dot product between the query and the top_k embeddings, equivalent to looping
        # over the queries and calculating 'rescore_embeddings[i] @ top_k_embeddings[i].T'
        rescored_scores = np.einsum("ij,ikj->ik", rescore_embeddings, top_k_embeddings)
        rescored_indices = np.argsort(-rescored_scores)[:, :top_k]
        indices = indices[np.arange(len(query_embeddings))[:, None], rescored_indices]
        scores = rescored_scores[np.arange(len(query_embeddings))[:, None], rescored_indices]

    delta_t = time.time() - start_t

    outputs = (
        [
            [
                {"corpus_id": int(neighbor), "score": float(score)}
                for score, neighbor in zip(scores[query_id], indices[query_id])
            ]
            for query_id in range(len(query_embeddings))
        ],
        delta_t,
    )
    if output_index:
        outputs = (*outputs, corpus_index)
    return outputs


def quantize_embeddings(
    embeddings: Union[Tensor, np.ndarray],
    precision: Literal["float32", "int8", "uint8", "binary", "ubinary"],
    ranges: Optional[np.ndarray] = None,
    calibration_embeddings: Optional[np.ndarray] = None,
) -> np.ndarray:
    """
    Quantizes embeddings to a lower precision. This can be used to reduce the memory footprint and increase the
    speed of similarity search. The supported precisions are "float32", "int8", "uint8", "binary", and "ubinary".

    Args:
        embeddings: Unquantized (e.g. float) embeddings with to quantize
            to a given precision
        precision: The precision to convert to. Options are "float32",
            "int8", "uint8", "binary", "ubinary".
        ranges (Optional[np.ndarray]): Ranges for quantization of
            embeddings. This is only used for int8 quantization, where
            the ranges refers to the minimum and maximum values for each
            dimension. So, it's a 2D array with shape (2,
            embedding_dim). Default is None, which means that the ranges
            will be calculated from the calibration embeddings.
        calibration_embeddings (Optional[np.ndarray]): Embeddings used
            for calibration during quantization. This is only used for
            int8 quantization, where the calibration embeddings can be
            used to compute ranges, i.e. the minimum and maximum values
            for each dimension. Default is None, which means that the
            ranges will be calculated from the query embeddings. This is
            not recommended.

    Returns:
        Quantized embeddings with the specified precision
    """
    if isinstance(embeddings, Tensor):
        embeddings = embeddings.cpu().numpy()
    elif isinstance(embeddings, list):
        if isinstance(embeddings[0], Tensor):
            embeddings = [embedding.cpu().numpy() for embedding in embeddings]
        embeddings = np.array(embeddings)
    if embeddings.dtype in (np.uint8, np.int8):
        raise Exception("Embeddings to quantize must be float rather than int8 or uint8.")

    if precision == "float32":
        return embeddings.astype(np.float32)

    if precision.endswith("int8"):
        # Either use the 1. provided ranges, 2. the calibration dataset or 3. the provided embeddings
        if ranges is None:
            if calibration_embeddings is not None:
                ranges = np.vstack((np.min(calibration_embeddings, axis=0), np.max(calibration_embeddings, axis=0)))
            else:
                if embeddings.shape[0] < 100:
                    logger.warning(
                        f"Computing {precision} quantization buckets based on {len(embeddings)} embedding{'s' if len(embeddings) != 1 else ''}."
                        f" {precision} quantization is more stable with `ranges` calculated from more embeddings "
                        "or a `calibration_embeddings` that can be used to calculate the buckets."
                    )
                ranges = np.vstack((np.min(embeddings, axis=0), np.max(embeddings, axis=0)))
        starts = ranges[0, :]
        steps = (ranges[1, :] - ranges[0, :]) / 255

        if precision == "uint8":
            return ((embeddings - starts) / steps).astype(np.uint8)
        elif precision == "int8":
            return ((embeddings - starts) / steps - 128).astype(np.int8)

    if precision == "binary":
        return (np.packbits(embeddings > 0).reshape(embeddings.shape[0], -1) - 128).astype(np.int8)

    if precision == "ubinary":
        return np.packbits(embeddings > 0).reshape(embeddings.shape[0], -1)

    raise ValueError(f"Precision {precision} is not supported")