torch-similarity-search 0.0.2

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
IndexIVFFlat	IVFFlatIndex	✅ Supported
IndexIVFPQ	IVFPQIndex	Planned
IndexFlat	FlatIndex	Planned

API Reference

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" or "ip" (inner product)
    nprobe=10,        # Clusters to search at query time
    k=10,             # Default k for forward() method
)

Methods:

Method	Description
train(vectors)	Train cluster centroids via k-means. Requires (n, dim) tensor with n >= nlist.
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.

Properties:

Property	Description
ntotal	Number of indexed vectors
nprobe	Clusters to probe during search (settable, higher = more accurate)
k	Default k for forward() (settable)
is_trained	Whether index has been trained

from_faiss(index)

Convert a trained FAISS index to PyTorch.

from torch_similarity_search import from_faiss

torch_index = from_faiss(faiss_index)  # Returns IVFFlatIndex

Supported: faiss.IndexIVFFlat (L2 and inner product metrics)

Requirements

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

License

MIT