mod cooccurrence;
mod entities;
mod relations;
mod segment;
use crate::types::{Options, TripleCandidate};
use crate::Result;
use std::collections::BTreeMap;
pub use entities::extract_entities;
pub use relations::normalize_relation;
pub fn extract_candidate_triples(text: &str, opts: &Options) -> Result<Vec<TripleCandidate>> {
if text.is_empty() {
return Ok(vec![]);
}
let min_confidence = opts.min_confidence.unwrap_or(0.0);
if !(0.0..=1.0).contains(&min_confidence) {
return Err(crate::NarrativeGraphError::InvalidConfidenceThreshold(
min_confidence,
));
}
let mut candidates = Vec::new();
for (sent_text, sent_start) in segment::split_sentences(text) {
let entities = extract_entities(sent_text, &opts.aliases);
if entities.is_empty() {
continue;
}
let relations = relations::extract_relations(sent_text, &entities, &opts.ontology);
for rel in relations {
let confidence = cooccurrence::score_confidence(rel.base, rel.gap);
if confidence >= min_confidence {
let span = [sent_start + rel.span[0], sent_start + rel.span[1]];
candidates.push(TripleCandidate {
subject: rel.subject,
relation: rel.relation,
object: rel.object,
confidence,
span,
rule: rel.rule,
});
}
}
}
dedup_candidates(&mut candidates);
Ok(candidates)
}
fn dedup_candidates(candidates: &mut Vec<TripleCandidate>) {
let mut best: BTreeMap<(String, String, String), TripleCandidate> = BTreeMap::new();
for candidate in candidates.drain(..) {
let key = (
candidate.subject.clone(),
candidate.relation.clone(),
candidate.object.clone(),
);
best.entry(key)
.and_modify(|best_cand| {
if candidate.confidence > best_cand.confidence {
*best_cand = candidate.clone();
}
})
.or_insert(candidate);
}
*candidates = best.into_values().collect();
}