get_relation_prediction_df

get_relation_prediction_df(model, triples_factory, head_label, tail_label, *, relations=None, **kwargs)[source]

Predict relations for the given head and tail (given by label).

Parameters

  • model (Model) – A PyKEEN model

  • triples_factory (TriplesFactory) – the training triples factory

  • head_label (str) – the string label for the head entity

  • tail_label (str) – the string label for the tail entity

  • relations (Optional[Sequence[str]]) – restrict relation prediction to the given relations

  • kwargs – additional keyword-based parameters passed to get_prediction_df().

Return type

DataFrame

Returns

shape: (k, 3) A dataframe for relation predictions. Contains either the k highest scoring triples, or all possible triples if k is None

The following example shows that after you train a model on the Nations dataset, you can score all relations w.r.t. a given head entity and tail entity.

>>> from pykeen.pipeline import pipeline
>>> from pykeen.models.predict import get_relation_prediction_df
>>> result = pipeline(
...     dataset='Nations',
...     model='RotatE',
... )
>>> df = get_relation_prediction_df(result.model, 'brazil', 'uk', triples_factory=result.training)