class EmbeddingsInCanonicalShapeTests(unittest.TestCase):
"""Test get_embedding_in_canonical_shape()."""
#: The number of embeddings
num_embeddings: int = 3
#: The embedding dimension
embedding_dim: int = 2
def setUp(self) -> None:
"""Initialize embedding."""
self.embedding = Embedding(num_embeddings=self.num_embeddings,
embedding_dim=self.embedding_dim)
self.generator = torch.manual_seed(42)
self.embedding._embeddings.weight.data = torch.rand(
self.num_embeddings,
self.embedding_dim,
generator=self.generator,
)
def test_no_indices(self):
"""Test getting all embeddings."""
emb = self.embedding.get_in_canonical_shape(indices=None)
# check shape
assert emb.shape == (1, self.num_embeddings, self.embedding_dim)
# check values
exp = self.embedding(indices=None).view(1, self.num_embeddings,
self.embedding_dim)
assert torch.allclose(emb, exp)
def _test_with_indices(self, indices: torch.Tensor) -> None:
"""Help tests with index."""
emb = self.embedding.get_in_canonical_shape(indices=indices)
# check shape
num_ind = indices.shape[0]
assert emb.shape == (num_ind, 1, self.embedding_dim)
# check values
exp = torch.stack([self.embedding(i) for i in indices],
dim=0).view(num_ind, 1, self.embedding_dim)
assert torch.allclose(emb, exp)
def test_with_consecutive_indices(self):
"""Test to retrieve all embeddings with consecutive indices."""
indices = torch.arange(self.num_embeddings, dtype=torch.long)
self._test_with_indices(indices=indices)
def test_with_indices_with_duplicates(self):
"""Test to retrieve embeddings at random positions with duplicate indices."""
indices = torch.randint(
self.num_embeddings,
size=(2 * self.num_embeddings, ),
dtype=torch.long,
generator=self.generator,
)
self._test_with_indices(indices=indices)
def setUp(self) -> None:
"""Initialize embedding."""
self.embedding = Embedding(num_embeddings=self.num_embeddings, embedding_dim=self.embedding_dim)
self.generator = torch.manual_seed(42)
self.embedding._embeddings.weight.data = torch.rand(
self.num_embeddings,
self.embedding_dim,
generator=self.generator,
)
def test_score_hrt_manual(self):
"""Manually test interaction function of TransR."""
# entity embeddings
weights = torch.as_tensor(data=[[2.0, 2.0], [3.0, 3.0]],
dtype=torch.float)
entity_embeddings = Embedding(
num_embeddings=2,
embedding_dim=2,
)
entity_embeddings._embeddings.weight.data.copy_(weights)
self.instance.entity_embeddings = entity_embeddings
# relation embeddings
relation_weights = torch.as_tensor(data=[[4.0, 4], [5.0, 5.0]],
dtype=torch.float)
relation_embeddings = Embedding(
num_embeddings=2,
embedding_dim=2,
)
relation_embeddings._embeddings.weight.data.copy_(relation_weights)
self.instance.relation_embeddings = relation_embeddings
relation_projection_weights = torch.as_tensor(
data=[[5.0, 5.0, 6.0, 6.0], [7.0, 7.0, 8.0, 8.0]],
dtype=torch.float)
relation_projection_embeddings = Embedding(
num_embeddings=2,
embedding_dim=4,
)
relation_projection_embeddings._embeddings.weight.data.copy_(
relation_projection_weights)
self.instance.relation_projections = relation_projection_embeddings
# Compute Scores
batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long)
scores = self.instance.score_hrt(hrt_batch=batch)
self.assertEqual(scores.shape[0], 2)
self.assertEqual(scores.shape[1], 1)
first_score = scores[0].item()
# second_score = scores[1].item()
self.assertAlmostEqual(first_score, -32, delta=0.01)
def test_score_hrt_manual(self):
"""Manually test interaction function of TransD."""
# entity embeddings
weights = torch.as_tensor(data=[[2.0, 2.0], [4.0, 4.0]],
dtype=torch.float)
entity_embeddings = Embedding(
num_embeddings=2,
embedding_dim=2,
)
entity_embeddings._embeddings.weight.data.copy_(weights)
projection_weights = torch.as_tensor(data=[[3.0, 3.0], [2.0, 2.0]],
dtype=torch.float)
entity_projection_embeddings = Embedding(
num_embeddings=2,
embedding_dim=2,
)
entity_projection_embeddings._embeddings.weight.data.copy_(
projection_weights)
self.instance.entity_representations = torch.nn.ModuleList(
[entity_embeddings, entity_projection_embeddings])
# relation embeddings
relation_weights = torch.as_tensor(data=[[4.0], [4.0]],
dtype=torch.float)
relation_embeddings = Embedding(
num_embeddings=2,
embedding_dim=1,
)
relation_embeddings._embeddings.weight.data.copy_(relation_weights)
relation_projection_weights = torch.as_tensor(data=[[5.0], [3.0]],
dtype=torch.float)
relation_projection_embeddings = Embedding(
num_embeddings=2,
embedding_dim=1,
)
relation_projection_embeddings._embeddings.weight.data.copy_(
relation_projection_weights)
self.instance.relation_representations = torch.nn.ModuleList(
[relation_embeddings, relation_projection_embeddings])
# Compute Scores
batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long)
scores = self.instance.score_hrt(hrt_batch=batch)
self.assertEqual(scores.shape[0], 2)
self.assertEqual(scores.shape[1], 1)
first_score = scores[0].item()
self.assertAlmostEqual(first_score, -16, delta=0.01)
# Use different dimension for relation embedding: relation_dim > entity_dim
# relation embeddings
relation_weights = torch.as_tensor(data=[[3.0, 3.0, 3.0],
[3.0, 3.0, 3.0]],
dtype=torch.float)
relation_embeddings = Embedding(
num_embeddings=2,
embedding_dim=3,
)
relation_embeddings._embeddings.weight.data.copy_(relation_weights)
relation_projection_weights = torch.as_tensor(data=[[4.0, 4.0, 4.0],
[4.0, 4.0, 4.0]],
dtype=torch.float)
relation_projection_embeddings = Embedding(
num_embeddings=2,
embedding_dim=3,
)
relation_projection_embeddings._embeddings.weight.data.copy_(
relation_projection_weights)
self.instance.relation_representations = torch.nn.ModuleList(
[relation_embeddings, relation_projection_embeddings])
# Compute Scores
batch = torch.as_tensor(data=[[0, 0, 0]], dtype=torch.long)
scores = self.instance.score_hrt(hrt_batch=batch)
self.assertAlmostEqual(scores.item(), -27, delta=0.01)
batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long)
scores = self.instance.score_hrt(hrt_batch=batch)
self.assertEqual(scores.shape[0], 2)
self.assertEqual(scores.shape[1], 1)
first_score = scores[0].item()
second_score = scores[1].item()
self.assertAlmostEqual(first_score, -27, delta=0.01)
self.assertAlmostEqual(second_score, -27, delta=0.01)
# Use different dimension for relation embedding: relation_dim < entity_dim
# entity embeddings
weights = torch.as_tensor(data=[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
dtype=torch.float)
entity_embeddings = Embedding(
num_embeddings=2,
embedding_dim=3,
)
entity_embeddings._embeddings.weight.data.copy_(weights)
projection_weights = torch.as_tensor(data=[[2.0, 2.0, 2.0],
[2.0, 2.0, 2.0]],
dtype=torch.float)
entity_projection_embeddings = Embedding(
num_embeddings=2,
embedding_dim=3,
)
entity_projection_embeddings._embeddings.weight.data.copy_(
projection_weights)
self.instance.entity_representations = torch.nn.ModuleList(
[entity_embeddings, entity_projection_embeddings])
# relation embeddings
relation_weights = torch.as_tensor(data=[[3.0, 3.0], [3.0, 3.0]],
dtype=torch.float)
relation_embeddings = Embedding(
num_embeddings=2,
embedding_dim=2,
)
relation_embeddings._embeddings.weight.data.copy_(relation_weights)
relation_projection_weights = torch.as_tensor(data=[[4.0, 4.0],
[4.0, 4.0]],
dtype=torch.float)
relation_projection_embeddings = Embedding(
num_embeddings=2,
embedding_dim=2,
)
relation_projection_embeddings._embeddings.weight.data.copy_(
relation_projection_weights)
self.instance.relation_representations = torch.nn.ModuleList(
[relation_embeddings, relation_projection_embeddings])
# Compute Scores
batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long)
scores = self.instance.score_hrt(hrt_batch=batch)
self.assertEqual(scores.shape[0], 2)
self.assertEqual(scores.shape[1], 1)
first_score = scores[0].item()
second_score = scores[1].item()
self.assertAlmostEqual(first_score, -18, delta=0.01)
self.assertAlmostEqual(second_score, -18, delta=0.01)
