hypersteeldb 0.5.5

//! Multi-Facet Orthogonal SPLADE projector — runs the baked ONNX graph
//! (encoder + K facet heads + SPLADE pool → `vecs [1,K,V]`) and reads the top active vocab terms per
//! facet, emitting facet-prefixed infon tokens (`actor/…`, `instrument/…`, `target/…`,
//! `constraint/…`). Ported from `src/ingest/splade/index.ts`.

use ort::session::Session;
use ort::value::Tensor;
use serde::Deserialize;
use std::error::Error;
use std::path::Path;
use tokenizers::Tokenizer;

type Res<T> = Result<T, Box<dyn Error + Send + Sync>>;

#[derive(Deserialize)]
struct FacetsCfg {
    facets: Vec<String>,
    vocab_size: usize,
}

#[derive(Debug, Clone)]
pub struct FacetTerm {
    pub token: String,
    pub term: String,
    pub weight: f32,
}

pub struct SpladeProjector {
    session: Session,
    tok: Tokenizer,
    facets: Vec<String>,
    vocab_size: usize,
    bpe: bool,
}

// trigger words too generic to be visual evidence / to snap a whole-word token onto
const TRIG_STOP: &[&str] = &[
    "the", "a", "an", "of", "to", "in", "on", "for", "and", "or", "is", "are", "was", "were", "be",
    "been", "this", "that", "these", "those", "with", "by", "from", "as", "at", "it", "its", "their",
    "our", "your", "we", "they", "them", "then", "than", "so", "such", "can", "may", "will", "would",
    "could", "should", "not", "no", "into", "over", "under", "more", "most", "many", "much", "few",
    "some", "any", "all", "each", "using", "use", "used",
];
// function/structure words that must never be emitted as a facet token (precision mask over the vocab)
const FACET_STOP: &[&str] = &[
    "and", "the", "of", "to", "in", "on", "for", "is", "are", "was", "were", "be", "been", "this",
    "that", "these", "those", "not", "all", "any", "based", "item", "source", "record", "note",
    "summary", "review", "section", "abstract", "its", "our", "their", "with", "by", "from", "as",
    "at", "we", "they", "it", "a", "an", "or", "new",
];

/// Snap a wordpiece position to the WHOLE word — merge leading non-## and trailing `##` continuations —
/// and return the word's char span + text. Mirrors the TS `wordSpanAt`.
fn word_span_at(pieces: &[String], offsets: &[(usize, usize)], pos: usize, text: &str) -> Option<(usize, usize, String)> {
    if pos >= pieces.len() {
        return None;
    }
    let mut s = pos;
    let mut e = pos;
    while s > 0 && pieces[s].starts_with("##") {
        s -= 1;
    }
    while e + 1 < pieces.len() && pieces[e + 1].starts_with("##") {
        e += 1;
    }
    let (a, b) = (offsets.get(s)?.0, offsets.get(e)?.1);
    if b <= a {
        return None;
    }
    Some((a, b, text.get(a..b)?.to_string()))
}

fn slug_ascii(s: &str) -> String {
    let mut out = String::new();
    let mut dash = false;
    for ch in s.chars() {
        if ch.is_ascii_alphanumeric() {
            out.push(ch.to_ascii_lowercase());
            dash = false;
        } else if !out.is_empty() && !dash {
            out.push('-');
            dash = true;
        }
    }
    while out.ends_with('-') {
        out.pop();
    }
    out
}

fn slug_unicode(s: &str) -> String {
    let s = s.strip_prefix('\u{2581}').unwrap_or(s); // strip ▁ metaspace marker
    let mut out = String::new();
    let mut dash = false;
    for ch in s.chars() {
        if ch.is_alphanumeric() {
            out.extend(ch.to_lowercase());
            dash = false;
        } else if !out.is_empty() && !dash {
            out.push('-');
            dash = true;
        }
    }
    while out.ends_with('-') {
        out.pop();
    }
    out
}

impl SpladeProjector {
    /// `dir` holds `splade.onnx`, `facets.json`, and `tokenizer/tokenizer.json`. `multilingual` selects
    /// the BPE term-filtering branch (word-start ▁ tokens); otherwise the English WordPiece branch.
    pub fn load(dir: &Path, multilingual: bool) -> Res<SpladeProjector> {
        let cfg: FacetsCfg = serde_json::from_slice(&std::fs::read(dir.join("facets.json"))?)?;
        let session = Session::builder()?.commit_from_file(dir.join("splade.onnx"))?;
        let mut tok = Tokenizer::from_file(dir.join("tokenizer/tokenizer.json"))?;
        tok.with_truncation(Some(tokenizers::TruncationParams {
            max_length: 64,
            ..Default::default()
        }))?;
        Ok(SpladeProjector { session, tok, facets: cfg.facets, vocab_size: cfg.vocab_size, bpe: multilingual })
    }

    pub fn facets(&self) -> &[String] {
        &self.facets
    }

    /// Project text → up to `top_k` faceted terms per facet, pruning activations ≤ `min_weight`.
    pub fn project(&mut self, text: &str, top_k: usize, min_weight: f32) -> Res<Vec<FacetTerm>> {
        let enc = self.tok.encode(text, true)?;
        let ids: Vec<i64> = enc.get_ids().iter().map(|&x| x as i64).collect();
        let t = ids.len();
        let attn: Vec<i64> = vec![1; t];
        let id_t = Tensor::from_array(([1usize, t], ids))?;
        let at_t = Tensor::from_array(([1usize, t], attn))?;
        let outputs = self
            .session
            .run(ort::inputs!["input_ids" => id_t, "attention_mask" => at_t])?;
        let (_shape, vecs) = outputs["vecs"].try_extract_tensor::<f32>()?; // [1, K, V]
        // `src` [1, K, V] = the input token position whose logit fired each vocab dim (argmax trace),
        // used to snap a coarse wordpiece back to its whole trigger word. Optional: older exports omit it.
        let src: Option<&[i64]> = outputs.get("src").and_then(|v| v.try_extract_tensor::<i64>().ok()).map(|(_, s)| s);
        // input wordpieces + their char/byte spans, for the whole-word trace (WordPiece branch)
        let in_pieces: Vec<String> = enc.get_tokens().to_vec();
        let offsets: Vec<(usize, usize)> = enc.get_offsets().to_vec();
        let v = self.vocab_size;

        let mut out = Vec::new();
        for (f, facet) in self.facets.iter().enumerate() {
            let base = f * v;
            let mut hits: Vec<(usize, f32)> = Vec::new();
            for id in 0..v {
                let w = vecs[base + id];
                if w > min_weight {
                    hits.push((id, w));
                }
            }
            hits.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

            let mut seen = std::collections::HashSet::new();
            let mut n = 0;
            for (id, w) in hits {
                let wp = match self.tok.id_to_token(id as u32) {
                    Some(s) => s,
                    None => continue,
                };
                let s = if self.bpe {
                    if !wp.starts_with('\u{2581}') || wp.starts_with('<') || wp.starts_with('[') {
                        continue;
                    }
                    slug_unicode(&wp)
                } else {
                    if wp.starts_with("##")
                        || wp.starts_with('[')
                        || wp.chars().filter(|c| c.is_ascii_alphabetic()).count() < 3
                        || FACET_STOP.contains(&wp.as_str())
                    {
                        continue;
                    }
                    // HIGH-RESOLUTION token: when the head fired on a GROUNDED wordpiece (a fragment of a
                    // clean whole word in the text — "reds" of "redshift"), emit the WHOLE WORD, not the
                    // fragment; keep the raw vocab term only for pure expansions with no clean trigger word.
                    let mut chosen = slug_ascii(&wp);
                    if let Some(src) = src {
                        let pos = *src.get(base + id).unwrap_or(&-1);
                        if pos >= 0 {
                            if let Some((_, _, word)) = word_span_at(&in_pieces, &offsets, pos as usize, text) {
                                let wl = word.to_lowercase();
                                let clean = word.chars().count() >= 3
                                    && word.chars().all(|c| c.is_ascii_alphabetic())
                                    && !TRIG_STOP.contains(&wl.as_str());
                                let prefix: String = wp.chars().take(3).flat_map(|c| c.to_lowercase()).collect();
                                if clean && wl.contains(&prefix) {
                                    chosen = slug_ascii(&word);
                                }
                            }
                        }
                    }
                    chosen
                };
                if s.len() < 2 || !seen.insert(s.clone()) {
                    continue;
                }
                out.push(FacetTerm {
                    token: format!("{facet}/{s}"),
                    term: wp.trim_start_matches('\u{2581}').to_string(),
                    weight: (w * 1000.0).round() / 1000.0,
                });
                n += 1;
                if n >= top_k {
                    break;
                }
            }
        }
        Ok(out)
    }
}