besskge.scoring.ComplEx

class besskge.scoring.ComplEx(negative_sample_sharing, sharding, n_relation_type, embedding_size, entity_initializer=[<function init_KGE_normal>], relation_initializer=[<function init_KGE_normal>], inverse_relations=False)[source]

ComplEx scoring function [TWR+16].

Initialize ComplEx model.

Parameters:

  • negative_sample_sharing (bool) – see DistanceBasedScoreFunction.__init__()

  • sharding (Sharding) – Entity sharding.

  • n_relation_type (int) – Number of relation types in the knowledge graph.

  • embedding_size (int) – Complex size of entity and relation embeddings.

  • entity_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for entity embeddings.

  • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization function or table for relation embeddings.

  • inverse_relations (bool) – If True, learn embeddings for inverse relations. Default: False.

broadcasted_dot_product(v1, v2)

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)

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

score_heads(head_emb, relation_id, tail_emb)[source]

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.

score_tails(head_emb, relation_id, tail_emb)[source]

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.

score_triple(head_emb, relation_id, tail_emb)[source]

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