besskge.scoring.MatrixDecompositionScoreFunction

class besskge.scoring.MatrixDecompositionScoreFunction(negative_sample_sharing)[source]

Base class for matrix-decomposition scoring functions.

Initialize matrix-decomposition scoring function.

Parameters:

negative_sample_sharing (bool) – see BaseScoreFunction

broadcasted_dot_product(v1, v2)[source]

Broadcasted dot product of queries against sets of entities.

For each query and candidate, computes the dot product of the embeddings.

Parameters:

  • v1 (Tensor) – shape: (batch_size, embedding_size) Batch queries.

  • v2 (Tensor) – shape: (B, n_neg, embedding_size) with B = 1, batch_size Negative embeddings to score against queries.

Return type:

Tensor

Returns:

shape: (batch_size, B * n_neg) if BaseScoreFunction.negative_sample_sharing else (batch_size, n_neg)

entity_embedding: Parameter

Entity embedding table

forward(head_emb, relation_id, tail_emb)

see BaseScoreFunction.score_triple()

Return type:

Tensor

Parameters:
negative_sample_sharing: bool

Share negative entities to construct negative samples

reduce_embedding(v)[source]

Sum reduction along the embedding dimension.

Parameters:

v (Tensor) – shape: (*, embedding_size) The tensor to reduce.

Return type:

Tensor

Returns:

shape: (*,) Sum reduction.

relation_embedding: Parameter

Relation embedding table

abstract score_heads(head_emb, relation_id, tail_emb)

Score sets of head entities against fixed (r,t) queries.

Parameters:

  • head_emb (Tensor) – shape: (B, n_heads, embedding_size) with B = 1, batch_size Embeddings of head entities.

  • relation_id (Tensor) – shape: (batch_size,) IDs of relation types in batch.

  • tail_emb (Tensor) – shape: (batch_size, embedding_size) Embedding of tail entities in batch.

Return type:

Tensor

Returns:

shape: (batch_size, B * n_heads) if BaseScoreFunction.negative_sample_sharing else (batch_size, n_heads). Scores of broadcasted triples.

abstract score_tails(head_emb, relation_id, tail_emb)

Score sets of tail entities against fixed (h,r) queries.

Parameters:

  • head_emb (Tensor) – shape: (batch_size, embedding_size) Embeddings of head entities in batch.

  • relation_id (Tensor) – shape: (batch_size,) IDs of relation types in batch.

  • tail_emb (Tensor) – shape: (B, n_tails, embedding_size) with B = 1, batch_size Embedding of tail entities.

Return type:

Tensor

Returns:

shape: (batch_size, B * n_tails) if BaseScoreFunction.negative_sample_sharing else (batch_size, n_tails) Scores of broadcasted triples.

abstract score_triple(head_emb, relation_id, tail_emb)

Score a batch of (h,r,t) triples.

Parameters:

  • head_emb (Tensor) – shape: (batch_size, embedding_size) Embeddings of head entities in batch.

  • relation_id (Tensor) – shape: (batch_size,) IDs of relation types in batch.

  • tail_emb (Tensor) – shape: (batch_size, embedding_size) Embeddings of tail entities in batch.

Return type:

Tensor

Returns:

shape: (batch_size,) Scores of a batch of triples.

sharding: Sharding

Sharding of entities

update_sharding(new_sharding)

Change the sharding of the entity embedding table.

Parameters:

new_sharding (Sharding) – The new entity sharding.

Return type:

None