def test_construction_single(self):
entity_dump_dir = "test/data/preprocessing/base/entity_data/entity_mappings"
entity_symbols = EntitySymbols(load_dir=entity_dump_dir,
alias_cand_map_file="alias2qids.json")
fmt_types = {"trie1": "qid_cand_with_score"}
max_values = {"trie1": 3}
input_dicts = {"trie1": entity_symbols.get_alias2qids()}
record_trie = RecordTrieCollection(load_dir=None, input_dicts=input_dicts, vocabulary=entity_symbols.get_qid2eid(),
fmt_types=fmt_types, max_values=max_values)
truealias2qids = {
'alias1': [["Q1", 10], ["Q4", 6]],
'multi word alias2': [["Q2", 5], ["Q1", 3], ["Q4", 2]],
'alias3': [["Q1", 30]],
'alias4': [["Q4", 20], ["Q3", 15], ["Q2", 1]]
}
for al in truealias2qids:
self.assertEqual(truealias2qids[al], record_trie.get_value("trie1", al))
for al in truealias2qids:
self.assertEqual([x[0] for x in truealias2qids[al]], record_trie.get_value("trie1", al, getter=lambda x: x[0]))
for al in truealias2qids:
self.assertEqual([x[1] for x in truealias2qids[al]], record_trie.get_value("trie1", al, getter=lambda x: x[1]))
self.assertEqual(set(truealias2qids.keys()), set(record_trie.get_keys("trie1")))
def test_load_and_save(self):
entity_dump_dir = "test/data/preprocessing/base/entity_data/entity_mappings"
entity_symbols = EntitySymbols(load_dir=entity_dump_dir,
alias_cand_map_file="alias2qids.json")
fmt_types = {"trie1": "qid_cand_with_score"}
max_values = {"trie1": 3}
input_dicts = {"trie1": entity_symbols.get_alias2qids()}
record_trie = RecordTrieCollection(load_dir=None, input_dicts=input_dicts, vocabulary=entity_symbols.get_qid2eid(),
fmt_types=fmt_types, max_values=max_values)
record_trie.dump(save_dir=os.path.join(entity_dump_dir, "record_trie"))
record_trie_loaded = RecordTrieCollection(load_dir=os.path.join(entity_dump_dir, "record_trie"))
self.assertEqual(record_trie._fmt_types, record_trie_loaded._fmt_types)
self.assertEqual(record_trie._max_values, record_trie_loaded._max_values)
self.assertEqual(record_trie._stoi, record_trie_loaded._stoi)
np.testing.assert_array_equal(record_trie._itos, record_trie_loaded._itos)
self.assertEqual(record_trie._record_tris, record_trie_loaded._record_tris)
def test_construction_double(self):
truealias2qids = {
'alias1': [["Q1", 10], ["Q4", 6]],
'multi word alias2': [["Q2", 5], ["Q1", 3], ["Q4", 2]],
'alias3': [["Q1", 30]],
'alias4': [["Q4", 20], ["Q3", 15], ["Q2", 1]]
}
truealias2typeids = {
'Q1': [1, 2, 3],
'Q2': [8],
'Q3': [4],
'Q4': [2,4]
}
truerelations = {
'Q1': set(["Q2", "Q3"]),
'Q2': set(["Q1"]),
'Q3': set(["Q2"]),
'Q4': set(["Q1", "Q2", "Q3"])
}
entity_dump_dir = "test/data/preprocessing/base/entity_data/entity_mappings"
entity_symbols = EntitySymbols(load_dir=entity_dump_dir,
alias_cand_map_file="alias2qids.json")
fmt_types = {"trie1": "qid_cand_with_score", "trie2": "type_ids", "trie3": "kg_relations"}
max_values = {"trie1": 3, "trie2": 3, "trie3": 3}
input_dicts = {"trie1": entity_symbols.get_alias2qids(), "trie2": truealias2typeids, "trie3": truerelations}
record_trie = RecordTrieCollection(load_dir=None, input_dicts=input_dicts, vocabulary=entity_symbols.get_qid2eid(),
fmt_types=fmt_types, max_values=max_values)
for al in truealias2qids:
self.assertEqual(truealias2qids[al], record_trie.get_value("trie1", al))
for qid in truealias2typeids:
self.assertEqual(truealias2typeids[qid], record_trie.get_value("trie2", qid))
for qid in truerelations:
self.assertEqual(truerelations[qid], record_trie.get_value("trie3", qid))