steeldb/

spans.rs

//! **Span boundary repair** — turn model predictions over sub-word pieces into spans over whole words.
//!
//! A token classifier predicts over wordpieces, not words. "Registeel" is tokenised as something like
//! `Reg ##ist ##eel`, and if the `B-ENT` label lands on `##ist` the decoded span is the three characters
//! `ist`. The prediction was right about *where* the entity is and wrong about where it *starts*, which is a
//! boundary problem, not a classification problem — so it is repaired here rather than trained around.
//!
//! Left unrepaired this is corrosive rather than merely untidy: fragments become gazetteer entries, index
//! tokens and clustering inputs, so an entity is stored under a name that appears nowhere in the text and can
//! never be matched again. A real run of the exporter produced `itar` for Registeel, `topolis City` for
//! Sootopolis City and `202` for 2025, and the resulting facet codebook was debris.
//!
//! Kept free of feature gates and model dependencies so every consumer shares one implementation: the ONNX
//! tagger, the tuned tagger, and offline export all need the same repair.

/// Expand `[start, end)` so both edges sit on word boundaries in `text`.
///
/// Byte offsets are snapped to character boundaries first, so a span cutting a multi-byte character cannot
/// panic when it is sliced. A word is alphanumeric plus `'` and `-`, which keeps "Ward's" and "11-minute"
/// whole.
pub fn snap_to_words(text: &str, start: usize, end: usize) -> (usize, usize) {
    let is_word = |c: char| c.is_alphanumeric() || c == '\'' || c == '-';

    let mut s = start.min(text.len());
    let mut e = end.clamp(s, text.len());
    while s > 0 && !text.is_char_boundary(s) {
        s -= 1;
    }
    while e < text.len() && !text.is_char_boundary(e) {
        e += 1;
    }

    // walk left while the character before the span continues the same word
    while s > 0 {
        match text[..s].chars().next_back() {
            Some(c) if is_word(c) => s -= c.len_utf8(),
            _ => break,
        }
    }
    // walk right while the character after the span continues the word
    while e < text.len() {
        match text[e..].chars().next() {
            Some(c) if is_word(c) => e += c.len_utf8(),
            _ => break,
        }
    }
    (s, e)
}

/// Snap a list of `(start, end, kind)` ranges and merge neighbours of the same kind that snapping brought
/// into contact.
///
/// Merging is required, not cosmetic: two pieces of one word are frequently labelled separately, and once both
/// expand to the whole word they become duplicates covering identical text. Ranges separated only by
/// whitespace are also joined, so a multi-word mention predicted piecewise arrives as one span.
pub fn snap_and_merge(text: &str, ranges: &[(usize, usize, String)]) -> Vec<(usize, usize, String)> {
    let mut snapped: Vec<(usize, usize, String)> = ranges
        .iter()
        .map(|(s, e, k)| {
            let (ss, se) = snap_to_words(text, *s, *e);
            (ss, se, k.clone())
        })
        .filter(|(s, e, _)| e > s)
        .collect();

    snapped.sort_by(|a, b| a.0.cmp(&b.0).then(b.1.cmp(&a.1)));

    let mut out: Vec<(usize, usize, String)> = Vec::new();
    for (s, e, kind) in snapped {
        // same kind, and either overlapping or separated only by whitespace → one span
        let joins = out.last().is_some_and(|(_, pe, pk)| {
            *pk == kind && (s <= *pe || text.get(*pe..s).is_some_and(|gap| gap.chars().all(char::is_whitespace)))
        });
        if joins {
            if let Some((_, pe, _)) = out.last_mut() {
                *pe = (*pe).max(e);
            }
        } else {
            out.push((s, e, kind));
        }
    }
    out
}

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

    #[test]
    fn a_mid_word_prediction_recovers_the_whole_word() {
        let t = "Registeel took the win";
        // the model labelled only "ist" inside Registeel
        let (s, e) = snap_to_words(t, 3, 6);
        assert_eq!(&t[s..e], "Registeel");
    }

    #[test]
    fn a_four_digit_year_is_not_truncated() {
        let t = "held in 2025 at the venue";
        let (s, e) = snap_to_words(t, 8, 11); // "202"
        assert_eq!(&t[s..e], "2025");
    }

    #[test]
    fn an_already_aligned_span_is_unchanged() {
        let t = "Sootopolis City hosted it";
        let (s, e) = snap_to_words(t, 0, 15);
        assert_eq!(&t[s..e], "Sootopolis City");
    }

    #[test]
    fn multibyte_text_does_not_panic_and_slices_cleanly() {
        // é and ° are multi-byte; a byte-wise walk panics here
        for t in ["a Pokémon named Aggron", "recorded 28 °C at the site", "Café Ecruteak"] {
            for start in 0..t.len() {
                for end in start..t.len() {
                    let (s, e) = snap_to_words(t, start, end);
                    assert!(t.get(s..e).is_some(), "span {s}..{e} must slice {t:?}");
                }
            }
        }
    }

    #[test]
    fn apostrophes_and_hyphens_stay_inside_one_word() {
        let t = "Cynthia Ward's 11-minute battle";
        let (s, e) = snap_to_words(t, 8, 12); // inside "Ward's"
        assert_eq!(&t[s..e], "Ward's");
        let (s2, e2) = snap_to_words(t, 15, 17); // inside "11-minute"
        assert_eq!(&t[s2..e2], "11-minute");
    }

    #[test]
    fn two_pieces_of_one_word_merge_to_a_single_span() {
        let t = "Sootopolis City";
        // predicted as two fragments of the same kind
        let ranges = vec![
            (0usize, 4usize, "GEO".to_string()),   // "Soot"
            (4usize, 10usize, "GEO".to_string()),  // "opolis"
        ];
        let out = snap_and_merge(t, &ranges);
        assert_eq!(out.len(), 1, "{out:?}");
        assert_eq!(&t[out[0].0..out[0].1], "Sootopolis");
    }

    #[test]
    fn a_multi_word_mention_predicted_piecewise_becomes_one_span() {
        let t = "at Sootopolis City today";
        let ranges = vec![
            (3usize, 8usize, "GEO".to_string()),   // "Sooto"
            (14usize, 18usize, "GEO".to_string()), // "City"
        ];
        let out = snap_and_merge(t, &ranges);
        assert_eq!(out.len(), 1, "{out:?}");
        assert_eq!(&t[out[0].0..out[0].1], "Sootopolis City");
    }

    #[test]
    fn different_kinds_are_never_merged() {
        let t = "Aggron 1082 m";
        let ranges = vec![
            (0usize, 6usize, "ENT".to_string()),
            (7usize, 13usize, "QTY".to_string()),
        ];
        let out = snap_and_merge(t, &ranges);
        assert_eq!(out.len(), 2, "adjacent spans of different kinds must stay separate: {out:?}");
    }

    #[test]
    fn spans_separated_by_real_words_stay_separate() {
        let t = "Aggron and also Salamence";
        let ranges = vec![
            (0usize, 6usize, "ENT".to_string()),
            (16usize, 25usize, "ENT".to_string()),
        ];
        let out = snap_and_merge(t, &ranges);
        assert_eq!(out.len(), 2, "{out:?}");
    }
}