def _translate_triples(
triples: MappedTriples,
entity_translation: torch.LongTensor,
relation_translation: torch.LongTensor,
) -> MappedTriples:
"""
Translate triples given vectorized translations for entities and relations.
:param triples: shape: (num_triples, 3)
The original triples
:param entity_translation: shape: (num_old_entity_ids,)
The translation from old to new entity IDs.
:param relation_translation: shape: (num_old_relation_ids,)
The translation from old to new relation IDs.
:return: shape: (num_triples, 3)
The translated triples.
"""
triples = torch.stack(
[
trans[column]
for column, trans in zip(
triples.t(),
(entity_translation, relation_translation, entity_translation),
)
],
dim=-1,
)
assert (triples >= 0).all()
return triples
def mapped_triples_to_sparse_matrices(
mapped_triples: MappedTriples,
num_relations: int,
) -> Tuple[scipy.sparse.spmatrix, scipy.sparse.spmatrix]:
"""Compute relation representations as sparse matrices of entity pairs.
.. note ::
Both sets, head-tail-set, tail-head-set, have to be created at once since they need to share the same entity
pair to Id mapping.
:param mapped_triples:
The input triples.
:param num_relations:
The number of input relations
:return: shape: (num_relations, num_entity_pairs)
head-tail-set, tail-head-set matrices as {0, 1} integer matrices.
"""
num_triples = mapped_triples.shape[0]
# compute unique pairs in triples *and* inverted triples for consistent pair-to-id mapping
extended_mapped_triples = torch.cat(
[
mapped_triples,
mapped_triples.flip(-1),
],
dim=0,
)
pairs, pair_id = extended_mapped_triples[:, [0, 2]].unique(dim=0, return_inverse=True)
n_pairs = pairs.shape[0]
forward, backward = pair_id.split(num_triples)
relations = mapped_triples[:, 1]
rel = _to_one_hot(rows=relations, cols=forward, shape=(num_relations, n_pairs))
inv = _to_one_hot(rows=relations, cols=backward, shape=(num_relations, n_pairs))
return rel, inv
def _test_get_candidate_set_size(
self,
mapped_triples: MappedTriples,
restrict_entities_to: Optional[Collection[int]],
restrict_relations_to: Optional[Collection[int]],
additional_filter_triples: Union[None, MappedTriples,
List[MappedTriples]],
num_entities: Optional[int],
):
"""Test get_candidate_set_size."""
df = get_candidate_set_size(
mapped_triples=mapped_triples,
restrict_entities_to=restrict_entities_to,
restrict_relations_to=restrict_relations_to,
additional_filter_triples=additional_filter_triples,
num_entities=num_entities,
)
# return type
assert isinstance(df, pandas.DataFrame)
# columns
assert set(df.columns) == {
"index",
LABEL_HEAD,
LABEL_RELATION,
LABEL_TAIL,
f"{LABEL_HEAD}_candidates",
f"{LABEL_TAIL}_candidates",
}
# value range
if not restrict_entities_to and not restrict_relations_to:
numpy.testing.assert_array_equal(
df["index"], numpy.arange(mapped_triples.shape[0]))
numpy.testing.assert_array_equal(
df[[LABEL_HEAD, LABEL_RELATION, LABEL_TAIL]].values,
mapped_triples.numpy(),
)
for candidate_column in (f"{LABEL_HEAD}_candidates",
f"{LABEL_TAIL}_candidates"):
numpy.testing.assert_array_less(-1, df[candidate_column])
numpy.testing.assert_array_less(df[candidate_column],
self.dataset.num_entities)