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)