get_prediction_df

get_prediction_df(model, triples_factory, *, head_label=None, relation_label=None, tail_label=None, targets=None, add_novelties=True, remove_known=False, testing=None, mode=None)[source]

Get predictions for the head, relation, and/or tail combination.

Note

Exactly one of head_label, relation_label and tail_label should be None. This is the position which will be predicted.

Parameters

  • model (Model) – A PyKEEN model

  • triples_factory (TriplesFactory) – the training triples factory

  • head_label (Optional[str]) – the head entity label. If None, predict heads

  • relation_label (Optional[str]) – the relation label. If None, predict relations

  • tail_label (Optional[str]) – the tail entity label. If None, predict tails

  • targets (Optional[Sequence[str]]) – restrict prediction to these targets

  • add_novelties (bool) – should the dataframe include a column denoting if the ranked head entities correspond to novel triples?

  • remove_known (bool) – should non-novel triples (those appearing in the training set) be shown with the results? On one hand, this allows you to better assess the goodness of the predictions - you want to see that the non-novel triples generally have higher scores. On the other hand, if you’re doing hypothesis generation, they may pose as a distraction. If this is set to True, then non-novel triples will be removed and the column denoting novelty will be excluded, since all remaining triples will be novel. Defaults to false.

  • testing (Optional[LongTensor]) – the mapped_triples from the testing triples factory (TriplesFactory.mapped_triples)

  • mode (Optional[Literal[‘training’, ‘validation’, ‘testing’]]) – The pass mode, which is None in the transductive setting and one of “training”, “validation”, or “testing” in the inductive setting.

Return type

DataFrame

Returns

shape: (k, 3) A dataframe with columns based on the settings or a tensor. Contains either the k highest scoring triples, or all possible triples if k is None