narrative_graph/heuristic/
mod.rs

1mod cooccurrence;
2mod entities;
3mod relations;
4mod segment;
5
6use crate::types::{Options, TripleCandidate};
7use crate::Result;
8use std::collections::BTreeMap;
9
10pub use entities::extract_entities;
11pub use relations::normalize_relation;
12
13/// Extract candidate (subject, relation, object) triples from prose.
14/// Runs the heuristic pipeline: entity detection, relation labeling, co-occurrence scoring, and confidence calculation.
15pub fn extract_candidate_triples(text: &str, opts: &Options) -> Result<Vec<TripleCandidate>> {
16    if text.is_empty() {
17        return Ok(vec![]);
18    }
19
20    let min_confidence = opts.min_confidence.unwrap_or(0.0);
21    if !(0.0..=1.0).contains(&min_confidence) {
22        return Err(crate::NarrativeGraphError::InvalidConfidenceThreshold(
23            min_confidence,
24        ));
25    }
26
27    let mut candidates = Vec::new();
28
29    for (sent_text, sent_start) in segment::split_sentences(text) {
30        let entities = extract_entities(sent_text, &opts.aliases);
31
32        if entities.is_empty() {
33            continue;
34        }
35
36        // Look for relations between entity pairs in the same clause
37        let relations = relations::extract_relations(sent_text, &entities, &opts.ontology);
38
39        // Score by rule strength and entity proximity within the sentence
40        for rel in relations {
41            let confidence = cooccurrence::score_confidence(rel.base, rel.gap);
42
43            if confidence >= min_confidence {
44                let span = [sent_start + rel.span[0], sent_start + rel.span[1]];
45
46                candidates.push(TripleCandidate {
47                    subject: rel.subject,
48                    relation: rel.relation,
49                    object: rel.object,
50                    confidence,
51                    span,
52                    rule: rel.rule,
53                });
54            }
55        }
56    }
57
58    // Deduplicate: keep highest confidence for each (subj, rel, obj) triple
59    dedup_candidates(&mut candidates);
60
61    Ok(candidates)
62}
63
64fn dedup_candidates(candidates: &mut Vec<TripleCandidate>) {
65    let mut best: BTreeMap<(String, String, String), TripleCandidate> = BTreeMap::new();
66
67    for candidate in candidates.drain(..) {
68        let key = (
69            candidate.subject.clone(),
70            candidate.relation.clone(),
71            candidate.object.clone(),
72        );
73        best.entry(key)
74            .and_modify(|best_cand| {
75                if candidate.confidence > best_cand.confidence {
76                    *best_cand = candidate.clone();
77                }
78            })
79            .or_insert(candidate);
80    }
81
82    *candidates = best.into_values().collect();
83}
narrative_graph/heuristic/mod.rs

narrative_graph/heuristic/
mod.rs
