hypersteeldb 0.5.5

//! SPO span tagger — runs the exported BERT token-classifier (`spo.onnx`) and BIO-decodes
//! ENT / REL / QTY / GEO / TIME / IGNORE spans. Ported from `src/ingest/spo/tagger.ts`.
//!
//! We use the `tokenizers` crate's native char offsets + special-token mask instead of the TS manual
//! ▁/## reconstruction — correct for any script (EN/JA/KO) and far less fragile.

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

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

#[derive(Debug, Clone)]
pub struct Span {
    pub kind: String,
    pub start: usize,
    pub end: usize,
    pub text: String,
}

pub struct SpoTagger {
    session: Session,
    tok: Tokenizer,
    labels: Vec<String>,
}

impl SpoTagger {
    /// `dir` holds `spo.onnx`, `labels.json`, and `tokenizer/tokenizer.json` (the mmBERT ML bundle).
    pub fn load(dir: &Path) -> Res<SpoTagger> {
        let session = Session::builder()?.commit_from_file(dir.join("spo.onnx"))?;
        let mut tok = Tokenizer::from_file(dir.join("tokenizer/tokenizer.json"))?;
        // match training max_length (spo_tagger.py = 256); else trailing tokens are silently dropped.
        tok.with_truncation(Some(tokenizers::TruncationParams {
            max_length: 256,
            ..Default::default()
        }))?;
        let labels: Vec<String> = serde_json::from_slice(&std::fs::read(dir.join("labels.json"))?)?;
        Ok(SpoTagger { session, tok, labels })
    }

    pub fn tag(&mut self, sentence: &str) -> Res<Vec<Span>> {
        let enc = self.tok.encode(sentence, true)?;
        let ids: Vec<i64> = enc.get_ids().iter().map(|&x| x as i64).collect();
        let attn: Vec<i64> = enc.get_attention_mask().iter().map(|&x| x as i64).collect();
        let t = ids.len();

        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, logits) = outputs["logits"].try_extract_tensor::<f32>()?;
        // shape = [1, T, L]
        let l = *shape.last().unwrap() as usize;

        let offsets = enc.get_offsets();
        let special = enc.get_special_tokens_mask();

        let mut spans: Vec<Span> = Vec::new();
        let mut cur: Option<Span> = None;
        for ti in 0..t {
            if special[ti] == 1 {
                continue;
            }
            let (a, b) = offsets[ti];
            if b <= a {
                continue;
            }
            // argmax label at this token
            let base = ti * l;
            let mut arg = 0usize;
            let mut best = f32::NEG_INFINITY;
            for li in 0..l {
                let v = logits[base + li];
                if v > best {
                    best = v;
                    arg = li;
                }
            }
            let lab = &self.labels[arg];
            if let Some(kind) = lab.strip_prefix("B-") {
                if let Some(s) = cur.take() {
                    spans.push(s);
                }
                cur = Some(Span { kind: kind.to_string(), start: a, end: b, text: String::new() });
            } else if let Some(kind) = lab.strip_prefix("I-") {
                match cur.as_mut() {
                    Some(s) if s.kind == kind => s.end = b,
                    _ => {
                        if let Some(s) = cur.take() {
                            spans.push(s);
                        }
                    }
                }
            } else {
                // "O"
                if let Some(s) = cur.take() {
                    spans.push(s);
                }
            }
        }
        if let Some(s) = cur.take() {
            spans.push(s);
        }
        // The offsets above are WORDPIECE offsets, so a B- label landing mid-word yields a fragment:
        // "Registeel" decodes as "itar", "Sootopolis City" as "topolis City", "2025" as "202". The prediction
        // is right about where the entity is and wrong about where it starts, which is a boundary problem and
        // is repaired here. Unrepaired, fragments flow into the gazetteer, the index and the facet codebook,
        // so an entity ends up stored under a name that appears nowhere in the text.
        let ranges: Vec<(usize, usize, String)> =
            spans.iter().map(|s| (s.start, s.end, s.kind.clone())).collect();
        let mut spans: Vec<Span> = crate::spans::snap_and_merge(sentence, &ranges)
            .into_iter()
            .map(|(start, end, kind)| Span {
                kind,
                start,
                end,
                text: sentence.get(start..end).unwrap_or("").to_string(),
            })
            .collect();

        // IGNORE spans are non-content; drop them.
        spans.retain(|s| s.kind != "IGNORE" && !s.text.trim().is_empty());
        Ok(spans)
    }
}