use std::collections::{HashMap, HashSet};
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TermCluster {
pub label: String,
pub terms: Vec<String>,
pub coverage: f64,
pub cohesion: f64,
}
pub fn common_nouns(docs: &[String]) -> HashSet<String> {
let mut lower_seen: HashSet<String> = HashSet::new();
for doc in docs {
for raw in doc.split(|c: char| !c.is_alphanumeric() && c != '-' && c != '\'') {
let t = raw.trim_matches('-');
if t.len() < 3 {
continue;
}
if t.chars().next().map(|c| c.is_lowercase()).unwrap_or(false) {
lower_seen.insert(t.to_lowercase());
}
}
}
lower_seen
}
fn tokenize(s: &str) -> Vec<String> {
s.split(|c: char| !c.is_alphanumeric() && c != '-' && c != '\'')
.map(|w| w.trim_matches('-').to_lowercase())
.filter(|w| {
w.len() >= 3
&& w.len() <= 28
&& w.chars().next().map(|c| c.is_alphabetic()).unwrap_or(false)
&& !w.chars().all(|c| c.is_ascii_digit())
})
.collect()
}
const STOP: &[&str] = &[
"the", "and", "for", "with", "was", "were", "this", "that", "from", "into", "are", "has", "had", "his",
"her", "its", "not", "but", "all", "any", "may", "can", "will", "each", "than", "then", "during",
"under", "over", "also", "which", "while", "their", "there", "been", "being", "who", "when", "what",
"how", "why", "per", "via", "such", "more", "most", "less", "other", "some", "one", "two", "three",
"against", "after", "before", "between", "both", "out", "off", "own", "same", "too", "very", "just",
"him", "she", "they", "them", "these", "those", "have", "does", "did", "doing", "would", "could",
"should", "must", "shall", "about", "above", "below", "again", "further", "once", "here", "only",
"remains", "stands", "recorded", "reported", "held", "took", "made", "including", "included",
];
fn salient_terms(docs: &[String], n_terms: usize) -> Vec<(String, HashSet<usize>)> {
let mut incidence: HashMap<String, HashSet<usize>> = HashMap::new();
let mut tf: HashMap<String, usize> = HashMap::new();
for (i, doc) in docs.iter().enumerate() {
for w in tokenize(doc) {
if STOP.contains(&w.as_str()) || GENERIC.contains(&w.as_str()) || LOCATIVES.contains(&w.as_str()) {
continue;
}
*tf.entry(w.clone()).or_default() += 1;
incidence.entry(w).or_default().insert(i);
}
}
let n = docs.len().max(1) as f64;
let min_df = 2usize;
let max_df = ((n * 0.85).ceil() as usize).max(min_df + 1);
let mut scored: Vec<(String, f64)> = incidence
.iter()
.filter(|(_, docs_in)| docs_in.len() >= min_df && docs_in.len() <= max_df)
.map(|(term, _)| (term.clone(), *tf.get(term).unwrap_or(&1) as f64))
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal).then(a.0.cmp(&b.0)));
scored.truncate(n_terms);
scored
.into_iter()
.map(|(term, _)| {
let docs_in = incidence.remove(&term).unwrap_or_default();
(term, docs_in)
})
.collect()
}
pub fn discover_motifs(docs: &[String], n_terms: usize, k: usize) -> Vec<(String, Vec<String>)> {
let (terms, vecs) = term_vectors(docs, n_terms);
if terms.len() < 2 || k == 0 {
return Vec::new();
}
let (_protos, assign, _cost) = crate::text::ot::codebook(&vecs, k, MOTIF_EPS);
let groups = assign.iter().copied().max().map_or(0, |m| m + 1);
let mut members: Vec<Vec<String>> = vec![Vec::new(); groups];
for (i, &a) in assign.iter().enumerate() {
members[a].push(terms[i].clone());
}
let commons = common_nouns(docs);
members
.into_iter()
.filter(|g| !g.is_empty())
.filter_map(|g| {
let name = g.iter().find(|t| commons.contains(*t)).or_else(|| g.first())?.clone();
(!name.is_empty()).then_some((name, g))
})
.collect()
}
const MOTIF_EPS: f32 = 0.03;
const DISCOVER_EPS: f32 = 0.05;
pub fn term_vectors(docs: &[String], n_terms: usize) -> (Vec<String>, Vec<Vec<f32>>) {
let picked = salient_terms(docs, n_terms);
let n = docs.len().max(1);
let mut names = Vec::with_capacity(picked.len());
let mut vecs = Vec::with_capacity(picked.len());
for (term, docs_in) in picked {
let mut v = vec![0f32; n];
for i in &docs_in {
if *i < n {
v[*i] = 1.0;
}
}
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt() + 1e-9;
vecs.push(v.into_iter().map(|x| x / norm).collect());
names.push(term);
}
(names, vecs)
}
pub fn quantity_spans(doc: &str) -> Vec<(usize, usize, String)> {
const UNITS: &[(&str, &str)] = &[
("mm", "length_mm"), ("cm", "length_cm"), ("km", "length_km"), ("m", "length_m"),
("kg", "mass_kg"), ("g", "mass_g"), ("t", "mass_t"),
("°c", "temp_c"), ("°f", "temp_f"), ("c", "temp_c"),
("minutes", "minutes"), ("minute", "minutes"), ("min", "minutes"),
("hours", "hours"), ("hour", "hours"), ("seconds", "seconds"),
("mm/yr", "rainfall_mm"), ("%", "percent"),
];
let b = doc.as_bytes();
let mut out = Vec::new();
let mut i = 0usize;
while i < b.len() {
if b[i].is_ascii_digit() && (i == 0 || !(b[i - 1] as char).is_alphanumeric()) {
let mut start = i;
if i > 0 && b[i - 1] == b'-' {
let before_sign = i >= 2 && !(b[i - 2] as char).is_alphanumeric() && b[i - 2] != b'-';
if i == 1 || before_sign {
start = i - 1;
}
}
let mut j = i;
while j < b.len() && (b[j].is_ascii_digit() || b[j] == b'.' || b[j] == b',') {
j += 1;
}
let num_end = j;
let mut k = j;
while k < b.len() && (b[k] == b' ' || b[k] == b'-') {
k += 1;
}
if k < b.len() && doc.is_char_boundary(k) {
let rest = &doc[k..];
let unit_len = rest
.char_indices()
.take_while(|(_, c)| c.is_alphabetic() || *c == '°' || *c == '%' || *c == '/')
.map(|(bi, c)| bi + c.len_utf8())
.last()
.unwrap_or(0);
if unit_len > 0 {
let unit = rest[..unit_len].to_lowercase();
let mut best: Option<(&str, usize)> = None;
for (u, field) in UNITS {
if unit == *u && best.map(|(_, l)| u.len() > l).unwrap_or(true) {
best = Some((field, u.len()));
}
}
if let Some((field, ulen)) = best {
let end = k + ulen;
if doc.is_char_boundary(start) && doc.is_char_boundary(end) {
out.push((start, end, field.to_string()));
i = end;
continue;
}
}
}
}
i = num_end.max(i + 1);
continue;
}
i += 1;
}
out
}
pub fn contains_term(hay: &str, needle: &str) -> bool {
!word_spans(hay, needle).is_empty()
}
const REL_VERBS: &[(&str, &str)] = &[
("defeated", "defeated"), ("beat", "defeated"), ("faced", "faced"), ("met", "faced"),
("documented", "documented"), ("recorded", "recorded"), ("measured", "measured"),
("observed", "observed"), ("found", "observed"), ("held", "held_at"), ("hosted", "held_at"),
("used", "used"), ("led", "used"), ("answered", "answered_with"), ("commanded", "commanded"),
("permitted", "permitted"), ("banned", "banned"), ("restricted", "restricted"),
("competed", "competed_in"), ("entered", "competed_in"), ("won", "won"), ("secured", "secured"),
("contributes", "contributes_to"), ("supplies", "supplies"), ("operates", "operates"),
];
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Relation {
pub verb: String,
pub actor: String,
pub target: String,
}
pub fn relation_spans(doc: &str, mentions: &[String]) -> Vec<Relation> {
let mut out: Vec<Relation> = Vec::new();
for sentence in doc.split(['.', ';', '!', '?', '\n']) {
if sentence.trim().is_empty() {
continue;
}
let mut found: Vec<(usize, usize, String)> = Vec::new();
for m in mentions {
for (s, e) in word_spans(sentence, m) {
if !found.iter().any(|(fs, fe, _)| s >= *fs && e <= *fe) {
found.push((s, e, m.clone()));
}
}
}
for (s, e, name) in local_mentions(sentence) {
if !found.iter().any(|(fs, fe, _)| s < *fe && e > *fs) {
found.push((s, e, name));
}
}
if found.len() < 2 {
continue;
}
found.sort_by_key(|(s, _, _)| *s);
for (raw, canon) in REL_VERBS {
for (vs, ve) in word_spans(sentence, raw) {
let actor = found.iter().filter(|(_, e, _)| *e <= vs).next_back();
let target = found.iter().find(|(s, _, _)| *s >= ve);
if let (Some((_, _, a)), Some((_, _, t))) = (actor, target) {
if a != t {
out.push(Relation { verb: canon.to_string(), actor: a.clone(), target: t.clone() });
}
}
}
}
}
out.dedup_by(|a, b| a.verb == b.verb && a.actor == b.actor && a.target == b.target);
out
}
pub fn local_mentions(sentence: &str) -> Vec<(usize, usize, String)> {
let is_word = |c: char| c.is_alphanumeric() || c == '\'' || c == '-';
let mut words: Vec<(usize, usize, &str)> = Vec::new();
let mut cur: Option<usize> = None;
for (i, c) in sentence.char_indices() {
if is_word(c) {
if cur.is_none() {
cur = Some(i);
}
} else if let Some(st) = cur.take() {
words.push((st, i, &sentence[st..i]));
}
}
if let Some(st) = cur {
words.push((st, sentence.len(), &sentence[st..]));
}
let mut runs: Vec<(usize, usize)> = Vec::new();
let mut run: Option<(usize, usize)> = None;
let mut prev_end: Option<usize> = None;
for (wi, (st, en, w)) in words.iter().enumerate() {
let capped = w.chars().next().map(|c| c.is_uppercase()).unwrap_or(false) && w.chars().count() > 1;
let punctuated = prev_end
.map(|pe| sentence[pe..*st].chars().any(|c| !c.is_whitespace()))
.unwrap_or(false);
if punctuated {
if let Some(r) = run.take() {
runs.push(r);
}
}
if capped {
run = Some(match run {
Some((rs, _)) => (rs, wi),
None => (wi, wi),
});
} else if let Some(r) = run.take() {
runs.push(r);
}
prev_end = Some(*en);
}
if let Some(r) = run {
runs.push(r);
}
const DETERMINERS: &[&str] = &[
"the", "a", "an", "this", "that", "these", "those", "their", "its", "his", "her", "our", "your", "my",
"it", "they", "we", "he", "she", "there", "then", "when", "where", "what", "which", "who",
];
let mut out: Vec<(usize, usize, String)> = Vec::new();
for (first, last) in runs {
let mut first = first;
if first == 0 && DETERMINERS.contains(&words[0].2.to_lowercase().as_str()) {
first += 1;
}
if first > last {
continue;
}
let (rs, re) = (words[first].0, words[last].1);
out.push((rs, re, sentence[rs..re].to_string()));
}
out
}
pub fn temporal_spans(doc: &str) -> Vec<(usize, usize, String)> {
let b = doc.as_bytes();
let mut out: Vec<(usize, usize, String)> = Vec::new();
let mut i = 0usize;
while i + 1 < b.len() {
if (b[i] == b'Q' || b[i] == b'q') && b[i + 1].is_ascii_digit() {
let q = (b[i + 1] - b'0') as u32;
let starts_word = i == 0 || !(b[i - 1] as char).is_alphanumeric();
if (1..=4).contains(&q) && starts_word {
let end = i + 2;
let tail = &doc[end..].trim_start();
let year: Option<u32> = tail
.split(|c: char| !c.is_ascii_digit())
.next()
.filter(|t| t.len() == 4)
.and_then(|t| t.parse().ok())
.filter(|y| (1900..2200).contains(y));
let token = match year {
Some(y) => format!("time/{y}/q{q}"),
None => format!("time/q{q}"),
};
if doc.is_char_boundary(i) && doc.is_char_boundary(end) {
out.push((i, end, token));
}
i = end;
continue;
}
}
i += 1;
}
let mut j = 0usize;
while j + 3 < b.len() {
if b[j].is_ascii_digit() {
let before_ok = j == 0 || !(b[j - 1] as char).is_alphanumeric();
let end = j + 4;
let after_ok = end >= b.len() || !(b[end] as char).is_alphanumeric();
if before_ok && after_ok && b[j..end].iter().all(|c| c.is_ascii_digit()) {
if let Ok(y) = doc[j..end].parse::<u32>() {
if (1900..2200).contains(&y) && !out.iter().any(|(s, e, _)| j >= *s && end <= *e) {
out.push((j, end, format!("time/{y}")));
}
}
j = end;
continue;
}
}
j += 1;
}
out.sort_by_key(|(s, _, _)| *s);
out
}
pub fn mine_gazetteer(docs: &[String], min_count: usize) -> Vec<String> {
let mut counts: HashMap<String, usize> = HashMap::new();
for doc in docs {
for sentence in doc.split(['.', '\n', ';', '!', '?']) {
let words: Vec<&str> = sentence.split_whitespace().collect();
let mut run: Vec<&str> = Vec::new();
let mut first = true;
for w in words {
let clean = w.trim_matches(|c: char| !c.is_alphanumeric() && c != '\'' && c != '-');
let cap = clean
.chars()
.next()
.map(|c| c.is_uppercase())
.unwrap_or(false)
&& clean.len() > 1;
if cap && !(first && run.is_empty()) {
run.push(clean);
} else {
if run.len() >= 2 {
*counts.entry(run.join(" ")).or_default() += 1;
}
run.clear();
}
first = false;
}
if run.len() >= 2 {
*counts.entry(run.join(" ")).or_default() += 1;
}
}
}
let mut kept: Vec<String> =
counts.into_iter().filter(|(_, n)| *n >= min_count).map(|(s, _)| s).collect();
kept.sort_by(|a, b| b.len().cmp(&a.len()).then(a.cmp(b)));
kept
}
pub const LOCATIVES: &[&str] = &[
"near", "nearby", "across", "along", "around", "through", "throughout", "toward", "towards",
"beside", "behind", "beyond", "upon", "onto", "inside", "outside", "amid", "among", "amongst",
"beneath", "underneath", "alongside", "opposite", "past", "since", "until", "till", "unto",
];
pub const GENERIC: &[&str] = &[
"within", "presence", "data", "contributes", "distribution", "environmental", "understanding",
"preferences", "observation", "site", "location", "conditions", "period", "mean", "documented",
"recorded", "measured", "reported", "described", "including", "various", "distinctive", "populations",
"ecological", "information", "details", "features", "aspects", "elements", "factors", "values",
"results", "analysis", "summary", "overview", "context", "purposes", "requirements",
];
pub fn word_spans(hay: &str, needle: &str) -> Vec<(usize, usize)> {
let mut out = Vec::new();
if needle.is_empty() {
return out;
}
let lower = hay.to_lowercase();
let pat = needle.to_lowercase();
let bytes = lower.as_bytes();
let mut from = 0usize;
while let Some(rel) = lower[from..].find(&pat) {
let s = from + rel;
let e = s + pat.len();
let before_ok = s == 0 || !(bytes[s - 1] as char).is_alphanumeric();
let after_ok = e >= bytes.len() || !(bytes[e] as char).is_alphanumeric();
if before_ok && after_ok && hay.is_char_boundary(s) && hay.is_char_boundary(e) {
out.push((s, e));
}
from = s + pat.len().max(1);
if from >= lower.len() {
break;
}
}
out
}
pub fn salient(docs: &[String], n_terms: usize) -> Vec<(String, f64, usize)> {
let picked = salient_terms(docs, n_terms);
let n = docs.len().max(1) as f64;
picked
.into_iter()
.map(|(term, docs_in)| {
let d = docs_in.len();
let idf = ((n + 1.0) / (d as f64 + 1.0)).ln() + 1.0;
(term, idf, d)
})
.collect()
}
pub fn name_cluster(
members: &[String],
tf_in_cluster: &HashMap<String, usize>,
clusters_containing: &HashMap<String, usize>,
n_clusters: usize,
cluster_size: usize,
commons: &HashSet<String>,
exclusivity: &HashMap<String, f64>,
) -> String {
let n = n_clusters.max(1) as f64;
let size = cluster_size.max(1) as f64;
const MIN_EXCLUSIVITY: f64 = 0.8;
let confined: Vec<String> = members
.iter()
.filter(|t| exclusivity.get(*t).copied().unwrap_or(1.0) >= MIN_EXCLUSIVITY)
.cloned()
.collect();
let pool: &[String] = if confined.is_empty() { members } else { &confined };
let mut scored: Vec<(String, f64)> = pool
.iter()
.map(|t| {
let tf = *tf_in_cluster.get(t).unwrap_or(&1) as f64;
let dfc = *clusters_containing.get(t).unwrap_or(&1) as f64;
let idf = ((n + 1.0) / (dfc + 1.0)).ln() + 1.0;
let coverage = (tf / size).min(1.0);
let common_bonus = if commons.contains(t) { 1.3 } else { 1.0 };
(t.clone(), coverage * common_bonus + 0.12 * idf)
})
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal).then(a.0.cmp(&b.0)));
scored.first().map(|(t, _)| t.clone()).unwrap_or_default()
}
fn cosine(a: &HashSet<usize>, b: &HashSet<usize>) -> f64 {
if a.is_empty() || b.is_empty() {
return 0.0;
}
let inter = a.intersection(b).count() as f64;
inter / ((a.len() as f64).sqrt() * (b.len() as f64).sqrt())
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TreeNode {
pub label: String,
pub height: f64,
pub size: usize,
pub children: Vec<TreeNode>,
}
pub fn discover_hierarchy(docs: &[String], n_terms: usize, n_clusters: usize) -> Option<TreeNode> {
let terms = salient_terms(docs, n_terms.clamp(2, 400));
if terms.len() < 2 {
return None;
}
let cut = n_clusters.clamp(1, terms.len());
let commons = common_nouns(docs);
let mut active: Vec<(Vec<usize>, HashSet<usize>, TreeNode)> = terms
.iter()
.enumerate()
.map(|(i, (t, d))| {
(vec![i], d.clone(), TreeNode { label: t.clone(), height: 0.0, size: 1, children: Vec::new() })
})
.collect();
let avg_linkage = |a: &[usize], b: &[usize]| -> f64 {
let mut sum = 0.0;
for x in a {
for y in b {
sum += cosine(&terms[*x].1, &terms[*y].1);
}
}
sum / (a.len() * b.len()) as f64
};
let mut labelled = false;
while active.len() > 1 {
if active.len() == cut && !labelled {
labelled = true;
let n_c = active.len();
let mut tf: HashMap<String, usize> = HashMap::new();
let mut containing: HashMap<String, usize> = HashMap::new();
for (members, docs_in, _) in &active {
let mut seen = HashSet::new();
for m in members {
let t = &terms[*m].0;
let c = docs_in.iter().filter_map(|i| docs.get(*i)).filter(|d| !word_spans(d, t).is_empty()).count();
*tf.entry(t.clone()).or_default() += c.max(1);
seen.insert(t.clone());
}
for t in seen {
*containing.entry(t).or_default() += 1;
}
}
let excl: Vec<HashMap<String, f64>> = (0..active.len())
.map(|ci| {
let others: HashSet<usize> = active
.iter()
.enumerate()
.filter(|(cj, _)| *cj != ci)
.flat_map(|(_, (_, docs_j, _))| docs_j.iter().copied())
.collect();
active[ci]
.0
.iter()
.map(|m| {
let (term, term_docs) = &terms[*m];
let own = term_docs.iter().filter(|d| !others.contains(d)).count() as f64;
(term.clone(), own / term_docs.len().max(1) as f64)
})
.collect()
})
.collect();
for (ci, (members, docs_in, node)) in active.iter_mut().enumerate() {
let names: Vec<String> = members.iter().map(|m| terms[*m].0.clone()).collect();
node.label =
name_cluster(&names, &tf, &containing, n_c, docs_in.len(), &commons, &excl[ci]);
}
}
let mut best: Option<(usize, usize, f64)> = None;
for i in 0..active.len() {
for j in (i + 1)..active.len() {
let s = avg_linkage(&active[i].0, &active[j].0);
if best.map(|(_, _, bs)| s > bs).unwrap_or(true) {
best = Some((i, j, s));
}
}
}
let Some((i, j, sim)) = best else { break };
let (mj, dj, nj) = active.remove(j);
let (_, _, ni) = &active[i];
let merged = TreeNode {
label: String::new(),
height: sim,
size: ni.size + nj.size,
children: vec![active[i].2.clone(), nj],
};
active[i].0.extend(mj);
active[i].1.extend(dj);
active[i].2 = merged;
}
active.into_iter().next().map(|(_, _, n)| n)
}
pub fn discover(docs: &[String], n_terms: usize, n_clusters: usize) -> Vec<TermCluster> {
let terms = salient_terms(docs, n_terms.clamp(2, 400));
if terms.len() < 2 || n_clusters == 0 {
return Vec::new();
}
let n_docs = docs.len().max(1);
let vecs: Vec<Vec<f32>> = terms
.iter()
.map(|(_, docs_in)| {
let mut v = vec![0f32; n_docs];
for i in docs_in {
if *i < n_docs {
v[*i] = 1.0;
}
}
let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt() + 1e-9;
v.into_iter().map(|x| x / norm).collect()
})
.collect();
let (_protos, assign, _cost) = crate::text::ot::codebook(&vecs, n_clusters, DISCOVER_EPS);
let k = assign.iter().copied().max().map_or(0, |m| m + 1);
let mut clusters: Vec<(Vec<usize>, HashSet<usize>)> = vec![(Vec::new(), HashSet::new()); k];
for (i, &a) in assign.iter().enumerate() {
clusters[a].0.push(i);
clusters[a].1.extend(terms[i].1.iter().copied());
}
clusters.retain(|(members, _)| !members.is_empty());
let commons = common_nouns(docs);
let n = docs.len().max(1) as f64;
let n_clusters_final = clusters.len();
let mut tf_in_cluster: HashMap<String, usize> = HashMap::new();
let mut clusters_containing: HashMap<String, usize> = HashMap::new();
for (members, docs_in) in &clusters {
let mut seen: HashSet<String> = HashSet::new();
for m in members {
let term = &terms[*m].0;
let count = docs_in
.iter()
.filter_map(|i| docs.get(*i))
.filter(|d| word_spans(d, term).len() > 0)
.count();
*tf_in_cluster.entry(term.clone()).or_default() += count.max(1);
seen.insert(term.clone());
}
for t in seen {
*clusters_containing.entry(t).or_default() += 1;
}
}
let exclusivity: Vec<HashMap<String, f64>> = (0..clusters.len())
.map(|ci| {
let others: HashSet<usize> = clusters
.iter()
.enumerate()
.filter(|(cj, _)| *cj != ci)
.flat_map(|(_, (_, docs_j))| docs_j.iter().copied())
.collect();
clusters[ci]
.0
.iter()
.map(|m| {
let (term, term_docs) = &terms[*m];
let own = term_docs.iter().filter(|d| !others.contains(d)).count() as f64;
(term.clone(), own / term_docs.len().max(1) as f64)
})
.collect()
})
.collect();
let mut out: Vec<TermCluster> = clusters
.iter()
.cloned()
.enumerate()
.map(|(ci, (members, docs_in))| {
let mut member_terms: Vec<String> = members.iter().map(|m| terms[*m].0.clone()).collect();
let mut sum = 0.0;
let mut pairs = 0usize;
for a in 0..members.len() {
for b in (a + 1)..members.len() {
sum += cosine(&terms[members[a]].1, &terms[members[b]].1);
pairs += 1;
}
}
let cohesion = if pairs == 0 { 1.0 } else { sum / pairs as f64 };
let label = name_cluster(
&member_terms,
&tf_in_cluster,
&clusters_containing,
n_clusters_final,
docs_in.len(),
&commons,
&exclusivity[ci],
);
member_terms.retain(|t| *t != label);
member_terms.insert(0, label.clone());
TermCluster { label, terms: member_terms, coverage: docs_in.len() as f64 / n, cohesion }
})
.filter(|c| !c.label.is_empty() && c.terms.len() > 1)
.collect();
out.sort_by(|a, b| b.coverage.partial_cmp(&a.coverage).unwrap_or(std::cmp::Ordering::Equal));
out
}
#[cfg(test)]
mod tests {
use super::*;
fn corpus() -> Vec<String> {
let battles = [
"Morty Shade defeated Wallace Gale in a battle at the tournament venue",
"Bea Strike defeated Falkner Gale in a battle at the tournament venue",
"Iris Draco defeated Nessa Reef in a battle at the tournament venue",
"Marnie Dusk defeated Juan Tide in a battle at the tournament venue",
];
let surveys = [
"The survey recorded elevation and rainfall across the habitat region",
"The survey recorded elevation and temperature across the habitat region",
"A survey measured rainfall and elevation within the habitat region",
"A survey measured temperature and elevation within the habitat region",
];
battles.iter().chain(surveys.iter()).map(|s| s.to_string()).collect()
}
#[test]
fn discovers_the_two_domains_without_any_field_markers() {
let docs = corpus();
let clusters = discover(&docs, 60, 2);
assert_eq!(clusters.len(), 2, "{clusters:#?}");
let joined: Vec<String> = clusters.iter().map(|c| c.terms.join(" ")).collect();
let battle_cluster = joined.iter().find(|t| t.contains("battle")).expect(&format!("{joined:?}"));
let survey_cluster = joined.iter().find(|t| t.contains("survey")).expect(&format!("{joined:?}"));
assert!(!battle_cluster.contains("survey"), "battle cluster leaked survey terms: {battle_cluster}");
assert!(!survey_cluster.contains("battle"), "survey cluster leaked battle terms: {survey_cluster}");
assert!(survey_cluster.contains("elevation"), "{survey_cluster}");
}
#[test]
fn labels_are_drawn_from_the_cluster_and_are_distinctive() {
let docs = corpus();
for c in discover(&docs, 60, 2) {
assert!(c.terms.contains(&c.label), "label must be a member: {c:?}");
assert_eq!(c.terms[0], c.label, "label should lead the member list");
assert!(!STOP.contains(&c.label.as_str()), "label is a stopword: {c:?}");
assert!(c.coverage > 0.0 && c.coverage <= 1.0, "{c:?}");
}
}
#[test]
fn prose_with_no_shared_vocabulary_yields_nothing_rather_than_noise() {
let docs: Vec<String> = ["alpha beta", "gamma delta", "epsilon zeta"].iter().map(|s| s.to_string()).collect();
assert!(discover(&docs, 40, 3).is_empty());
}
#[test]
fn is_deterministic() {
let docs = corpus();
let a = discover(&docs, 60, 3);
let b = discover(&docs, 60, 3);
assert_eq!(
a.iter().map(|c| c.terms.join(",")).collect::<Vec<_>>(),
b.iter().map(|c| c.terms.join(",")).collect::<Vec<_>>()
);
}
#[test]
fn quantities_are_extracted_with_their_units() {
let doc = "recorded at an elevation of 1082 m. The mean temperature was 28 °C and it ran 7 minutes.";
let q = quantity_spans(doc);
let got: Vec<(&str, &str)> = q.iter().map(|(s, e, f)| (&doc[*s..*e], f.as_str())).collect();
assert!(got.contains(&("1082 m", "length_m")), "{got:?}");
assert!(got.contains(&("28 °C", "temp_c")), "{got:?}");
assert!(got.contains(&("7 minutes", "minutes")), "{got:?}");
}
#[test]
fn km_is_not_read_as_m() {
let doc = "a range of 500 km across";
let q = quantity_spans(doc);
assert_eq!(q.len(), 1, "{q:?}");
assert_eq!(q[0].2, "length_km");
assert_eq!(&doc[q[0].0..q[0].1], "500 km");
}
#[test]
fn gazetteer_keeps_whole_mentions_not_their_parts() {
let docs: Vec<String> = [
"A survey in Sootopolis City recorded Aggron near the crater",
"Another survey in Sootopolis City found more Aggron there",
"The Indigo Invitational was held in Sootopolis City again",
].iter().map(|s| s.to_string()).collect();
let g = mine_gazetteer(&docs, 2);
assert!(g.contains(&"Sootopolis City".to_string()), "{g:?}");
assert!(!g.contains(&"Sootopolis".to_string()), "{g:?}");
assert!(!g.contains(&"City".to_string()), "{g:?}");
assert!(g.iter().all(|x| x.split_whitespace().count() >= 2), "{g:?}");
}
#[test]
fn sentence_initial_capitals_do_not_become_entities() {
let docs: Vec<String> = [
"The survey found nothing. The survey ended early",
"The survey found nothing. The survey ended early",
].iter().map(|s| s.to_string()).collect();
let g = mine_gazetteer(&docs, 2);
assert!(!g.iter().any(|x| x.starts_with("The ")), "{g:?}");
}
#[test]
fn generic_words_are_not_selected_as_vocabulary() {
let docs: Vec<String> = (0..4)
.map(|i| format!("survey {i} recorded data within the location and the distribution of species"))
.collect();
let picked: Vec<String> = salient(&docs, 40).into_iter().map(|(t, _, _)| t).collect();
for g in ["data", "within", "location", "distribution"] {
assert!(!picked.contains(&g.to_string()), "generic term leaked: {g} in {picked:?}");
}
assert!(picked.contains(&"survey".to_string()) || picked.contains(&"species".to_string()), "{picked:?}");
}
#[test]
fn temporal_buckets_are_extracted_and_qualified() {
let doc = "held in Q3 2026 at the venue, following the 2025 season";
let t = temporal_spans(doc);
let toks: Vec<&str> = t.iter().map(|(_, _, x)| x.as_str()).collect();
assert!(toks.contains(&"time/2026/q3"), "{toks:?}");
assert!(toks.contains(&"time/2025"), "{toks:?}");
for (s, e, _) in &t {
assert!(doc.get(*s..*e).is_some(), "bad span {s}..{e}");
}
}
#[test]
fn a_four_digit_number_that_is_not_a_year_is_ignored() {
let t = temporal_spans("an elevation of 2369 m");
assert!(t.iter().all(|(_, _, x)| x != "time/2369"), "{t:?}");
}
#[test]
fn relations_carry_direction_from_word_order() {
let mentions: Vec<String> = ["Morty Shade", "Wallace Gale"].iter().map(|s| s.to_string()).collect();
let r = relation_spans("Morty Shade defeated Wallace Gale at the venue", &mentions);
assert_eq!(r.len(), 1, "{r:?}");
assert_eq!(r[0].verb, "defeated");
assert_eq!(r[0].actor, "Morty Shade");
assert_eq!(r[0].target, "Wallace Gale");
let rev = relation_spans("Wallace Gale defeated Morty Shade at the venue", &mentions);
assert_eq!(rev[0].actor, "Wallace Gale");
assert_eq!(rev[0].target, "Morty Shade");
}
#[test]
fn no_relation_is_invented_across_a_sentence_boundary() {
let mentions: Vec<String> = ["Morty Shade", "Wallace Gale"].iter().map(|s| s.to_string()).collect();
let r = relation_spans("Morty Shade defeated someone. Wallace Gale watched", &mentions);
assert!(r.is_empty(), "should not link across sentences: {r:?}");
}
#[test]
fn a_single_mention_yields_no_relation() {
let mentions: Vec<String> = vec!["Morty Shade".to_string()];
assert!(relation_spans("Morty Shade defeated everyone", &mentions).is_empty());
}
#[test]
fn a_one_off_name_can_still_be_a_relation_participant() {
let r = relation_spans("At the tournament, Juan Tide defeated Cynthia Ward in a close battle", &[]);
assert!(!r.is_empty(), "should read the relation from local names: {r:?}");
let d = r.iter().find(|x| x.verb == "defeated").expect("a defeated relation");
assert_eq!(d.actor, "Juan Tide");
assert_eq!(d.target, "Cynthia Ward");
}
#[test]
fn local_mentions_skip_the_sentence_initial_capital() {
let m = local_mentions("Juan Tide defeated Cynthia Ward");
let names: Vec<&str> = m.iter().map(|(_, _, n)| n.as_str()).collect();
assert!(names.iter().any(|n| n.contains("Cynthia Ward")), "{names:?}");
assert!(!names.iter().any(|n| n.starts_with("Juan Tide defeated")), "{names:?}");
}
#[test]
fn local_mentions_survive_multibyte_text() {
for text in [
"A Pokémon named Aggron was recorded at 28 °C by Cynthia Ward",
"Café Ecruteak hosted Juan Tide and Bea Strike",
"28 °C — Sootopolis City",
] {
let m = local_mentions(text);
for (s, e, name) in &m {
assert_eq!(&text[*s..*e], name, "offsets must slice cleanly");
}
}
}
#[test]
fn relations_survive_multibyte_text() {
let r = relation_spans("At the venue, Juan Tide defeated Cynthia Ward and a Pokémon at 28 °C", &[]);
assert!(r.iter().any(|x| x.actor == "Juan Tide"), "{r:?}");
}
#[test]
fn punctuation_breaks_a_capitalised_run() {
let m = local_mentions("held in Violet City, Johto, Juan Tide defeated Cynthia Ward");
let names: Vec<&str> = m.iter().map(|(_, _, n)| n.as_str()).collect();
assert!(names.contains(&"Violet City"), "{names:?}");
assert!(names.contains(&"Juan Tide"), "{names:?}");
assert!(!names.iter().any(|n| n.contains(',')), "a name must not span punctuation: {names:?}");
let r = relation_spans("held in Violet City, Johto, Juan Tide defeated Cynthia Ward", &[]);
let d = r.iter().find(|x| x.verb == "defeated").expect("a defeated relation");
assert_eq!(d.actor, "Juan Tide", "nearest clean name, not a comma-joined run");
assert_eq!(d.target, "Cynthia Ward");
}
#[test]
fn a_negative_quantity_keeps_its_sign() {
let doc = "the mean temperature was -7 °C that winter";
let q = quantity_spans(doc);
let (s, e, f) = q.first().expect("a quantity").clone();
assert_eq!(f, "temp_c");
assert_eq!(&doc[s..e], "-7 °C", "the sign is part of the number");
}
#[test]
fn a_hyphen_between_words_is_not_a_minus_sign() {
let doc = "an 11-minute battle";
let q = quantity_spans(doc);
let (s, e, _) = q.first().expect("a quantity").clone();
assert_eq!(&doc[s..e], "11-minute".split('-').next().unwrap().to_owned() + "-minute");
assert!(!doc[s..e].starts_with('-'), "must not read the compound hyphen as a sign: {:?}", &doc[s..e]);
}
#[test]
fn a_range_hyphen_is_not_a_minus_sign() {
let doc = "between 5-10 m of clearance";
for (s, e, _) in quantity_spans(doc) {
assert!(!doc[s..e].starts_with('-'), "range hyphen read as a sign: {:?}", &doc[s..e]);
}
}
#[test]
fn a_term_spanning_two_domains_cannot_name_either() {
let docs: Vec<String> = [
"Morty Shade defeated Wallace Gale in a battle at Ecruteak City",
"Bea Strike defeated Falkner Gale in a battle at Ecruteak City",
"Iris Draco defeated Nessa Reef in a battle at Ecruteak City",
"The survey recorded elevation across the habitat near Sootopolis City",
"The survey recorded rainfall across the habitat near Sootopolis City",
"A survey measured elevation within the habitat near Sootopolis City",
]
.iter()
.map(|s| s.to_string())
.collect();
let clusters = discover(&docs, 60, 2);
assert!(!clusters.is_empty(), "the two domains should still be found");
for c in &clusters {
assert_ne!(c.label, "city", "a term common to both domains named one of them: {c:?}");
}
assert!(clusters.iter().any(|c| c.terms.iter().any(|t| t == "city")), "{clusters:#?}");
}
#[test]
fn discovery_transports_every_salient_term_to_some_facet() {
let docs = corpus();
let clusters = discover(&docs, 60, 2);
let placed: usize = clusters.iter().map(|c| c.terms.len()).sum();
assert!(placed >= 6, "expected the salient terms to be placed, got {placed}: {clusters:#?}");
}
#[test]
fn a_function_word_cannot_become_a_category() {
let docs: Vec<String> = [
"Morty Shade defeated Wallace Gale at Ecruteak City in 2025.",
"Bea Strike defeated Iris Draco at Ecruteak City in 2025.",
"Lance Wing defeated Karen Dusk at Ecruteak City in 2025.",
"A survey recorded Aggron near Sootopolis City at 28 degrees.",
"A survey recorded Salamence near Sootopolis City at 31 degrees.",
"A survey recorded Metagross near Sootopolis City at 19 degrees.",
"Milotic is not permitted in Series 1 play for the 2025 season.",
"Registeel is not permitted in Series 1 play for the 2025 season.",
]
.iter()
.map(|s| s.to_string())
.collect();
for c in discover(&docs, 60, 4) {
assert!(
!LOCATIVES.contains(&c.label.as_str()) && !STOP.contains(&c.label.as_str()),
"a function word named a category: {c:?}"
);
assert!(
!c.terms.iter().any(|t| LOCATIVES.contains(&t.as_str())),
"a function word was clustered as a signal word: {c:?}"
);
}
}
#[test]
fn the_word_lists_do_not_overlap() {
for w in LOCATIVES {
assert!(!STOP.contains(w), "{w} is in both STOP and LOCATIVES");
assert!(!GENERIC.contains(w), "{w} is in both GENERIC and LOCATIVES");
}
}
}