steeldb/

tagger_discover.rs

//! **Stage 1-2 of vocabulary discovery: the tagger and the codebook.**
//!
//! This produces a **raw spec**, not an ontology, and the distinction is the whole design:
//!
//! 1. the SPO tagger reads each sentence and marks typed spans, each arriving already labelled by kind;
//! 2. every span becomes a 128-dimensional point via the model2vec embedder;
//! 3. **within each kind separately**, k-means proposes prototypes and Sinkhorn optimal transport assigns spans
//!    to them, each cluster named by its most frequent surface term.
//!
//! What comes out is a list of raw **entity-value clusters** — and a cluster label is a *value*, not a facet.
//! The reference's own committed spec makes this plain: its raw clusters are `missile`, `bo`, `guidance
//! control`, `ph`, `ge`, and only after curation does the ontology read `weapon-platform`, `control-system`,
//! with those raw terms demoted to examples. Treating a cluster label as a facet name yields exactly that
//! junk — `ph` and `ge` as retrieval dimensions.
//!
//! So this module deliberately stops short of naming facets. Turning raw clusters into canonical MECE facet
//! *types* — merging synonyms, dropping boilerplate, naming the kind — is judgment, and it belongs to
//! [`crate::learn`]'s curation step.
//!
//! The six dimensions of V are handled by kind, as the reference does:
//!
//! | kind | treatment |
//! |---|---|
//! | `ENT`, `GEO` | clustered into raw entity-value clusters, for curation |
//! | `REL` | clustered into raw relation clusters, for curation |
//! | `TIME` | not clustered — the deterministic temporal extractor buckets these exactly |
//! | `QTY` | routed OUT of the codebook: a measurement is a value on a scale, not a kind of thing |
//! | `IGNORE` | dropped by the tagger |
//!
//! Native-only: the tagger is ONNX and the embedder's tokenizer links a C library, so neither targets wasm32.

use crate::discover_ontology::OtDiscover;
use crate::text::{tagger::SpoTagger, Model2Vec};

/// The kinds that form entity-value clusters. `REL` is clustered separately; `QTY`/`TIME` never enter a
/// codebook.
const ENTITY_KINDS: &[&str] = &["ENT", "GEO"];

/// One raw cluster: a label drawn from its most representative member, and the members themselves.
///
/// The label is a *value*, not a facet name. It exists to give the curator something to refer to.
#[derive(Debug, Clone)]
pub struct RawCluster {
    /// the most representative member, used only as a handle for curation
    pub label: String,
    /// the cluster's members, most representative first — the surface forms
    pub terms: Vec<String>,
}

/// What the tagger and the codebook produce: candidate clusters awaiting curation.
#[derive(Debug, Clone, Default)]
pub struct RawSpec {
    /// entity-value clusters from the `ENT` and `GEO` kinds
    pub entity_clusters: Vec<RawCluster>,
    /// relation clusters from the `REL` kind
    pub relation_clusters: Vec<RawCluster>,
    /// how many documents were read
    pub documents: usize,
    /// spans routed away from the codebook, by kind — reported so a caller can see the routing happened
    pub routed_away: Vec<(String, usize)>,
}

impl RawSpec {
    /// Every entity surface form seen, longest first so gazetteer matching prefers the full mention.
    ///
    /// These are the *observed* spans, independent of how curation groups them, so the gazetteer survives
    /// whatever the curator decides to merge or drop.
    pub fn surfaces(&self) -> Vec<String> {
        let mut out: Vec<String> = Vec::new();
        for c in &self.entity_clusters {
            for t in &c.terms {
                if !out.contains(t) {
                    out.push(t.clone());
                }
            }
        }
        out.sort_by(|a, b| b.chars().count().cmp(&a.chars().count()).then(a.cmp(b)));
        out
    }
}

/// How discovery could not proceed.
#[derive(Debug)]
pub enum TaggerError {
    /// a model file was not found; carries the resolver's instruction, which names repo, revision and size
    Model(String),
    /// the tagger or embedder failed on input
    Inference(String),
}

impl std::fmt::Display for TaggerError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TaggerError::Model(m) => write!(f, "model not available: {m}"),
            TaggerError::Inference(m) => write!(f, "tagger failed: {m}"),
        }
    }
}

impl std::error::Error for TaggerError {}

struct Span {
    kind: String,
    text: String,
    vec: Vec<f32>,
}

/// Tag a sample, embed the spans, and build a per-kind codebook.
///
/// `k_ent`/`k_rel` are the reference's `ke`/`kr` (12 and 8) — upper bounds, which the solver clamps down when a
/// kind has few spans.
pub fn discover(sample: &[String], k_ent: usize, k_rel: usize) -> Result<RawSpec, TaggerError> {
    let tagger_dir = crate::models::resolve("spo-tagger").map_err(|e| TaggerError::Model(e.to_string()))?;
    let m2v_dir = crate::models::resolve("model2vec").map_err(|e| TaggerError::Model(e.to_string()))?;
    let mut tagger = SpoTagger::load(&tagger_dir).map_err(|e| TaggerError::Inference(e.to_string()))?;
    let embedder = Model2Vec::load(&m2v_dir).map_err(|e| TaggerError::Inference(e.to_string()))?;

    let mut spans: Vec<Span> = Vec::new();
    let mut routed: std::collections::BTreeMap<String, usize> = std::collections::BTreeMap::new();
    for doc in sample {
        for sentence in doc.split(['.', ';', '!', '?', '\n']) {
            let s = sentence.trim();
            if s.is_empty() {
                continue;
            }
            let tagged = tagger.tag(s).map_err(|e| TaggerError::Inference(e.to_string()))?;
            for sp in tagged {
                // QTY is a value on a scale and TIME is bucketed exactly without a model, so neither belongs in
                // a semantic codebook. Counted rather than silently discarded, so the routing is visible.
                if sp.kind == "QTY" || sp.kind == "TIME" {
                    *routed.entry(sp.kind).or_default() += 1;
                    continue;
                }
                if let Some(vec) = embedder.embed(&sp.text) {
                    spans.push(Span { kind: sp.kind, text: sp.text, vec });
                }
            }
        }
    }

    let mut entity_clusters = Vec::new();
    for kind in ENTITY_KINDS {
        let (terms, embs) = by_kind(&spans, kind);
        entity_clusters.extend(codebook(terms, embs, k_ent));
    }
    let (rel_terms, rel_embs) = by_kind(&spans, "REL");
    let relation_clusters = codebook(rel_terms, rel_embs, k_rel);

    Ok(RawSpec {
        entity_clusters,
        relation_clusters,
        documents: sample.len(),
        routed_away: routed.into_iter().collect(),
    })
}

fn by_kind(spans: &[Span], kind: &str) -> (Vec<String>, Vec<Vec<f32>>) {
    let mut terms = Vec::new();
    let mut embs = Vec::new();
    for s in spans.iter().filter(|s| s.kind == kind) {
        terms.push(s.text.clone());
        embs.push(s.vec.clone());
    }
    (terms, embs)
}

/// Cluster one kind: k-means++ prototypes, Sinkhorn assignment, members ordered by closeness.
fn codebook(terms: Vec<String>, embs: Vec<Vec<f32>>, k: usize) -> Vec<RawCluster> {
    if terms.len() < 4 {
        return Vec::new();
    }
    let d = OtDiscover::new(terms, embs, k);
    // 200 Sinkhorn iterations, matching the reference; assignment is the argmax of the transport plan
    let (assign, _cost) = d.assign(200);
    d.clusters(&assign, 12)
        .into_iter()
        .map(|c| RawCluster { label: c.label, terms: c.terms })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn surfaces_are_deduplicated_and_longest_first() {
        // The gazetteer is built from observed spans, not from curated facets, so it survives whatever the
        // curator merges or drops. Longest-first ordering is what makes a full mention beat its own prefix.
        let spec = RawSpec {
            entity_clusters: vec![
                RawCluster {
                    label: "city".into(),
                    terms: vec!["Sootopolis".into(), "Sootopolis City".into(), "Sootopolis".into()],
                },
                RawCluster { label: "region".into(), terms: vec!["Johto".into()] },
            ],
            relation_clusters: Vec::new(),
            documents: 2,
            routed_away: Vec::new(),
        };
        let s = spec.surfaces();
        assert_eq!(s, vec!["Sootopolis City", "Sootopolis", "Johto"], "{s:?}");
    }

    #[test]
    fn a_raw_cluster_label_is_not_treated_as_a_facet_name() {
        // Guards the design this module exists to respect: nothing here promotes a cluster label to a facet.
        // The reference's raw spec contains `ph` and `ge`; only curation turns clusters into named types, so a
        // RawSpec must expose clusters and surfaces and NOT a `categories` list.
        let spec = RawSpec::default();
        assert!(spec.entity_clusters.is_empty() && spec.surfaces().is_empty());
    }
}