def test_trans_e_predict(self):
"""Test TransE's predict function."""
trans_e = TransE(**TRANS_E_CONFIG)
predictions = trans_e.predict(triples=TEST_TRIPLES)
self.assertEqual(len(predictions), len(TEST_TRIPLES))
self.assertTrue(type(predictions.shape[0]), float)
def test_trans_h_predict(self):
"""Test TransH's predict function."""
trans_h = TransE(config=TRANS_H_CONFIG)
predictions = trans_h.predict(triples=TEST_TRIPLES)
self.assertEqual(len(predictions), len(TEST_TRIPLES))
self.assertTrue(type(predictions.shape[0]), float)
def test_trans_h_predict(self):
"""Test TransH's predict function."""
trans_h = TransE(**TRANS_H_CONFIG)
trans_h.num_entities = TRANS_H_CONFIG[NUM_ENTITIES]
trans_h.num_relations = TRANS_H_CONFIG[NUM_RELATIONS]
predictions = trans_h.predict(triples=TEST_TRIPLES)
self.assertEqual(len(predictions), len(TEST_TRIPLES))
self.assertTrue(type(predictions.shape[0]), float)
def test_compute_scores_trans_e(self):
"""Test that TransE's socore function computes the scores correct."""
trans_e = TransE(config=TRANS_E_CONFIG)
h_embs = torch.tensor([[1., 1.], [2., 2.]], dtype=torch.float)
r_embs = torch.tensor([[1., 1.], [2., 2.]], dtype=torch.float)
t_embs = torch.tensor([[2., 2.], [4., 4.]], dtype=torch.float)
scores = trans_e._compute_scores(h_embs, r_embs,
t_embs).cpu().numpy().tolist()
self.assertEqual(scores, [0., 0.])
def test_instantiate_trans_e(self):
"""Test that TransE can be instantiated."""
trans_e = TransE(**TRANS_E_CONFIG)
trans_e.num_entities = TRANS_E_CONFIG[NUM_ENTITIES]
trans_e.num_relations = TRANS_E_CONFIG[NUM_RELATIONS]
self.assertIsNotNone(trans_e)
self.assertEqual(trans_e.num_entities, 5)
self.assertEqual(trans_e.num_relations, 5)
self.assertEqual(trans_e.embedding_dim, 5)
self.assertEqual(trans_e.l_p_norm_entities, 2)
self.assertEqual(trans_e.scoring_fct_norm, 1)
self.assertEqual(trans_e.margin_loss, 4)
class TranseUnit(EnsembleUnit):
def __init__(self,
relation_embeddings,
num_entities,
num_relations,
entity_embeddings,
preferred_device='cpu'):
super(TranseUnit, self).__init__()
if preferred_device == 'cuda':
preferred_device = 'gpu'
self.model = TransE(preferred_device=preferred_device,
random_seed=1234,
embedding_dim=entity_embeddings.shape[1],
margin_loss=1,
normalization_of_entities=1)
self.model.num_entities = num_entities
self.model.num_relations = num_relations
self.model.relation_embeddings = nn.Embedding.from_pretrained(
relation_embeddings) # noqa
self.model.entity_embeddings = nn.Embedding.from_pretrained(
entity_embeddings) # noqa
def forward(self, x: torch.tensor) -> torch.tensor:
"""
Compute the scores for given s, p and every o.
:param x: torch.tensor, dtype: torch.long, shape: (batch_size, 2)
The s, p indices.
:return:
scores: torch.tensor, dtype: torch.float, shape: (batch_size, num_entities) # noqa
The scores for every possible entity as o.
"""
batch_size = x.shape[0]
head_embeddings = self.model.entity_embeddings(x[:, 0:1])
head_embeddings = head_embeddings.view(batch_size, 1, -1)
relation_embeddings = self.model.relation_embeddings(x[:, 1:2])
relation_embeddings = relation_embeddings.view(batch_size, 1, -1)
num_entities, entity_embedding_size = self.model.entity_embeddings.weight.shape # noqa
tail_embeddings = self.model.entity_embeddings.weight
tail_embeddings = tail_embeddings.view(1, num_entities,
entity_embedding_size)
sum_res = head_embeddings + relation_embeddings - tail_embeddings
distances = torch.norm(sum_res, dim=-1, p=1)
return distances
def __init__(self,
relation_embeddings,
num_entities,
num_relations,
entity_embeddings,
preferred_device='cpu'):
super(TranseUnit, self).__init__()
if preferred_device == 'cuda':
preferred_device = 'gpu'
self.model = TransE(preferred_device=preferred_device,
random_seed=1234,
embedding_dim=entity_embeddings.shape[1],
margin_loss=1,
normalization_of_entities=1)
self.model.num_entities = num_entities
self.model.num_relations = num_relations
self.model.relation_embeddings = nn.Embedding.from_pretrained(
relation_embeddings) # noqa
self.model.entity_embeddings = nn.Embedding.from_pretrained(
entity_embeddings) # noqa