lattix 0.8.1

Knowledge graph data structures
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

//! Passage retrieval with Personalized PageRank over a toy knowledge graph.
//!
//! Run:
//! `cargo run --example ppr_retrieval`

use std::collections::{HashMap, HashSet};

use lattix::algo::ppr::{personalized_pagerank, PprConfig};
use lattix::{EntityId, KnowledgeGraph, Triple};

#[derive(Debug)]
struct Passage {
    id: &'static str,
    text: &'static str,
    entities: &'static [&'static str],
}

fn main() {
    let passages = [
        Passage {
            id: "p_ada",
            text: "Ada Lovelace wrote notes about the Analytical Engine.",
            entities: &["Ada Lovelace", "Analytical Engine"],
        },
        Passage {
            id: "p_engine",
            text: "The Analytical Engine was designed by Charles Babbage.",
            entities: &["Analytical Engine", "Charles Babbage"],
        },
        Passage {
            id: "p_babbage",
            text: "Charles Babbage studied mechanical computation.",
            entities: &["Charles Babbage", "mechanical computation"],
        },
        Passage {
            id: "p_hopper",
            text: "Grace Hopper worked on COBOL compilers.",
            entities: &["Grace Hopper", "COBOL", "compilers"],
        },
    ];

    let mut kg = KnowledgeGraph::new();
    for (subject, predicate, object) in [
        ("Ada Lovelace", "wrote_notes_about", "Analytical Engine"),
        ("Analytical Engine", "designed_by", "Charles Babbage"),
        ("Charles Babbage", "studied", "mechanical computation"),
        ("Grace Hopper", "worked_on", "COBOL"),
        ("Grace Hopper", "worked_on", "compilers"),
    ] {
        kg.add_triple(Triple::new(subject, predicate, object));
    }

    let query = "who designed the machine Ada Lovelace wrote notes about?";
    let seed = "Ada Lovelace";

    let ppr_scores = personalized_pagerank(
        &kg,
        seed,
        PprConfig {
            damping: 0.85,
            max_iterations: 100,
            tolerance: 1e-9,
        },
    );

    let graph_ranked = rank_by_graph(&passages, &ppr_scores, &[seed]);
    let lexical_ranked = rank_by_lexical(&passages, query);

    println!("query: {query}");
    println!("seed:  {seed}");

    println!("\nlexical baseline:");
    for (passage, score) in lexical_ranked.iter().take(3) {
        println!("  {:<10} score={score:<2} {}", passage.id, passage.text);
    }

    println!("\ngraph retrieval:");
    for (passage, score) in graph_ranked.iter().take(3) {
        println!("  {:<10} score={score:.5} {}", passage.id, passage.text);
    }

    assert_eq!(lexical_ranked[0].0.id, "p_ada");
    assert_eq!(graph_ranked[0].0.id, "p_engine");
    assert!(
        ppr_scores
            .get("Charles Babbage")
            .copied()
            .unwrap_or_default()
            > ppr_scores.get("Grace Hopper").copied().unwrap_or_default()
    );
}

fn rank_by_graph<'a>(
    passages: &'a [Passage],
    scores: &HashMap<EntityId, f64>,
    excluded_entities: &[&str],
) -> Vec<(&'a Passage, f64)> {
    let excluded_entities = excluded_entities.iter().copied().collect::<HashSet<_>>();
    let mut ranked = passages
        .iter()
        .map(|passage| {
            let score = passage
                .entities
                .iter()
                .filter(|entity| !excluded_entities.contains(*entity))
                .map(|entity| scores.get(*entity).copied().unwrap_or_default())
                .sum();
            (passage, score)
        })
        .collect::<Vec<_>>();
    ranked.sort_by(compare_scored_passages);
    ranked
}

fn rank_by_lexical<'a>(passages: &'a [Passage], query: &str) -> Vec<(&'a Passage, usize)> {
    let query_terms = tokenize(query);
    let mut ranked = passages
        .iter()
        .map(|passage| {
            let score = tokenize(passage.text).intersection(&query_terms).count();
            (passage, score)
        })
        .collect::<Vec<_>>();
    ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.id.cmp(b.0.id)));
    ranked
}

fn tokenize(text: &str) -> HashSet<String> {
    text.split(|ch: char| !ch.is_ascii_alphanumeric())
        .filter(|term| term.len() > 2)
        .map(str::to_ascii_lowercase)
        .collect()
}

fn compare_scored_passages(left: &(&Passage, f64), right: &(&Passage, f64)) -> std::cmp::Ordering {
    right
        .1
        .partial_cmp(&left.1)
        .unwrap_or(std::cmp::Ordering::Equal)
        .then_with(|| left.0.id.cmp(right.0.id))
}