Metadata-Version: 2.4
Name: coala_embeddings
Version: 0.1.3
Summary: Coarse-to-fine Optimized Asymmetric Loss Architecture for fast and efficient embedding retrieval.
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: torch
Requires-Dist: faiss-cpu
Requires-Dist: beir

<img width="2048" height="512" alt="image (9)" src="https://github.com/user-attachments/assets/fe8d3ea8-6b7d-4b13-84c8-abed60aac8f5" />

# Coala Embeddings

#### A 20 million document index takes 80 GB of memory. Coala can fit it into under 1 GB—retaining over 94% search quality.

Coala Embeddings is a compression framework that shrinks embeddings by up to 99% by training a lightweight neural network along with quantization. It enables compression of any existing embedding model in seconds, without sacrificing retrieval accuracy.

## The main features are
- Fast search: Even on low end CPUs, the speed can reach over 1000 queries per second.
- Excellent Quality: Our novel neural network and losses are up to 27% better quality than previous SOTA.
- Lightweight code: The repository only has around 1k lines of code and easy to install using pip.
- Extreme RAM savings: can compress embeddings up to 256x, saving significantly more RAM than just pure quantization (TurboQuant, RabitQ, etc.)


## Usage

#### Simple 1 line installation
```python
pip install coala_embeddings
```

#### Simple 3 line usage

```python
from coala_embeddings.manager import ProposedIndex

# 1. Initialize with a sample of your dataset (4,000+ documents recommended)
index = ProposedIndex(train_embeddings=train_data, directory="document_embeddings")

# 2. Add all your document embeddings to the compressed index
index.add_documents(documents)

# 3. Search to get the top 10 most relevant document IDs
scores, doc_ids = index.search(query, fine_k=10)
```
Note: Train_data, documents, query should be formatted as a numpy array with shape like this (num_documents, dimension). Documents/Network are automatically saved to directory.

#### Loading documents
```python
from coala_embeddings.manager import ProposedIndex
index = ProposedIndex(load=True, directory="document_embeddings")

## If you don't want to load documents but load the neural network compresser only, then do the below
## encoder_path = 'document_embeddings/encoder.pkl' ## Make sure to change the path if it's not correct.
## index = ProposedIndex(encoder_path=encoder_path, load=True, directory="new_doc_embeddings", scratch_doc=True)
```

## Advanced Usage

#### Usage with description of inputs
```python
import numpy as np
from coala_embeddings.manager import ProposedIndex

# 1. Initialize with GPU support and custom dimensions
index = ProposedIndex(
    train_embeddings=train_data,  # Pre-training sample (~4k+ vectors)
    original_dim=1024,            # Input vector dimension
    slow_dim=256,                 # Re-ranking stage dimension(stored on disk)
    fast_dim=32,                  # Fast IVF candidate generation dimension(stored on memory)
    device='cpu',                 # Acceleration ('cuda' or 'cpu')
    probes=256,                   # Centroids for IVF index (increase for 1M+ docs)
    directory="document_embeddings",       # Path to save/load index binaries
    num_threads=8                 # CPU threads for FAISS
)

# 2. Add full document corpus
index.add_documents(documents)

# 3. Fine-tune two-stage search performance
scores, doc_ids = index.search(
    query,
    coarse_k=200,  # Number of candidates retrieved in Stage 1 (Fast)
    fine_k=10,     # Final re-ranked results returned in Stage 2 (Slow)
    nprobe=16      # IVF centroids to visit during search. Higher = better quality but slower.
)
```

#### Load and search
```python
from coala_embeddings.manager import ProposedIndex

# Reload existing index files directly from disk
index = ProposedIndex(
    load=True,
    directory="document_embeddings",
    device='cuda'
)

# Ready to search immediately
scores, doc_ids = index.search(query, fine_k=10, nprobe=16)
```

## Tips

- RAM Management for Indexing: Build your index on a high-RAM machine. Indexing operations (like IVF clustering and quantization) require significantly more RAM during setup than during search. Once built, you can easily load and run it on a low-RAM device.
- Hardware Acceleration: Utilize a GPU for the initial training phase. For loading and searching, CPU execution is fine and often faster than transferring data back and forth to a GPU per query.
- Training Corpus Size: Make sure you provide at least 4,000 document vectors in your training corpus (`train_embeddings`) for optimal compression quality.
- Hands-on Examples: Check out the `notebooks/` directory for step-by-step real world examples using popular models like mxbai-embed-large and nomic-embed-text-v1.5.

## Frequently Asked Questions

#### Q: How does this work?
Modern embedding models produce high-dimensional vectors (768 to 2048+ dimensions). We use a lightweight neural network trained with novel loss functions to project these embeddings into two paths:
* Low-dim path (32-64 dim): Stored in memory to perform an ultra-fast coarse search (for example, fetching top 200 candidates). The loss used for training optimizes for coarse recall.
* High-dim path (~256 dim): Stored on disk to re-rank those candidates into the final top 10 results. The loss used for training optimizes fine retrieval.

Both paths are further compressed using 4-bit quantization.

#### Q: How does it achieve 64x memory savings?
By keeping only the ultra-compressed low-dimensional index in RAM while leaving the higher-dimensional re-ranking index on disk. Because disk reads are only triggered for the top candidate results, you get massive memory savings without a performance issue.

#### Q: How does it compare to SVD, MRL, or other techniques?
It significantly outperforms them. Traditional post-hoc methods like SVD degrade quickly at extreme compression levels (e.g., 32 dimensions), and MRL requires expensive model re-training. Coala delivers superior retrieval quality post-hoc on your existing models in seconds. You can see this in the next section.

## Benchmarks
As seen in the image, our method (colored green) is consistently #1, infact up to 27% better then SVD. The method is very close to no-compression (colored black), reaching 95% of it's quality.
<img width="1600" height="513" alt="image" src="https://github.com/user-attachments/assets/95d58b8c-d051-4ea7-977f-8406f82c7f4c" />

[View full benchmarks](benchmarks.md)

## Roadmap
- [ ] Improve efficiency and memory.
- [ ] Huggingface space demo with real world model.
- [ ] Release paper for technique and more benchmarks.
      
## Acknowledgments
- [Faiss](https://github.com/facebookresearch/faiss): Very useful to create optimized IVF indexes and quantize code.
- [Beir](https://github.com/beir-cellar/beir): Simple and effective benchmarking code.
- [Matryoshoka Paper](https://arxiv.org/abs/2205.13147): Useful paper with several ideas used here.
  
## Final Notes
The code is licensed under the Apache-2.0 license. See LICENSE for details.

Stars would be appreciated, thank you.

Email: yatharthsharma3501@gmail.com
