steeldb/text/
tagger.rs

1//! SPO span tagger — runs the exported BERT token-classifier (`spo.onnx`) and BIO-decodes
2//! ENT / REL / QTY / GEO / TIME / IGNORE spans. Ported from `src/ingest/spo/tagger.ts`.
3//!
4//! We use the `tokenizers` crate's native char offsets + special-token mask instead of the TS manual
5//! ▁/## reconstruction — correct for any script (EN/JA/KO) and far less fragile.
6
7use ort::session::Session;
8use ort::value::Tensor;
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(Debug, Clone)]
16pub struct Span {
17    pub kind: String,
18    pub start: usize,
19    pub end: usize,
20    pub text: String,
21}
22
23pub struct SpoTagger {
24    session: Session,
25    tok: Tokenizer,
26    labels: Vec<String>,
27}
28
29impl SpoTagger {
30    /// `dir` holds `spo.onnx`, `labels.json`, and `tokenizer/tokenizer.json` (the mmBERT ML bundle).
31    pub fn load(dir: &Path) -> Res<SpoTagger> {
32        let session = Session::builder()?.commit_from_file(dir.join("spo.onnx"))?;
33        let mut tok = Tokenizer::from_file(dir.join("tokenizer/tokenizer.json"))?;
34        // match training max_length (spo_tagger.py = 256); else trailing tokens are silently dropped.
35        tok.with_truncation(Some(tokenizers::TruncationParams {
36            max_length: 256,
37            ..Default::default()
38        }))?;
39        let labels: Vec<String> = serde_json::from_slice(&std::fs::read(dir.join("labels.json"))?)?;
40        Ok(SpoTagger { session, tok, labels })
41    }
42
43    pub fn tag(&mut self, sentence: &str) -> Res<Vec<Span>> {
44        let enc = self.tok.encode(sentence, true)?;
45        let ids: Vec<i64> = enc.get_ids().iter().map(|&x| x as i64).collect();
46        let attn: Vec<i64> = enc.get_attention_mask().iter().map(|&x| x as i64).collect();
47        let t = ids.len();
48
49        let id_t = Tensor::from_array(([1usize, t], ids))?;
50        let at_t = Tensor::from_array(([1usize, t], attn))?;
51        let outputs = self
52            .session
53            .run(ort::inputs!["input_ids" => id_t, "attention_mask" => at_t])?;
54        let (shape, logits) = outputs["logits"].try_extract_tensor::<f32>()?;
55        // shape = [1, T, L]
56        let l = *shape.last().unwrap() as usize;
57
58        let offsets = enc.get_offsets();
59        let special = enc.get_special_tokens_mask();
60
61        let mut spans: Vec<Span> = Vec::new();
62        let mut cur: Option<Span> = None;
63        for ti in 0..t {
64            if special[ti] == 1 {
65                continue;
66            }
67            let (a, b) = offsets[ti];
68            if b <= a {
69                continue;
70            }
71            // argmax label at this token
72            let base = ti * l;
73            let mut arg = 0usize;
74            let mut best = f32::NEG_INFINITY;
75            for li in 0..l {
76                let v = logits[base + li];
77                if v > best {
78                    best = v;
79                    arg = li;
80                }
81            }
82            let lab = &self.labels[arg];
83            if let Some(kind) = lab.strip_prefix("B-") {
84                if let Some(s) = cur.take() {
85                    spans.push(s);
86                }
87                cur = Some(Span { kind: kind.to_string(), start: a, end: b, text: String::new() });
88            } else if let Some(kind) = lab.strip_prefix("I-") {
89                match cur.as_mut() {
90                    Some(s) if s.kind == kind => s.end = b,
91                    _ => {
92                        if let Some(s) = cur.take() {
93                            spans.push(s);
94                        }
95                    }
96                }
97            } else {
98                // "O"
99                if let Some(s) = cur.take() {
100                    spans.push(s);
101                }
102            }
103        }
104        if let Some(s) = cur.take() {
105            spans.push(s);
106        }
107        // The offsets above are WORDPIECE offsets, so a B- label landing mid-word yields a fragment:
108        // "Registeel" decodes as "itar", "Sootopolis City" as "topolis City", "2025" as "202". The prediction
109        // is right about where the entity is and wrong about where it starts, which is a boundary problem and
110        // is repaired here. Unrepaired, fragments flow into the gazetteer, the index and the facet codebook,
111        // so an entity ends up stored under a name that appears nowhere in the text.
112        let ranges: Vec<(usize, usize, String)> =
113            spans.iter().map(|s| (s.start, s.end, s.kind.clone())).collect();
114        let mut spans: Vec<Span> = crate::spans::snap_and_merge(sentence, &ranges)
115            .into_iter()
116            .map(|(start, end, kind)| Span {
117                kind,
118                start,
119                end,
120                text: sentence.get(start..end).unwrap_or("").to_string(),
121            })
122            .collect();
123
124        // IGNORE spans are non-content; drop them.
125        spans.retain(|s| s.kind != "IGNORE" && !s.text.trim().is_empty());
126        Ok(spans)
127    }
128}
steeldb/text/tagger.rs

steeldb/text/
tagger.rs
