📿 🧩 Support disconnected nodes in the Relation tokenizer

[migalkin]

A quick fix for the sampler in relation tokenizer to support graphs where num_entities is larger than pool size.

A script to reproduce the error before/after the fix

import torch.nn
from pykeen.models import NodePiece
from pykeen.triples import CoreTriplesFactory

edges = torch.tensor([
    [0, 0, 1],
    [1, 1, 0],
    [3, 1, 0],
    [3, 2, 1]
], dtype=torch.long)  # node ID 2 is missing as a disconnected node

factory = CoreTriplesFactory.create(mapped_triples=edges, num_entities=4, num_relations=3, create_inverse_triples=True)
model = NodePiece(triples_factory=factory, num_tokens=2)

Generally, taking care of disconnected nodes in NodePiece in more complex scenarios might take much more effort

[mberr]

lgtm

could you add the small snippet from the PR description as a test?

[migalkin]

lgtm

could you add the small snippet from the PR description as a test?

To test_nn/test_node_piece or cases? This test would require a different triples factory than the current default NodePieceTestCase

[mberr]

lgtm
could you add the small snippet from the PR description as a test?

To test_nn/test_node_piece or cases? This test would require a different triples factory than the current default NodePieceTestCase

maybe like this? 59e829f

[migalkin]

Cool, works as intended! Thank you :)

[mberr]

Cool, works as intended! Thank you :)

thanks for writing the tests in the first place 😉