TopKScoreConsumer

class TopKScoreConsumer(k: int = 3, device: str | device | None = None)

Bases: ScoreConsumer

Collect top-k triples & scores.

Initialize the consumer.

Parameters:

• k (int) – the number of top-scored triples to collect

• device (str | device | None) – the model’s device

Attributes Summary

flatten

Methods Summary

__call__(batch, target, scores)

Consume scores for the given hr_batch.

Attributes Documentation

flatten: bool = False

Methods Documentation

__call__(batch: Tensor, target: Literal['head', 'relation', 'tail'], scores: Tensor) → None

Consume scores for the given hr_batch.

Parameters:

• batch (Tensor)

• target (Literal['head', 'relation', 'tail'])

• scores (Tensor)

Return type:

None