docbert-plaid 1.0.0

PLAID-style multi-vector index for docbert (ColBERT late-interaction)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

//! Benchmarks for the residual quantization codec.
//!
//! `encode_vector` runs at every-token-of-every-document scale during
//! index build; `decode_vector` runs once per candidate token at search
//! time. `train_quantizer` runs once per index build but on every
//! residual, so its scaling matters too.

#[path = "shared.rs"]
mod shared;

use std::hint::black_box;

use criterion::{
    BenchmarkId,
    Criterion,
    Throughput,
    criterion_group,
    criterion_main,
};
use docbert_plaid::codec::{EncodedVector, ResidualCodec, train_quantizer};

const DIM: usize = 128;
const K: usize = 256;

fn build_codec(nbits: u32) -> ResidualCodec {
    let centroids = shared::random_unit_vectors(0xC047, K, DIM);
    // Train cutoffs/weights on a representative sample of residuals.
    let sample: Vec<f32> = shared::random_unit_vectors(0xDA7A, 4096, DIM)
        .iter()
        .map(|v| v * 0.1)
        .collect();
    let (bucket_cutoffs, bucket_weights) = train_quantizer(sample, nbits);
    ResidualCodec {
        nbits,
        dim: DIM,
        centroids,
        bucket_cutoffs,
        bucket_weights,
    }
}

fn bench_encode_vector(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec/encode_vector");
    let vectors = shared::random_unit_vectors(0xE1, 1024, DIM);
    for &nbits in &[2u32, 4] {
        let codec = build_codec(nbits);
        group.throughput(Throughput::Elements(1));
        group.bench_with_input(
            BenchmarkId::from_parameter(format!("nbits={nbits}")),
            &nbits,
            |b, _| {
                let mut idx = 0usize;
                b.iter(|| {
                    let start = idx * DIM;
                    let v = &vectors[start..start + DIM];
                    idx = (idx + 1) % 1024;
                    codec.encode_vector(black_box(v))
                });
            },
        );
    }
    group.finish();
}

fn bench_decode_vector(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec/decode_vector");
    let vectors = shared::random_unit_vectors(0xE2, 1024, DIM);
    for &nbits in &[2u32, 4] {
        let codec = build_codec(nbits);
        let encoded: Vec<EncodedVector> = vectors
            .as_chunks::<DIM>()
            .0
            .iter()
            .map(|v| codec.encode_vector(v).unwrap())
            .collect();
        group.throughput(Throughput::Elements(1));
        group.bench_with_input(
            BenchmarkId::from_parameter(format!("nbits={nbits}")),
            &nbits,
            |b, _| {
                let mut idx = 0usize;
                b.iter(|| {
                    let ev = &encoded[idx];
                    idx = (idx + 1) % encoded.len();
                    codec.decode_vector(black_box(ev))
                });
            },
        );
    }
    group.finish();
}

fn bench_batch_encode_tokens(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec/batch_encode_tokens");
    for &n in &[1_000usize, 10_000, 100_000] {
        let tokens = shared::random_unit_vectors(0xE4, n, DIM);
        for &nbits in &[2u32, 4] {
            let codec = build_codec(nbits);
            group.throughput(Throughput::Elements(n as u64));
            group.sample_size(10);
            group.bench_with_input(
                BenchmarkId::from_parameter(format!("n={n},nbits={nbits}")),
                &(n, nbits),
                |b, _| {
                    b.iter(|| codec.batch_encode_tokens(black_box(&tokens)));
                },
            );
        }
    }
    group.finish();
}

fn bench_train_quantizer(c: &mut Criterion) {
    let mut group = c.benchmark_group("codec/train_quantizer");
    for &n in &[1_024usize, 16_384, 131_072] {
        let residuals: Vec<f32> = shared::random_unit_vectors(0xE3, n, 1)
            .iter()
            .map(|v| v * 0.1)
            .collect();
        group.throughput(Throughput::Elements(n as u64));
        group.bench_with_input(BenchmarkId::from_parameter(n), &n, |b, _| {
            b.iter(|| train_quantizer(black_box(residuals.clone()), 4));
        });
    }
    group.finish();
}

criterion_group!(
    benches,
    bench_encode_vector,
    bench_decode_vector,
    bench_batch_encode_tokens,
    bench_train_quantizer,
);
criterion_main!(benches);