FastPlaid

Use search.FastPlaid to initialize the engine, .create() to build an index from multi-vector embeddings, and .search() to retrieve results.

Device Selection:

• Pass device="cpu" to speed up initialization on CPU.
• Pass device="cuda" for GPU acceleration. If left blank, it auto-detects (including multi-GPU support).
• On GPU, index placement and batch sizes are automatically tuned based on available VRAM.

Output Format: .search() returns a list of lists, where each inner list contains tuples of (document_index, similarity_score) for the top top_k results per query.

import torch
from fast_plaid import search

# Initialize engine
fast_plaid = search.FastPlaid(index="index", device="cpu")

embedding_dim = 128

# Index 100 documents (300 tokens each, 128-dim)
fast_plaid.create(
    documents_embeddings=[torch.randn(300, embedding_dim) for _ in range(100)]
)

# Search for 2 queries (50 tokens each, 128-dim)
scores = fast_plaid.search(
    queries_embeddings=torch.randn(2, 50, embedding_dim),
    top_k=10,
)

print(scores)

FastPlaid features an optional metadata filtering engine powered by SQLite. When you provide a metadata parameter during .create() or .update() calls, FastPlaid stores this information in a searchable database within your index directory.

To filter search results based on metadata, use fast_plaid.filtering.where() to execute standard SQL conditions. This returns a list of embedding IDs (a subset) that you can pass to the subset parameter of the .search() method to pre-filter your vector search.

Key functions:

• fast_plaid.filtering.where(index, condition, parameters): Returns a subset of IDs matching the SQL condition.
• fast_plaid.search(..., subset=subset): Performs vector search restricted to the provided IDs.
• fast_plaid.filtering.get(index, subset): Retrieves the metadata for the specified subset of documents.

from datetime import date
import torch
from fast_plaid import filtering, search

# 1. Initialize the FastPlaid index
fast_plaid = search.FastPlaid(index="metadata_index")
embedding_dim = 128

# 2. Create initial documents with metadata
initial_embeddings = [torch.randn(10, embedding_dim) for _ in range(3)]
initial_metadata = [
    {"name": "Alice", "category": "A", "join_date": date(2023, 5, 17)},
    {"name": "Bob", "category": "B", "join_date": date(2021, 6, 21)},
    {"name": "Alex", "category": "A", "join_date": date(2023, 8, 1)},
]
fast_plaid.create(documents_embeddings=initial_embeddings, metadata=initial_metadata)

# 3. Update the index with new documents and metadata
new_embeddings = [torch.randn(10, embedding_dim) for _ in range(2)]
new_metadata = [
    {"name": "Charlie", "category": "B", "join_date": date(2020, 3, 15)},
    {
        "name": "Amanda",
        "category": "A",
        "join_date": date(2024, 1, 10),
        "status": "active",
    },
]
fast_plaid.update(documents_embeddings=new_embeddings, metadata=new_metadata)

# 4. Use filtering.where to get the corresponding rows of the FastPlaid index
subset = filtering.where(
    index="metadata_index",
    condition="name LIKE ? AND join_date > ?",
    parameters=("A%", "2022-12-31"),
)

# 5. Perform a search restricted to the filtered subset
query_embedding = torch.randn(1, 5, embedding_dim)
scores = fast_plaid.search(queries_embeddings=query_embedding, top_k=3, subset=subset)

print("Search results within the subset:")
print(scores)

# 6. Access to the metadata of the retrieved documents
for match in scores:
    print("Metadata of matched documents:")
    print(filtering.get(index="metadata_index", subset=[subset for subset, _ in match]))

The FastPlaid constructor accepts several parameters to tune performance and memory:

• index: (str) File path to the directory where the index is saved/loaded.
• device: (str | list[str] | None) Specifies computation devices (e.g., "cuda:0", "cpu", or ["cuda:0", "cuda:1"]). If multiple GPUs are provided, multiprocessing is used. Note: Wrap your code in if __name__ == "__main__": for proper multiprocessing.
• index_gpu_memory: (Literal["auto", "low", "medium", "high"]) Controls where document tensors live on GPU:
  • "low": Codes and residuals stay on CPU. Lowest VRAM, slowest search.
  • "medium": Codes on GPU, residuals on CPU.
  • "high": Both codes and residuals on GPU. Fastest search, highest VRAM.
  • "auto": Picks the highest tier within index_memory_fraction.
• index_memory_fraction: (float) Fraction of total device memory for "auto" placement (default 0.7).
• search_memory_fraction: (float) Fraction of free VRAM used for search chunk planning (default 0.5).

The index_gpu_memory flag in the search.FastPlaid constructor controls where large document tensors reside during query time. This affects search speed but does not change results. This setting has no effect when device="cpu".

• "auto" (default): Adaptive placement based on available VRAM. It picks the highest tier that fits within the index_memory_fraction (default 0.7).
• "high": Fastest option. Both codes and residuals live on the GPU, avoiding host-to-device copies.
• "medium": Codes are on GPU, but residuals (the larger tensor) stay on CPU.
• "low": Minimal VRAM usage. Both codes and residuals stay on CPU.

Migration Note: index_gpu_memory replaces the deprecated low_memory flag. low_memory=True is equivalent to index_gpu_memory="low", and low_memory=False is equivalent to index_gpu_memory="high".

# Default, adapts to each device
fast_plaid = search.FastPlaid(index="index")

# Force the index onto the GPU for maximum QPS
fast_plaid = search.FastPlaid(index="index", index_gpu_memory="high")

# Keep the index off the GPU, for instance to share the device with a model
fast_plaid = search.FastPlaid(index="index", index_gpu_memory="low")

Query the index using query embeddings to retrieve the most relevant documents.

Parameters:

• queries_embeddings: (torch.Tensor | list[torch.Tensor]) Shape (num_queries, num_tokens, dim) or a list of such tensors.
• top_k: (int) Number of documents to retrieve per query (default 10).
• batch_size: (int | "auto") Number of documents scored per chunk. "auto" uses search_memory_fraction to plan chunks.
• n_full_scores: (int) Number of candidate documents for full re-ranking. Higher values increase accuracy but increase computation (default 4096).
• n_ivf_probe: (int) Number of clusters to search (default 8). Higher values improve recall but increase time.
• subset: (list[list[int]] | list[int]) Restricts search to specific document IDs. IDs correspond to insertion order (starting at 0).

# Perform a search
queries = torch.randn(2, 50, 128)
results = fp.search(
    queries_embeddings=queries,
    top_k=5,
    n_full_scores=2048
)
# Returns: list[list[tuple[int, float]]] -> [[(doc_id, score), ...], ...]