def test_ermlp_predict(self):
"""Test ERMLP's predict function."""
ermlp = ERMLP(**ERMLP_CONFIG)
ermlp.num_entities = ERMLP_CONFIG[NUM_ENTITIES]
ermlp.num_relations = ERMLP_CONFIG[NUM_RELATIONS]
predictions = ermlp.predict(triples=TEST_TRIPLES)
self.assertEqual(len(predictions), len(TEST_TRIPLES))
self.assertTrue(type(predictions.shape[0]), float)
def test_instantiate_ermlp(self):
"""Test that ERMLP can be instantiated."""
ermlp = ERMLP(**ERMLP_CONFIG)
ermlp.num_entities = ERMLP_CONFIG[NUM_ENTITIES]
ermlp.num_relations = ERMLP_CONFIG[NUM_RELATIONS]
self.assertIsNotNone(ermlp)
self.assertEqual(ermlp.num_entities, 5)
self.assertEqual(ermlp.num_relations, 5)
self.assertEqual(ermlp.embedding_dim, 2)
self.assertEqual(ermlp.margin_loss, 4)