besskge.scoring.BoxE

class besskge.scoring.BoxE(negative_sample_sharing, scoring_norm, sharding, n_relation_type, embedding_size, entity_initializer=[<function uniform_>], relation_initializer=[<function uniform_>, <function init_uniform_norm>], apply_tanh=True, dist_func_per_dim=True, eps=1e-06, inverse_relations=False)[source]

BoxE scoring function [ACLS20].

Initialize BoxE model.

Parameters:

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

  • scoring_norm (int) – see DistanceBasedScoreFunction.__init__()

  • sharding (Sharding) – Entity sharding.

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

  • embedding_size (int) – Size of final entity embeddings.

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

  • relation_initializer (Union[Tensor, List[Callable[..., Tensor]]]) – Initialization functions or table for relation embeddings. If not passing a table, two functions are needed: the initializer for head/tail box centers and the initializer for (scalar) head/tail box sizes.

  • apply_tanh (bool) – If True, bound relation box sizes and bumped entity representations with tanh. Default: True.

  • dist_func_per_dim (bool) – If True, instead of selecting between the two BoxE distance functions based on whether the bumped representation is inside or outside the relation box, make the choice separately for each dimension of the embedding space. Default: True.

  • eps (float) – Softening parameter for geometric normalization of box widths. Default: 1e-6.

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

boxe_score(bumped_ht, center_ht, width_ht, box_size)[source]

BoxE score with relation broadcasting, for optimized batched negative scoring.

Parameters:

  • bumped_ht (Tensor) – shape: (batch_size, 2, emb_size) or (batch_size, n_negative, 2, emb_size) Bumped h/t entity embeddings (heads: bumped_ht[…,0,:], tails: bumped_ht[…,1,:]).

  • center_ht (Tensor) – shape: (batch_size, 2, emb_size) or (batch_size, 1, 2, emb_size) Centers of h/t relation boxes (heads: center_ht[…,0,:], tails: center_ht[…,1,:]).

  • width_ht (Tensor) – shape: (batch_size, 2, emb_size) or (batch_size, 1, 2, emb_size) Widths of h/t relation boxes, before normalization (heads: width_ht[…,0,:], tails: width_ht[…,1,:]).

  • box_size (Tensor) – shape: (batch_size, 2) or (batch_size, 1, 2) Parameter controlling the size of the (normalized) h/t relation boxes (heads: box_size[…,0], tails: box_size[…,1]).

Return type:

Tensor

Returns:

shape: shape: (batch_size,) or (batch_size, n_negative) Negative sum of head and tail scores, as defined in [ACLS20].

broadcasted_distance(v1, v2)

Broadcasted distances of queries against sets of entities.

For each query and candidate, the computes the p-distance 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)

p-norm reduction along embedding dimension.

Parameters:

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

Return type:

Tensor

Returns:

shape: (*,) p-norm 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