predict_triples_df(model, *, triples, triples_factory=None, batch_size=None)[source]

Predict on labeled or mapped triples.

  • model (Model) – The model.

  • triples (Union[None, LongTensor, ndarray, Tuple[str, str, str], Sequence[Tuple[str, str, str]]]) –

    shape: (num_triples, 3) The triples in one of the following formats:

    • A single label-based triple.

    • A list of label-based triples.

    • An array of label-based triples

    • An array of ID-based triples.

    • None. In this case, a triples factory has to be provided, and its triples will be used.

  • triples_factory (Optional[CoreTriplesFactory]) – The triples factory. Must be given if triples are label-based. If provided and triples are ID-based, add labels to result.

  • batch_size (Optional[int]) – The batch size to use. Use None for automatic memory optimization.

Return type



columns: head_id | relation_id | tail_id | score | * A dataframe with one row per triple.


ValueError – If label-based triples have been provided, but the triples factory does not provide a mapping.

The TransE model can be trained and used to predict a given triple.

>>> from pykeen.pipeline import pipeline
>>> result = pipeline(dataset="nations", model="TransE")
>>> from pykeen.models.predict import predict_triples_df
>>> df = predict_triples_df(
...     model=result.model,
...     triples=("uk", "conferences", "brazil"),
... )