From fc11f052c4dd4256ef898a7922d5c57cc4734a2b Mon Sep 17 00:00:00 2001 From: Stanislaw Adaszewski Date: Mon, 15 Jun 2020 11:42:54 +0200 Subject: [PATCH] Fix test_decode_layer_03. --- tests/icosagon/test_declayer.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tests/icosagon/test_declayer.py b/tests/icosagon/test_declayer.py index 17b74f9..26eb3cb 100644 --- a/tests/icosagon/test_declayer.py +++ b/tests/icosagon/test_declayer.py @@ -79,23 +79,24 @@ def test_decode_layer_03(): d = Data() d.add_node_type('Dummy 1', 100) d.add_node_type('Dummy 2', 100) - d.add_relation_type('Dummy Relation 1', 0, 1, + fam = d.add_relation_family('Dummy 1-Dummy 2', 0, 1, True) + fam.add_relation_type('Dummy Relation 1', torch.rand((100, 100), dtype=torch.float32).round().to_sparse()) prep_d = prepare_training(d, TrainValTest(.8, .1, .1)) in_layer = OneHotInputLayer(d) d_layer = DecagonLayer(in_layer.output_dim, 32, d) - dec_layer = DecodeLayer(input_dim=d_layer.output_dim, data=prep_d, keep_prob=1., - decoder_class={(0, 1): DEDICOMDecoder}, activation=lambda x: x) + dec_layer = DecodeLayer(input_dim=d_layer.output_dim, data=prep_d, + keep_prob=1., activation=lambda x: x) seq = torch.nn.Sequential(in_layer, d_layer, dec_layer) - pred_adj_matrices = seq(None) - assert isinstance(pred_adj_matrices, dict) - assert len(pred_adj_matrices) == 2 - assert isinstance(pred_adj_matrices[0, 1], list) - assert isinstance(pred_adj_matrices[1, 0], list) - assert len(pred_adj_matrices[0, 1]) == 1 - assert len(pred_adj_matrices[1, 0]) == 1 + pred = seq(None) + assert isinstance(pred, Predictions) + assert len(pred.relation_families) == 1 + assert isinstance(pred.relation_families[0], RelationFamilyPredictions) + assert isinstance(pred.relation_families[0].relation_types, list) + assert len(pred.relation_families[0].relation_types) == 1 + assert isinstance(pred.relation_families[0].relation_types[0], RelationPredictions) def test_decode_layer_04():