Metadata-Version: 2.4
Name: tq-search
Version: 0.1.0
Summary: High-performance PyTorch Vector Quantization Engine with Lloyd-Max and QJL sign residual correction.
Requires-Python: >=3.12
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.4.0

# tq-search

High-performance PyTorch Vector Quantization Engine with Lloyd-Max codebooks and QJL sign-bit residual correction.

---

## Overview

`tq-search` is a state-of-the-art vector compression and similarity search library built entirely in PyTorch. It implements the optimal quantization algorithms inspired by **TurboQuant** (Zandieh et al., ICLR 2026 / arXiv:2504.19874) and **QJL** (arXiv:2406.03482).

Unlike standard vector databases that require decompressing vector databases in RAM or VRAM to perform searches, `tq-search` performs **asymmetric inner product similarity scoring directly on compressed 5-bit representations**, reducing GPU memory usage by up to **4x–6x** while keeping search results highly accurate.

### Key Features
- **GPU-Accelerated**: Fully integrated with PyTorch, running seamlessly on CUDA and Apple Silicon MPS.
- **Batched Compression & Search**: Built for high-throughput similarity search, handling millions of candidate vectors concurrently.
- **Asymmetric Scoring**: Computes dot products directly on the compressed codebooks without decompressing residuals in memory.
- **Optimal Distortion**: Employs mathematically optimal Lloyd-Max codebooks fit to rotated coordinate distributions.

---

## Installation

Install the library directly from PyPI (once published):
```bash
pip install tq-search
```

Or install it locally in editable mode for development:
```bash
git clone https://github.com/barateza/tq-search.git
cd tq-search
pip install -e .
```

---

## Quickstart

Here is how to compress a database of embeddings and search them using `tq-search`:

```python
import torch
from tq_search import TurboQuantProd

# 1. Initialize the Quantizer
# We compress 1024-dimensional embeddings into 3-bit MSE + 1-bit QJL (4 bits total)
dim = 1024
bits = 4
quantizer = TurboQuantProd(d=dim, bits=bits, device="cuda")

# 2. Compress your database vectors (Shape: [N, dim])
database_vectors = torch.randn(10000, dim, device="cuda")
# L2-normalize vectors (recommended for maximum cosine similarity accuracy)
database_vectors = database_vectors / torch.norm(database_vectors, dim=-1, keepdim=True)

compressed = quantizer.quantize(database_vectors)
# 'compressed' is a lightweight dict containing:
# - mse_indices: [N, dim] (Stage 1 indices)
# - qjl_signs: [N, dim] (Stage 2 signs)
# - residual_norm: [N] (Dynamic residual scaling factor)

# 3. Perform Asymmetric Query Search (Query remains at full precision)
query = torch.randn(1, dim, device="cuda")
query = query / torch.norm(query, dim=-1, keepdim=True)

# Directly calculate dot products across all 10,000 compressed vectors in VRAM!
scores = quantizer.inner_product(query, compressed)
print("Inner Product Scores:", scores)
```

---

## Research Credits

This library is a high-performance search adaptation based on the following research papers:

1. **TurboQuant**: *"TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate"* (Zandieh et al., ICLR 2026) [arXiv:2504.19874](https://arxiv.org/pdf/2504.19874)
2. **QJL Residuals**: *"QJL: 1-Bit Quantized JL Transform for KV Cache Quantization"* (arXiv:2406.03482) [arXiv:2406.03482](https://arxiv.org/abs/2406.03482)

---

## License

MIT License. Feel free to use, modify, and distribute.
