Skip to main content
`Ctrl`+`K`

# ragoon.similarity_search

Classes

`SimilaritySearch`(model_name[, device, ndim, ...]) | A class dedicated to encoding text data, quantizing embeddings, and managing indices for efficient similarity search.

class ragoon.similarity_search.SimilaritySearch(model_name: str, device: str = 'cuda', ndim: int = 1024, metric: str = 'ip', dtype: str = 'i8')

  Bases: `object`

  A class dedicated to encoding text data, quantizing embeddings, and managing indices for efficient similarity search.

  model_name

    Name or identifier of the embedding model.

    Type:

      str

  device

    Computation device (‘cpu’ or ‘cuda’).

    Type:

      str

  ndim

    Dimension of the embeddings.

    Type:

      int

  metric

    Metric used for the index (‘ip’ for inner product, etc.).

    Type:

      str

  dtype

    Data type for the index (‘i8’ for int8, etc.).

    Type:

      str

  encode(corpus, normalize_embeddings=True)

    Encodes a list of text data into embeddings.

  quantize_embeddings(embeddings, quantization_type)

    Quantizes the embeddings for efficient storage and search.

  create_faiss_index(ubinary_embeddings, index_path)

    Creates and saves a FAISS binary index.

  create_usearch_index(int8_embeddings, index_path)

    Creates and saves a USEARCH integer index.

  load_usearch_index_view(index_path)

    Loads a USEARCH index as a view for memory-efficient operations.

  load_faiss_index(index_path)

    Loads a FAISS binary index for searching.

  search(query, top_k=10, rescore_multiplier=4)

    Performs a search operation against the indexed embeddings.

  Examples

      >>> instance = SimilaritySearch(
      ...    model_name="louisbrulenaudet/tsdae-lemone-mbert-base",
      ...    device="cuda",
      ...    ndim=768,
      ...    metric="ip",
      ...    dtype="i8"
      )
      >>> embeddings = instance.encode(corpus=dataset["output"])
      >>> ubinary_embeddings = instance.quantize_embeddings(
      ...    embeddings=embeddings,
      ...    quantization_type="ubinary"
      )
      >>> int8_embeddings = instance.quantize_embeddings(
      ...    embeddings=embeddings,
      ...    quantization_type="int8"
      )
      >>> instance.create_usearch_index(
      ...    int8_embeddings=int8_embeddings,
      ...    index_path="./usearch_int8.index"
      )
      >>> instance.create_faiss_index(
      ...    ubinary_embeddings=ubinary_embeddings,
      ...    index_path="./faiss_ubinary.index"
      )
      >>> top_k_scores, top_k_indices = instance.search(
      ...    query="Sont considérées comme ayant leur domicile fiscal en France au sens de l'article 4 A",
      ...    top_k=10,
      ...    rescore_multiplier=4
      )

  __init__(model_name: str, device: str = 'cuda', ndim: int = 1024, metric: str = 'ip', dtype: str = 'i8')

    Initializes the EmbeddingIndexer with the specified model, device, and index configurations.

    Parameters:

      - model_name (str) – The name or identifier of the SentenceTransformer model to use for embedding.

      - device (str, optional) – The computation device to use (‘cpu’ or ‘cuda’). Default is ‘cuda’.

      - ndim (int, optional) – The dimensionality of the embeddings. Default is 1024.

      - metric (str, optional) – The metric used for the index (‘ip’ for inner product). Default is ‘ip’.

      - dtype (str, optional) – The data type for the USEARCH index (‘i8’ for 8-bit integer). Default is ‘i8’.

  encode(corpus: list, normalize_embeddings: bool = True) → ndarray

    Encodes the given corpus into full-precision embeddings.

    Parameters:

      - corpus (list) – A list of sentences to be encoded.

      - normalize_embeddings (bool, optional) – Whether to normalize returned vectors to have length 1. In that case, the faster dot-product (util.dot_score) instead of cosine similarity can be used. Default is True.

    Returns:

      The full-precision embeddings of the corpus.

    Return type:

      np.ndarray

    Notes

    This method normalizes the embeddings and shows the progress bar during the encoding process.

  quantize_embeddings(embeddings: ndarray, quantization_type: str) → ndarray | bytearray

    Quantizes the given embeddings based on the specified quantization type (‘ubinary’ or ‘int8’).

    Parameters:

      - embeddings (np.ndarray) – The full-precision embeddings to be quantized.

      - quantization_type (str) – The type of quantization (‘ubinary’ for unsigned binary, ‘int8’ for 8-bit integers).

    Returns:

      The quantized embeddings.

    Return type:

      Union[np.ndarray, bytearray]

    Raises:

      ValueError – If an unsupported quantization type is provided.

  create_faiss_index(ubinary_embeddings: bytearray, index_path: str = None, save: bool = False) → None

    Creates and saves a FAISS binary index from ubinary embeddings.

    Parameters:

      - ubinary_embeddings (bytearray) – The ubinary-quantized embeddings.

      - index_path (str, optional) – The file path to save the FAISS binary index. Default is None.

      - save (bool, optional) – Indicator for saving the index. Default is False.

    Notes

    The dimensionality of the index is specified during the class initialization (default is 1024).

  create_usearch_index(int8_embeddings: ndarray, index_path: str = None, save: bool = False) → None

    Creates and saves a USEARCH integer index from int8 embeddings.

    Parameters:

      - int8_embeddings (np.ndarray) – The int8-quantized embeddings.

      - index_path (str, optional) – The file path to save the USEARCH integer index. Default is None.

      - save (bool, optional) – Indicator for saving the index. Default is False.

    Return type:

      None

    Notes

    The dimensionality and metric of the index are specified during class initialization.

  load_usearch_index_view(index_path: str) → any

    Loads a USEARCH index as a view for memory-efficient operations.

    Parameters:

      index_path (str) – The file path to the USEARCH index to be loaded as a view.

    Returns:

      A view of the USEARCH index for memory-efficient similarity search operations.

    Return type:

      object

    Notes

    Implementing this would depend on the specific USEARCH index handling library being used.

  load_faiss_index(index_path: str) → None

    Loads a FAISS binary index from a specified file path.

    This method loads a binary index created by FAISS into the class attribute binary_index, ready for performing similarity searches.

    Parameters:

      index_path (str) – The file path to the saved FAISS binary index.

    Return type:

      None

    Notes

    The loaded index is stored in the binary_index attribute of the class. Ensure that the index at index_path is compatible with the configurations (e.g., dimensions) used for this class instance.

  search(query: str, top_k: int = 10, rescore_multiplier: int = 4) → Tuple[List[float], List[int]]

    Performs a search operation against the indexed embeddings.

    Parameters:

      - query (str) – The query sentence/string to be searched.

      - top_k (int, optional) – The number of top results to return.

      - rescore_multiplier (int, optional) – The multiplier used to increase the initial retrieval size for re-scoring. Higher values can increase precision at the cost of performance.

    Returns:

      A tuple containing the scores and the indices of the top k results.

    Return type:

      Tuple[List[float], List[int]]

    Notes

    This method assumes that binary_index and int8_index are already loaded or created.