steeldb/projectors/

text_proj.rs

//! Native text projection — raw text → queryable Situation stream, entirely in Rust.
//! `TextEngine` loads the ONNX models once (expensive) and can project many files/sentences;
//! `TextProjector` wraps an engine + a single file for the `Projector` trait. Folder ingest loads
//! one `TextEngine` and reuses it across every text file.

use crate::projector::{slug, CorpusKind, Projector, Situation};
use crate::text::{Gazetteer, SpladeProjector, SpoTagger};
use std::path::{Path, PathBuf};

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

/// Non-discriminative facet leaves — they appear in most documents, so they add noise not signal to
/// the bitmap programs (they co-occur with everything and swamp rank/crosstab/cooccurs). The gazetteer
/// + whole-word tiers carry the real signal. Dropped at ingest. Ported from `pipeline.ts` SPLADE_NOISE.
const SPLADE_NOISE: &[&str] = &[
    "amazon", "aws", "cloud", "architecture", "pattern", "patterns", "management", "data", "time",
    "work", "works", "operations", "service", "services", "enterprise", "application", "applications",
    "production", "configuration", "solution", "solutions", "customer", "customers", "use", "uses",
    "using", "system", "systems", "platform", "platforms", "capability", "capabilities", "feature",
    "features", "support", "level", "provide", "provides", "available", "new", "business", "team",
    "teams", "organization", "organizations", "user", "users",
];

fn leaf_of(t: &str) -> &str {
    match t.find('/') {
        Some(i) => &t[i + 1..],
        None => t,
    }
}

/// Loaded text-projection models (SPO tagger + optional English SPLADE + optional gazetteer),
/// reusable across files.
pub struct TextEngine {
    tagger: SpoTagger,
    splade: Option<SpladeProjector>,
    gaz: Option<Gazetteer>,
    /// When set, multi-word entity spans discovered at ingest are learned into the gazetteer overlay and
    /// persisted here — the growing gazetteer. Later runs (ingest + query) recognise these entities.
    overlay_path: Option<PathBuf>,
    /// The **tuned** multi-head tagger (step 2). When present it supersedes the ONNX SPO tagger for
    /// dimensions 1/3/4/5: it emits *typed* hierarchical entity URIs (`org/…`) and the epistemic
    /// `state/…` cue that carries infon polarity, which the untyped ONNX tagger cannot. SPLADE (dim 6)
    /// and the gazetteer still contribute either way.
    #[cfg(feature = "native")]
    tuned: Option<crate::tagger_train::TunedTagger>,
}

impl TextEngine {
    pub fn load(ml_bundle: &Path, splade_dir: Option<&Path>) -> Res<TextEngine> {
        let tagger = SpoTagger::load(ml_bundle)?;
        let splade = match splade_dir {
            Some(d) => Some(SpladeProjector::load(d, false)?),
            None => None,
        };
        // High-res gazetteer: `STEELDB_GAZETTEER`, else `<splade_dir>/gazetteer.json` if present. Loaded
        // once and matched per sentence so index and query share the same whole-entity token space.
        let gaz = gazetteer_path(splade_dir).and_then(|p| Gazetteer::load(&p).ok());
        Ok(TextEngine {
            tagger,
            splade,
            gaz,
            overlay_path: None,
            #[cfg(feature = "native")]
            tuned: None,
        })
    }

    /// Enable the growing gazetteer: merge any prior overlay at `path`, then learn newly-discovered
    /// multi-word entities during ingest and persist them back with [`Self::save_overlay`].
    pub fn enable_growth(&mut self, path: impl Into<PathBuf>) {
        let path = path.into();
        let mut gaz = self.gaz.take().unwrap_or_else(Gazetteer::empty);
        gaz.merge_overlay(&path);
        self.gaz = Some(gaz);
        self.overlay_path = Some(path);
    }

    /// Persist the learned overlay (no-op unless growth was enabled). Returns the learned-entry count.
    pub fn save_overlay(&self) -> usize {
        if let (Some(gaz), Some(path)) = (self.gaz.as_ref(), self.overlay_path.as_ref()) {
            let _ = gaz.save_overlay(path);
            gaz.learned_len()
        } else {
            0
        }
    }

    /// Attach a tuned tagger (a `save()` directory from step 2) so ingest emits typed facets + polarity.
    /// `base_dir` is the encoder's HF snapshot (config.json) and `tokenizer` its `tokenizer.json`.
    #[cfg(feature = "native")]
    pub fn enable_tuned(&mut self, dir: &Path, base_dir: &Path, tokenizer: &Path, max_len: usize) -> Res<()> {
        self.tuned = Some(crate::tagger_train::TunedTagger::load(dir, base_dir, tokenizer, max_len)?);
        Ok(())
    }

    /// Attach Head C (step 5) to the tuned tagger so ingest emits bound relation tokens.
    #[cfg(feature = "native")]
    pub fn enable_relations(&mut self, dir: &Path, spec: &crate::vocabulary::VocabularySpace) -> Res<()> {
        match self.tuned.as_mut() {
            Some(t) => {
                t.enable_relations(dir, spec)?;
                Ok(())
            }
            None => Err("no tuned tagger — call enable_tuned first".into()),
        }
    }

    /// True when a tuned tagger is driving the entity/epistemic tiers.
    pub fn is_tuned(&self) -> bool {
        #[cfg(feature = "native")]
        {
            return self.tuned.is_some();
        }
        #[allow(unreachable_code)]
        false
    }

    /// Project one sentence into a full [`Situation`] — tokens, numeric fields, **and infon polarity**.
    /// This is the ingest entry point that preserves everything the tuned tagger knows; `project_sentence`
    /// remains for callers that only need tokens.
    pub fn project_situation(&mut self, sentence: &str) -> Situation {
        #[cfg(feature = "native")]
        if let Some(tt) = self.tuned.as_ref() {
            if let Ok(mut sit) = tt.project(sentence) {
                // dim 6 (latent motifs) and the gazetteer tier are orthogonal to the tuned tagger
                if let Some(splade) = self.splade.as_mut() {
                    if let Ok(terms) = splade.project(sentence, 8, 0.3) {
                        for t in terms {
                            if !SPLADE_NOISE.contains(&leaf_of(&t.token)) {
                                sit.tokens.push(t.token);
                            }
                        }
                    }
                }
                if let Some(gaz) = self.gaz.as_ref() {
                    for h in gaz.extract(sentence) {
                        sit.tokens.push(h.token);
                    }
                }
                sit.tokens.sort();
                sit.tokens.dedup();
                // polarity applies to every token of a negated/hedged clause, including the tiers above
                if let Some((_, level)) = sit.beliefs.first().map(|(k, v)| (k.clone(), *v)) {
                    sit.beliefs = sit.tokens.iter().map(|t| (t.clone(), level)).collect();
                }
                return sit;
            }
        }
        let (tokens, numbers) = self.project_sentence(sentence);
        Situation { tokens, display: vec![sentence.to_string()], numbers, beliefs: Vec::new() }
    }

    /// Project one sentence → (sorted de-duplicated tokens, numeric quantity fields). QTY spans that
    /// canonicalise (via `units`) become `(field, si_value)` pairs for `(num …)` range queries.
    pub fn project_sentence(&mut self, sentence: &str) -> (Vec<String>, Vec<(String, f64)>) {
        let mut tokens: Vec<String> = Vec::new();
        let mut numbers: Vec<(String, f64)> = Vec::new();
        if let Ok(spans) = self.tagger.tag(sentence) {
            for s in spans {
                let v = slug(&s.text);
                if !v.is_empty() {
                    let facet = facet_for(&s.kind);
                    let raw = format!("{facet}/{v}");
                    // Growing gazetteer + entity registration: for multi-word entity/place spans, register
                    // a CANONICAL token (content-word core, determiner/role variants collapsed) so the same
                    // entity's postings merge across surface variants, and future ingest + query recognise
                    // it. Emit the canonical token in place of the drifting raw span. Single words are
                    // already well-covered by the tagger/SPLADE tiers.
                    let mut token = raw;
                    if self.overlay_path.is_some() && matches!(s.kind.as_str(), "ENT" | "GEO") && s.text.split_whitespace().count() >= 2 {
                        if let Some(gaz) = self.gaz.as_mut() {
                            if let Some(canon) = gaz.register(&s.text, &facet) {
                                token = canon;
                            }
                        }
                    }
                    tokens.push(token);
                }
                if s.kind == "QTY" {
                    if let Some(nf) = crate::units::parse_quantity(&s.text) {
                        numbers.push(nf);
                    }
                }
            }
        }
        if let Some(splade) = self.splade.as_mut() {
            if let Ok(terms) = splade.project(sentence, 8, 0.3) {
                for t in terms {
                    // NOISE REGISTRATION: drop the ultra-generic, non-discriminative facet leaves.
                    if SPLADE_NOISE.contains(&leaf_of(&t.token)) {
                        continue;
                    }
                    tokens.push(t.token);
                }
            }
        }
        // GAZETTEER tier: high-resolution whole-entity `facet/value` tokens (aws-service/amazon-eks,
        // capability/zero-etl) — deterministic, span-grounded, the discriminative tokens the analytics
        // programs reason over.
        if let Some(gaz) = self.gaz.as_ref() {
            for h in gaz.extract(sentence) {
                tokens.push(h.token);
            }
        }
        tokens.sort();
        tokens.dedup();
        (tokens, numbers)
    }

    /// Project one sentence to a sorted, de-duplicated token set (tokens only).
    pub fn tokens_for(&mut self, sentence: &str) -> Vec<String> {
        self.project_sentence(sentence).0
    }

    /// Drive a sink with one Situation per sentence of `text`.
    pub fn project_text(&mut self, text: &str, sink: &mut dyn FnMut(Situation)) {
        for sent in sentences(text) {
            sink(self.project_situation(sent));
        }
    }
}

pub struct TextProjector {
    path: PathBuf,
    engine: TextEngine,
}

impl TextProjector {
    pub fn open(path: impl Into<PathBuf>, ml_bundle: &Path, splade_dir: Option<&Path>) -> Res<TextProjector> {
        Ok(TextProjector { path: path.into(), engine: TextEngine::load(ml_bundle, splade_dir)? })
    }
}

/// Split into sentences on Latin + CJK terminators and newlines.
pub fn sentences(text: &str) -> Vec<&str> {
    let mut out = Vec::new();
    let mut start = 0;
    for (i, ch) in text.char_indices() {
        if matches!(ch, '.' | '!' | '?' | '。' | '！' | '？' | '\n') {
            let end = i + ch.len_utf8();
            let s = text[start..end].trim();
            if s.len() >= 2 {
                out.push(s);
            }
            start = end;
        }
    }
    let tail = text[start..].trim();
    if tail.len() >= 2 {
        out.push(tail);
    }
    out
}

/// Resolve the gazetteer artifact: `STEELDB_GAZETTEER` env override, else `<splade_dir>/gazetteer.json`.
fn gazetteer_path(splade_dir: Option<&Path>) -> Option<PathBuf> {
    if let Ok(p) = std::env::var("STEELDB_GAZETTEER") {
        let p = PathBuf::from(p);
        if p.exists() {
            return Some(p);
        }
    }
    splade_dir.map(|d| d.join("gazetteer.json")).filter(|p| p.exists())
}

/// tagger span kind → token facet prefix
/// Map a tagger span kind to its infon facet prefix. The base tagger emits `ENT/REL/GEO/TIME/QTY`; a
/// domain-adapted tagger emits **typed** kinds (`ORG`, `ARTIFACT`, `PLATFORM`, `GENE`, …) — those pass
/// through as their own prefix (lowercased) so `org/toyota`, `artifact/battery_cell` etc. work directly
/// (Vocabulary-Space dim 1). Typed relations carry polarity as a suffix kind (`REL+`/`REL-` → `rel/…/+`).
fn facet_for(kind: &str) -> String {
    match kind {
        "ENT" => "ent".to_string(),
        "REL" => "rel".to_string(),
        "GEO" => "geo".to_string(),
        "TIME" => "time".to_string(),
        "QTY" => "qty".to_string(),
        "" | "IGNORE" | "O" => "misc".to_string(),
        // typed kind from a domain-adapted tagger → prefix is the kind itself, lowercased + slugged.
        other => other.chars().filter(|c| c.is_alphanumeric() || *c == '-').flat_map(|c| c.to_lowercase()).collect(),
    }
}

impl Projector for TextProjector {
    fn columns(&self) -> Vec<String> {
        vec!["sentence".to_string()]
    }
    fn kind(&self) -> CorpusKind {
        CorpusKind::Text
    }
    fn source(&self) -> String {
        self.path.display().to_string()
    }
    fn project(mut self: Box<Self>, sink: &mut dyn FnMut(Situation)) -> std::io::Result<()> {
        let text = std::fs::read_to_string(&self.path)?;
        self.engine.project_text(&text, sink);
        Ok(())
    }
}