class PartiallyRestrictedPredictionDataset(*, heads=None, relations=None, tails=None, target='tail')[source]

Bases: PredictionDataset

A dataset for scoring some links.

“Some links” is defined as

\[\mathcal{T}_{interest} = \mathcal{E}_{h} \times \mathcal{R}_{r} \times \mathcal{E}_{t}\]


For now, the target, i.e., position whose prediction method in the model is utilized, must be the full set of entities/relations.

Example .. code-block:: python

# train model; note: needs larger number of epochs to do something useful ;-) from pykeen.pipeline import pipeline result = pipeline(dataset=”nations”, model=”mure”, training_kwargs=dict(num_epochs=0))

# create prediction dataset, where the head entities is from a set of European countries, # and the relations are connected to tourism from pykeen.predict import PartiallyRestrictedPredictionDataset heads =[“netherlands”, “poland”, “uk”]) relations =[“reltourism”, “tourism”, “tourism3”]) dataset = PartiallyRestrictedPredictionDataset(heads=heads, relations=relations)

# calculate all scores for this restricted set, and keep k=3 largest from pykeen.predict import consume_scores, TopKScoreConsumer consumer = TopKScoreConsumer(k=3) consume_scores(result.model, ds, consumer) score_pack = consumer.finalize()

# add labels df =, score=score_pack.scores)

Initialize restricted prediction dataset.


NotImplementedError – if the target position is restricted, or any non-target position is not restricted