TopKScoreConsumer

class TopKScoreConsumer(k=3, device=None)[source]

Bases: ScoreConsumer

Collect top-k triples & scores.

Initialize the consumer.

Parameters:

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

  • device (Union[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, target, scores)[source]

Consume scores for the given hr_batch.

Return type:

None

Parameters:

  • batch (LongTensor) –

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

  • scores (FloatTensor) –