feat: late interaction and MUVERA

[lmmx]

Following the example of the fastembed Python lib (qdrant/fastembed#542) this PR adds support (via a new muvera feature) for both late interaction models (ColBERT and AnswerAIColbertSmallV1) and MUVERA

• Qdrant blog: https://qdrant.tech/articles/muvera-embeddings/

Working standalone demo:

• Included temporarily during development at examples/muvera_demo.rs, will be removed when finalised

Loading ColBERT model...
Document 0 shape: (11, 128)
Query 0 shape: (32, 128)
FDE size: 10240
Doc FDE shape: (10240,)
Query-Doc0 similarity (MUVERA): 513.4094
Query-Doc1 similarity (MUVERA): 86.4530
Query-Doc0 MaxSim score: 29.5734
Query-Doc1 MaxSim score: 9.9226

TODO:

• Document new models in README
  • colbert-ir/colbertv2.0
  • answerdotai/answerai-colbert-small-v1
• Add new models' weight tests
• Compiles with cargo check, address cargo clippy lints
• Match results from ColBERT
  • Put into the test test_muvera_with_colbert_model
• Run benchmarks on beir/scidocs for ColBERT (aiming to match these)
• Evaluate relative speed vs the Python implementation

Out of scope (?)

• Multimodal MUVERA
• Other models: Jina Colbert v2, ...

status:

• ✔️ Example Working at 635c01a
• 🚧 WIP (BEIR scidocs showing incompatible shapes)

[lmmx]

I have a libonnxruntime.so lying around from polars-fastembed-cuda (it could also be installed via venv from PyPI)

ORT_DYLIB_PATH=/home/louis/dev/polars-fastembed/polars-fastembed-cuda/python/polars_fastembed_cuda/libs/libonnxruntime.so cargo run --example beir_scidocs_benchmark --features "hf-hub,muvera,ort-load-dynamic" --release

BEIR scidocs bench:

Batch shape inconsistency error

Error: ShapeError/IncompatibleShape: incompatible shapes

workaround

diff --git a/examples/beir_scidocs_benchmark.rs b/examples/beir_scidocs_benchmark.rs
index 9c861c7..15e3024 100644
--- a/examples/beir_scidocs_benchmark.rs
+++ b/examples/beir_scidocs_benchmark.rs
@@ -144,18 +144,25 @@ fn main() -> Result<()> {
 
     // 3. Embed corpus
     println!("Embedding corpus...");
-    let corpus_embeddings = model.embed(&corpus_texts, Some(batch_size))?;
-    
-    // Create id -> index mapping
-    let corpus_id_to_idx: HashMap<&str, usize> = corpus_ids
-        .iter()
-        .enumerate()
-        .map(|(i, id)| (id.as_str(), i))
-        .collect();
+    let mut corpus_embeddings = Vec::with_capacity(corpus_texts.len());
+    for (i, text) in corpus_texts.iter().enumerate() {
+        if i % 1000 == 0 {
+            println!("  Embedded {}/{} documents", i, corpus_texts.len());
+        }
+        let emb = model.embed(&[text], None)?;
+        corpus_embeddings.push(emb.into_iter().next().unwrap());
+    }
 
     // 4. Embed queries
     println!("Embedding queries...");
-    let query_embeddings = model.query_embed(&query_texts, Some(batch_size))?;
+    let mut query_embeddings = Vec::with_capacity(query_texts.len());
+    for (i, text) in query_texts.iter().enumerate() {
+        if i % 100 == 0 {
+            println!("  Embedded {}/{} queries", i, query_texts.len());
+        }
+        let emb = model.query_embed(&[text], None)?;
+        query_embeddings.push(emb.into_iter().next().unwrap());
+    }
 
     // 5. Evaluate brute-force ColBERT (baseline)
     println!("\n--- Evaluating brute-force ColBERT ---");

Something's very off

Goal is to reproduce:

Vector: colbert, Recall@4: 0.4470, Recall@5: 0.4920, Recall@10: 0.7030
Vector: muvera-1280, Recall@4: 0.2350, Recall@5: 0.2610, Recall@10: 0.3700
Vector: muvera-2560, Recall@4: 0.2680, Recall@5: 0.3070, Recall@10: 0.4290
Vector: muvera-5120, Recall@4: 0.3020, Recall@5: 0.3370, Recall@10: 0.4770
Vector: muvera-10240, Recall@4: 0.3350, Recall@5: 0.3810, Recall@10: 0.5580
Vector: muvera-15360, Recall@4: 0.35, Recall@5: 0.4, Recall@10: 0.558
Vector: muvera-20480, Recall@4: 0.357, Recall@5: 0.415, Recall@10: 0.594

Parameters I was using:

r_reps, k_sim, dim_proj:
(20, 3, 8),
(20, 4, 8),
(20, 5, 8),
(20, 5, 16),
(30, 5, 16),
(40, 5, 16)

Getting:

--- Evaluating MUVERA-1280 (r=20, k=3, d=8) ---
Vector: muvera-1280, Recall@4: 0.0303, Recall@5: 0.0353, Recall@10: 0.0613
--- Evaluating MUVERA-2560 (r=20, k=4, d=8) ---
Vector: muvera-2560, Recall@4: 0.0463, Recall@5: 0.0548, Recall@10: 0.0827
--- Evaluating MUVERA-5120 (r=20, k=5, d=8) ---
Vector: muvera-5120, Recall@4: 0.0542, Recall@5: 0.0629, Recall@10: 0.0940
--- Evaluating MUVERA-10240 (r=20, k=5, d=16) ---
Vector: muvera-10240, Recall@4: 0.0687, Recall@5: 0.0814, Recall@10: 0.1173

[lmmx]

Note: it's been a month and I've got other stuff taking priority over this work, apologies that it got held up but for anyone else seeking to take it on, the BEIR benchmark would be the place to seek to check that this is working (if not clear from the discussion above), or feel free to make an entirely fresh attempt.

[Anush008]

Closing since stale.
Please feel free to re-open.