steeldb/text/
splade.rs

1//! Multi-Facet Orthogonal SPLADE projector — runs the baked ONNX graph
2//! (encoder + K facet heads + SPLADE pool → `vecs [1,K,V]`) and reads the top active vocab terms per
3//! facet, emitting facet-prefixed infon tokens (`actor/…`, `instrument/…`, `target/…`,
4//! `constraint/…`). Ported from `src/ingest/splade/index.ts`.
5
6use ort::session::Session;
7use ort::value::Tensor;
8use serde::Deserialize;
9use std::error::Error;
10use std::path::Path;
11use tokenizers::Tokenizer;
12
13type Res<T> = Result<T, Box<dyn Error + Send + Sync>>;
14
15#[derive(Deserialize)]
16struct FacetsCfg {
17    facets: Vec<String>,
18    vocab_size: usize,
19}
20
21#[derive(Debug, Clone)]
22pub struct FacetTerm {
23    pub token: String,
24    pub term: String,
25    pub weight: f32,
26}
27
28pub struct SpladeProjector {
29    session: Session,
30    tok: Tokenizer,
31    facets: Vec<String>,
32    vocab_size: usize,
33    bpe: bool,
34}
35
36// trigger words too generic to be visual evidence / to snap a whole-word token onto
37const TRIG_STOP: &[&str] = &[
38    "the", "a", "an", "of", "to", "in", "on", "for", "and", "or", "is", "are", "was", "were", "be",
39    "been", "this", "that", "these", "those", "with", "by", "from", "as", "at", "it", "its", "their",
40    "our", "your", "we", "they", "them", "then", "than", "so", "such", "can", "may", "will", "would",
41    "could", "should", "not", "no", "into", "over", "under", "more", "most", "many", "much", "few",
42    "some", "any", "all", "each", "using", "use", "used",
43];
44// function/structure words that must never be emitted as a facet token (precision mask over the vocab)
45const FACET_STOP: &[&str] = &[
46    "and", "the", "of", "to", "in", "on", "for", "is", "are", "was", "were", "be", "been", "this",
47    "that", "these", "those", "not", "all", "any", "based", "item", "source", "record", "note",
48    "summary", "review", "section", "abstract", "its", "our", "their", "with", "by", "from", "as",
49    "at", "we", "they", "it", "a", "an", "or", "new",
50];
51
52/// Snap a wordpiece position to the WHOLE word — merge leading non-## and trailing `##` continuations —
53/// and return the word's char span + text. Mirrors the TS `wordSpanAt`.
54fn word_span_at(pieces: &[String], offsets: &[(usize, usize)], pos: usize, text: &str) -> Option<(usize, usize, String)> {
55    if pos >= pieces.len() {
56        return None;
57    }
58    let mut s = pos;
59    let mut e = pos;
60    while s > 0 && pieces[s].starts_with("##") {
61        s -= 1;
62    }
63    while e + 1 < pieces.len() && pieces[e + 1].starts_with("##") {
64        e += 1;
65    }
66    let (a, b) = (offsets.get(s)?.0, offsets.get(e)?.1);
67    if b <= a {
68        return None;
69    }
70    Some((a, b, text.get(a..b)?.to_string()))
71}
72
73fn slug_ascii(s: &str) -> String {
74    let mut out = String::new();
75    let mut dash = false;
76    for ch in s.chars() {
77        if ch.is_ascii_alphanumeric() {
78            out.push(ch.to_ascii_lowercase());
79            dash = false;
80        } else if !out.is_empty() && !dash {
81            out.push('-');
82            dash = true;
83        }
84    }
85    while out.ends_with('-') {
86        out.pop();
87    }
88    out
89}
90
91fn slug_unicode(s: &str) -> String {
92    let s = s.strip_prefix('\u{2581}').unwrap_or(s); // strip ▁ metaspace marker
93    let mut out = String::new();
94    let mut dash = false;
95    for ch in s.chars() {
96        if ch.is_alphanumeric() {
97            out.extend(ch.to_lowercase());
98            dash = false;
99        } else if !out.is_empty() && !dash {
100            out.push('-');
101            dash = true;
102        }
103    }
104    while out.ends_with('-') {
105        out.pop();
106    }
107    out
108}
109
110impl SpladeProjector {
111    /// `dir` holds `splade.onnx`, `facets.json`, and `tokenizer/tokenizer.json`. `multilingual` selects
112    /// the BPE term-filtering branch (word-start ▁ tokens); otherwise the English WordPiece branch.
113    pub fn load(dir: &Path, multilingual: bool) -> Res<SpladeProjector> {
114        let cfg: FacetsCfg = serde_json::from_slice(&std::fs::read(dir.join("facets.json"))?)?;
115        let session = Session::builder()?.commit_from_file(dir.join("splade.onnx"))?;
116        let mut tok = Tokenizer::from_file(dir.join("tokenizer/tokenizer.json"))?;
117        tok.with_truncation(Some(tokenizers::TruncationParams {
118            max_length: 64,
119            ..Default::default()
120        }))?;
121        Ok(SpladeProjector { session, tok, facets: cfg.facets, vocab_size: cfg.vocab_size, bpe: multilingual })
122    }
123
124    pub fn facets(&self) -> &[String] {
125        &self.facets
126    }
127
128    /// Project text → up to `top_k` faceted terms per facet, pruning activations ≤ `min_weight`.
129    pub fn project(&mut self, text: &str, top_k: usize, min_weight: f32) -> Res<Vec<FacetTerm>> {
130        let enc = self.tok.encode(text, true)?;
131        let ids: Vec<i64> = enc.get_ids().iter().map(|&x| x as i64).collect();
132        let t = ids.len();
133        let attn: Vec<i64> = vec![1; t];
134        let id_t = Tensor::from_array(([1usize, t], ids))?;
135        let at_t = Tensor::from_array(([1usize, t], attn))?;
136        let outputs = self
137            .session
138            .run(ort::inputs!["input_ids" => id_t, "attention_mask" => at_t])?;
139        let (_shape, vecs) = outputs["vecs"].try_extract_tensor::<f32>()?; // [1, K, V]
140        // `src` [1, K, V] = the input token position whose logit fired each vocab dim (argmax trace),
141        // used to snap a coarse wordpiece back to its whole trigger word. Optional: older exports omit it.
142        let src: Option<&[i64]> = outputs.get("src").and_then(|v| v.try_extract_tensor::<i64>().ok()).map(|(_, s)| s);
143        // input wordpieces + their char/byte spans, for the whole-word trace (WordPiece branch)
144        let in_pieces: Vec<String> = enc.get_tokens().to_vec();
145        let offsets: Vec<(usize, usize)> = enc.get_offsets().to_vec();
146        let v = self.vocab_size;
147
148        let mut out = Vec::new();
149        for (f, facet) in self.facets.iter().enumerate() {
150            let base = f * v;
151            let mut hits: Vec<(usize, f32)> = Vec::new();
152            for id in 0..v {
153                let w = vecs[base + id];
154                if w > min_weight {
155                    hits.push((id, w));
156                }
157            }
158            hits.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
159
160            let mut seen = std::collections::HashSet::new();
161            let mut n = 0;
162            for (id, w) in hits {
163                let wp = match self.tok.id_to_token(id as u32) {
164                    Some(s) => s,
165                    None => continue,
166                };
167                let s = if self.bpe {
168                    if !wp.starts_with('\u{2581}') || wp.starts_with('<') || wp.starts_with('[') {
169                        continue;
170                    }
171                    slug_unicode(&wp)
172                } else {
173                    if wp.starts_with("##")
174                        || wp.starts_with('[')
175                        || wp.chars().filter(|c| c.is_ascii_alphabetic()).count() < 3
176                        || FACET_STOP.contains(&wp.as_str())
177                    {
178                        continue;
179                    }
180                    // HIGH-RESOLUTION token: when the head fired on a GROUNDED wordpiece (a fragment of a
181                    // clean whole word in the text — "reds" of "redshift"), emit the WHOLE WORD, not the
182                    // fragment; keep the raw vocab term only for pure expansions with no clean trigger word.
183                    let mut chosen = slug_ascii(&wp);
184                    if let Some(src) = src {
185                        let pos = *src.get(base + id).unwrap_or(&-1);
186                        if pos >= 0 {
187                            if let Some((_, _, word)) = word_span_at(&in_pieces, &offsets, pos as usize, text) {
188                                let wl = word.to_lowercase();
189                                let clean = word.chars().count() >= 3
190                                    && word.chars().all(|c| c.is_ascii_alphabetic())
191                                    && !TRIG_STOP.contains(&wl.as_str());
192                                let prefix: String = wp.chars().take(3).flat_map(|c| c.to_lowercase()).collect();
193                                if clean && wl.contains(&prefix) {
194                                    chosen = slug_ascii(&word);
195                                }
196                            }
197                        }
198                    }
199                    chosen
200                };
201                if s.len() < 2 || !seen.insert(s.clone()) {
202                    continue;
203                }
204                out.push(FacetTerm {
205                    token: format!("{facet}/{s}"),
206                    term: wp.trim_start_matches('\u{2581}').to_string(),
207                    weight: (w * 1000.0).round() / 1000.0,
208                });
209                n += 1;
210                if n >= top_k {
211                    break;
212                }
213            }
214        }
215        Ok(out)
216    }
217}
steeldb/text/splade.rs

steeldb/text/
splade.rs
