torch-similarity-search 0.0.4

PyTorch-native similarity search: convert FAISS indexes to nn.Module

torch-similarity-search

PyTorch-native similarity search. Convert trained FAISS indexes to pure nn.Module models for GPU inference.

Train with FAISS, deploy with PyTorch.

Why?

• No numpy overhead - FAISS requires numpy conversion; this library keeps tensors on GPU
• TorchScript export - Deploy without FAISS dependency, load with just torch.jit.load()
• GPU memory sharing - Index vectors stay in GPU memory alongside your embedding model
• Triton Inference Server ready - Export once, serve anywhere

Installation

pip install torch-similarity-search

For FAISS conversion support:

pip install torch-similarity-search faiss-cpu  # or faiss-gpu

Quick Start

Convert from FAISS

import faiss
import torch
import torch_similarity_search as tss

# Train with FAISS (your existing workflow)
quantizer = faiss.IndexFlatL2(128)
index = faiss.IndexIVFFlat(quantizer, 128, 100)
index.train(vectors)
index.add(vectors)

# Convert to PyTorch
model = tss.from_faiss(index)
model = model.cuda()
model.nprobe = 10

# Search with PyTorch tensors (no numpy!)
queries = torch.randn(32, 128, device="cuda")
distances, indices = model.search(queries, k=10)

Build from Scratch

import torch
from torch_similarity_search import IVFFlatIndex

# Create and train
index = IVFFlatIndex(dim=128, nlist=100, metric="l2")
training_vectors = torch.randn(10000, 128)
index.train(training_vectors)
index.add(training_vectors)

# Move to GPU
index = index.cuda()

# Search
queries = torch.randn(32, 128, device="cuda")
distances, indices = index.search(queries, k=10)

Export for Production

# Export to TorchScript (no torch_similarity_search needed to load!)
scripted = torch.jit.script(model)
scripted.save("index.pt")

# Load anywhere - just needs PyTorch
model = torch.jit.load("index.pt")
model = model.cuda()
distances, indices = model.search(queries, k=10)

Use with Embedding Models

# End-to-end GPU inference
class SearchModel(torch.nn.Module):
    def __init__(self, encoder, index):
        super().__init__()
        self.encoder = encoder
        self.index = index

    def forward(self, text_embeddings):
        # Everything stays on GPU
        return self.index.search(text_embeddings, k=10)

# Export the complete pipeline
model = SearchModel(encoder, index)
torch.jit.script(model).save("search_pipeline.pt")

Supported Index Types

FAISS Index	PyTorch Module	Status
IndexFlat	FlatIndex	✅ Supported
IndexIVFFlat	IVFFlatIndex	✅ Supported
IndexIVFPQ	IVFPQIndex	✅ Supported

API Reference

FlatIndex

Brute-force exact search - compares against all vectors. Best for small datasets or exact results.

from torch_similarity_search import FlatIndex

index = FlatIndex(
    dim=128,          # Vector dimensionality
    metric="l2",      # Distance metric: "l2", "ip" (inner product), or "cosine"
    k=10,             # Default k for forward() method
)

index.add(vectors)    # No training required
distances, indices = index.search(queries, k=10)

IVFFlatIndex

Inverted File Flat index - partitions vectors into clusters for fast approximate search.

from torch_similarity_search import IVFFlatIndex

index = IVFFlatIndex(
    dim=128,          # Vector dimensionality
    nlist=100,        # Number of clusters (higher = faster but less accurate)
    metric="l2",      # Distance metric: "l2", "ip" (inner product), or "cosine"
    nprobe=10,        # Clusters to search at query time
    k=10,             # Default k for forward() method
)

index.train(vectors)  # Train centroids first
index.add(vectors)
distances, indices = index.search(queries, k=10)

Common Methods (both index types):

Method	Description
add(vectors)	Add vectors to index. Accepts (n, dim) or (dim,) tensors.
search(queries, k)	Find k nearest neighbors. Returns (distances, indices) tensors.
forward(queries)	Same as search() but uses configured k. For TorchScript export.

IVFFlatIndex-specific:

Method/Property	Description
train(vectors)	Train cluster centroids via k-means. Requires n >= nlist.
nprobe	Clusters to probe during search (settable, higher = more accurate)
is_trained	Whether index has been trained

IVFPQIndex

Inverted File with Product Quantization - combines clustering with vector compression for memory-efficient approximate search. Best for large datasets where memory is a concern.

from torch_similarity_search import IVFPQIndex

index = IVFPQIndex(
    dim=128,          # Vector dimensionality (must be divisible by M)
    nlist=100,        # Number of IVF clusters
    M=8,              # Number of PQ subquantizers (compression factor)
    nbits=8,          # Bits per code (default: 8, meaning 256 centroids per subquantizer)
    metric="l2",      # Distance metric: "l2" or "ip" (inner product)
    nprobe=10,        # Clusters to search at query time
    k=10,             # Default k for forward() method
)

index.train(vectors)  # Train IVF centroids and PQ codebooks
index.add(vectors)
distances, indices = index.search(queries, k=10)

Compression: With M=8 and nbits=8, each 128-dim vector (512 bytes) is compressed to just 8 bytes - a 64x reduction in memory usage.

from_faiss(index)

Convert a FAISS index to PyTorch.

from torch_similarity_search import from_faiss

torch_index = from_faiss(faiss_index)  # Returns FlatIndex, IVFFlatIndex, or IVFPQIndex

Supported:

• faiss.IndexFlatL2, faiss.IndexFlatIP → FlatIndex
• faiss.IndexIVFFlat → IVFFlatIndex
• faiss.IndexIVFPQ → IVFPQIndex

Requirements

• Python 3.11+
• PyTorch 2.0+
• NumPy (for FAISS conversion only)
• FAISS (optional, for conversion only)

License

MIT