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
//! Adapter for constructing [`Dataset`] from [`petgraph`] graphs.
//!
//! Converts a `petgraph::Graph` (or `DiGraph`) where nodes are entity names
//! and edges are relation names into subsume's [`Triple`] / [`Dataset`] format.
use crate::dataset::{Dataset, Triple};
/// Convert a directed petgraph into a [`Dataset`] with all triples in the training split.
///
/// Node weights become entity names (via `ToString`), edge weights become
/// relation names. All edges become training triples; validation and test
/// splits are empty (callers can split manually).
///
/// # Example
///
/// ```rust,ignore
/// use petgraph::graph::DiGraph;
/// use subsume::petgraph_adapter::from_graph;
///
/// let mut g = DiGraph::new();
/// let dog = g.add_node("dog");
/// let animal = g.add_node("animal");
/// g.add_edge(dog, animal, "is_a");
///
/// let dataset = from_graph(&g);
/// assert_eq!(dataset.train.len(), 1);
/// ```
pub fn from_graph<N, E>(graph: &petgraph::graph::DiGraph<N, E>) -> Dataset
where
N: ToString,
E: ToString,
{
let triples: Vec<Triple> = graph
.edge_indices()
.filter_map(|e| {
let (src, dst) = graph.edge_endpoints(e)?;
let rel = graph.edge_weight(e)?;
Some(Triple::new(
graph[src].to_string(),
rel.to_string(),
graph[dst].to_string(),
))
})
.collect();
Dataset::new(triples, Vec::new(), Vec::new())
}