cleora 1.2.3

Cleora is a general-purpose model for efficient, scalable learning of stable and inductive entity embeddings for heterogeneous relational data.
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

use cleora::configuration::{Column, Configuration, FileType, OutputFormat};
use cleora::embedding::{calculate_embeddings, calculate_embeddings_mmap};
use cleora::persistence::embedding::EmbeddingPersistor;
use cleora::persistence::entity::InMemoryEntityMappingPersistor;
use cleora::pipeline::build_graphs;
use insta::assert_debug_snapshot;
use std::io;
use std::sync::Arc;

/// This test performs work for sample case and saves snapshot file.
/// Snapshot testing takes advantage of deterministic character of Cleora.
/// Any discrepancies between original snapshot results and current ones can be then
/// reviewed along with the code which introduced discrepancy.
///
/// Differing snapshot has to be renamed by removing .new from the name.
/// For more information, please review https://crates.io/crates/insta
///
/// Code executed performs roughly the same work as:
/// ./cleora -i files/samples/edgelist_1.tsv --columns="complex::reflexive::a b complex::c"
/// -d 128 -n 4 --relation-name=R1 -p 0
#[test]
fn test_build_graphs_and_create_embeddings() {
    let config = prepare_config();

    let in_memory_entity_mapping_persistor = InMemoryEntityMappingPersistor::default();
    let in_memory_entity_mapping_persistor = Arc::new(in_memory_entity_mapping_persistor);

    // build sparse matrices
    let sparse_matrices = build_graphs(&config, in_memory_entity_mapping_persistor.clone());

    let config = Arc::new(config);

    // embeddings for in-memory and mmap files calculation should be the same
    for sparse_matrix in sparse_matrices.into_iter() {
        let sparse_matrix = Arc::new(sparse_matrix);
        let snapshot_name = format!(
            "embeddings_{}_{}",
            sparse_matrix.col_a_name, sparse_matrix.col_b_name
        );

        let mut in_memory_embedding_persistor = InMemoryEmbeddingPersistor::default();
        // calculate embeddings in memory
        calculate_embeddings(
            config.clone(),
            sparse_matrix.clone(),
            in_memory_entity_mapping_persistor.clone(),
            &mut in_memory_embedding_persistor,
        );
        assert_debug_snapshot!(snapshot_name.clone(), in_memory_embedding_persistor);

        let mut in_memory_embedding_persistor = InMemoryEmbeddingPersistor::default();
        // calculate embeddings with mmap files
        calculate_embeddings_mmap(
            config.clone(),
            sparse_matrix.clone(),
            in_memory_entity_mapping_persistor.clone(),
            &mut in_memory_embedding_persistor,
        );
        assert_debug_snapshot!(snapshot_name, in_memory_embedding_persistor);
    }
}

fn prepare_config() -> Configuration {
    let columns = vec![
        Column {
            name: "a".to_string(),
            complex: true,
            reflexive: true,
            ..Column::default()
        },
        Column {
            name: "b".to_string(),
            ..Column::default()
        },
        Column {
            name: "c".to_string(),
            complex: true,
            ..Column::default()
        },
    ];

    let config = Configuration {
        produce_entity_occurrence_count: true,
        embeddings_dimension: 128,
        max_number_of_iteration: 4,
        seed: None,
        prepend_field: false,
        log_every_n: 10000,
        in_memory_embedding_calculation: true,
        input: vec!["files/samples/edgelist_1.tsv".to_string()],
        file_type: FileType::Tsv,
        output_format: OutputFormat::TextFile,
        output_dir: None,
        relation_name: "r1".to_string(),
        columns,
    };
    config
}

#[derive(Debug, Default)]
struct InMemoryEmbeddingPersistor {
    entity_count: u32,
    dimenstion: u16,
    entities: Vec<InMemoryEntity>,
}

#[derive(Debug)]
struct InMemoryEntity {
    entity: String,
    occur_count: u32,
    vector: Vec<f32>,
}

impl EmbeddingPersistor for InMemoryEmbeddingPersistor {
    fn put_metadata(&mut self, entity_count: u32, dimension: u16) -> Result<(), io::Error> {
        self.entity_count = entity_count;
        self.dimenstion = dimension;
        Ok(())
    }
    fn put_data(
        &mut self,
        entity: &str,
        occur_count: u32,
        vector: Vec<f32>,
    ) -> Result<(), io::Error> {
        let entity = entity.to_string();
        self.entities.push(InMemoryEntity {
            entity,
            occur_count,
            vector,
        });
        Ok(())
    }
    fn finish(&mut self) -> Result<(), io::Error> {
        Ok(())
    }
}