predict_triples_df

predict_triples_df(model, *, triples, triples_factory=None, batch_size=None, mode=None)

Predict on labeled or mapped triples.

Parameters

• 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.

• 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

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

Raises

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"),
...     triples_factory=result.training,
... )