Metadata-Version: 2.4
Name: torch-similarity-search
Version: 0.0.2
Summary: PyTorch-native similarity search: convert FAISS indexes to nn.Module
Requires-Python: >=3.10
Requires-Dist: faiss-gpu-cu12>=1.13.0
Requires-Dist: numpy>=1.20
Requires-Dist: torch>=2.0
Description-Content-Type: text/markdown

# torch-similarity-search

[![PyPI version](https://badge.fury.io/py/torch-similarity-search.svg)](https://pypi.org/project/torch-similarity-search/)
[![Python 3.11+](https://img.shields.io/badge/python-3.11+-blue.svg)](https://www.python.org/downloads/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)

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

```bash
pip install torch-similarity-search
```

For FAISS conversion support:

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

## Quick Start

### Convert from FAISS

```python
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

```python
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

```python
# 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

```python
# 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.

```python
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.

```python
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
