holographic-memory 0.6.3

Local-first search engine and agent memory: fast approximate nearest-neighbor and BM25 full-text search, encrypted stores, provable deletion, built on vector-symbolic architectures.
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

// Copyright 2024-2026 WritersLogic Contributors
// SPDX-License-Identifier: Apache-2.0

use anyhow::Result;
use candle_core::{DType, Device, Tensor, D};

use super::qwen3::Qwen3;
use super::{encode, encoder_batches, last_hidden, load, prompts, token_id, ModelSource};

/// Qwen3-Reranker: p(yes) from a softmax over the `no`/`yes` logits after the model-card
/// prompt; the query/document body is cut to `4096 - prefix - suffix` tokens.
pub struct Reranker {
    model: Qwen3,
    tokenizer: tokenizers::Tokenizer,
    device: Device,
    prefix: Vec<u32>,
    suffix: Vec<u32>,
    yes: u32,
    no: u32,
    batch: usize,
}

const MAX_TOKENS: usize = 4096;
const BATCH_TOKENS: usize = 12288;
/// Default padded batch on Metal/CUDA: on an M4 (Metal), 96 pairs scored in 10.6 s at batch 1
/// against 10.9 s at 2, 11.1 s at 4, 12.0 s at 8 and 15.3 s at 16, so padding does not pay
/// (`benchmarks/results/local_models_batching.json`). CUDA is unmeasured.
const GPU_BATCH: usize = 1;

impl Reranker {
    pub fn load(source: &ModelSource, device: &Device) -> Result<Self> {
        let l = load(source, device, super::dtype_for(device, false))?;
        let mut tokenizer = l.tokenizer;
        tokenizer.with_padding(None);
        // Only invalid truncation parameters are rejected; disabling it cannot fail.
        let _ = tokenizer.with_truncation(None);
        let prefix = encode(&tokenizer, prompts::RERANK_PREFIX, false)?;
        let suffix = encode(&tokenizer, prompts::RERANK_SUFFIX, false)?;
        Ok(Self {
            yes: token_id(&tokenizer, "yes")?,
            no: token_id(&tokenizer, "no")?,
            model: Qwen3::new(&l.config, l.vb, true)?,
            tokenizer,
            device: device.clone(),
            prefix,
            suffix,
            batch: super::default_encode_batch(device, GPU_BATCH),
        })
    }

    /// Sets the most sequences per forward pass (at least 1).
    pub fn set_batch(&mut self, batch: usize) {
        self.batch = batch.max(1);
    }

    /// Relevance in [0, 1] for each (query, document) pair, in input order.
    pub fn score<Q: AsRef<str>, T: AsRef<str>>(&self, pairs: &[(Q, T)]) -> Result<Vec<f32>> {
        let cap = MAX_TOKENS - self.prefix.len() - self.suffix.len();
        let ids: Vec<Vec<u32>> = pairs
            .iter()
            .map(|(q, d)| {
                let mut body = encode(
                    &self.tokenizer,
                    &prompts::rerank_body(q.as_ref(), d.as_ref()),
                    true,
                )?;
                body.truncate(cap);
                let mut x = self.prefix.clone();
                x.extend(body);
                x.extend(&self.suffix);
                Ok(x)
            })
            .collect::<Result<_>>()?;
        let lens: Vec<usize> = ids.iter().map(Vec::len).collect();
        let mut out = vec![0.0; pairs.len()];
        for batch in encoder_batches(&self.device, &lens, BATCH_TOKENS, self.batch) {
            let logits = self
                .model
                .logits(&last_hidden(&self.model, &ids, &batch)?)?;
            let no = logits.narrow(D::Minus1, self.no as usize, 1)?;
            let yes = logits.narrow(D::Minus1, self.yes as usize, 1)?;
            let pair = Tensor::cat(&[no, yes], 1)?.to_dtype(DType::F32)?;
            let p = candle_nn::ops::log_softmax(&pair, 1)?
                .narrow(1, 1, 1)?
                .exp()?;
            for (&i, v) in batch.iter().zip(p.flatten_all()?.to_vec1::<f32>()?) {
                out[i] = v;
            }
        }
        Ok(out)
    }
}