dicee.models.complex

Classes

ConEx	Convolutional ComplEx Knowledge Graph Embeddings
AConEx	Additive Convolutional ComplEx Knowledge Graph Embeddings
ComplEx	Base class for all Knowledge Graph Embedding models.
RotatE	RotatE: knowledge graph embedding by relational rotation in complex space.

Module Contents

class dicee.models.complex.ConEx(args)

Bases: dicee.models.base_model.BaseKGE

Convolutional ComplEx Knowledge Graph Embeddings

name = 'ConEx'

conv2d

fc_num_input

fc1

norm_fc1

bn_conv2d

feature_map_dropout

residual_convolution(C_1: Tuple[torch.Tensor, torch.Tensor], C_2: Tuple[torch.Tensor, torch.Tensor]) → torch.FloatTensor

Compute residual score of two complex-valued embeddings. :param C_1: a tuple of two pytorch tensors that corresponds complex-valued embeddings :param C_2: a tuple of two pytorch tensors that corresponds complex-valued embeddings :return:

forward_k_vs_all(x: torch.Tensor) → torch.FloatTensor

Score a (head, relation) batch against every entity.

Sub-classes must override this method. The default implementation raises ValueError to make missing overrides obvious at runtime.

Returns:

Shape (batch_size, num_entities) score matrix.

Return type:

torch.FloatTensor

forward_triples(x: torch.Tensor) → torch.FloatTensor

Score a batch of (head, relation, tail) index triples.

Parameters:

x (torch.LongTensor) – Shape (batch_size, 3) integer tensor where each row is [head_idx, relation_idx, tail_idx].

Returns:

Shape (batch_size,) triple scores.

Return type:

torch.FloatTensor

forward_k_vs_sample(x: torch.Tensor, target_entity_idx: torch.Tensor)

Score a (head, relation) batch against a sampled subset of entities.

Used by KvsSample and 1vsSample datasets. Sub-classes that support sample-based labelling must override this method.

Returns:

Shape (batch_size, k) score matrix where k is the number of sampled target entities.

Return type:

torch.FloatTensor

class dicee.models.complex.AConEx(args)

Bases: dicee.models.base_model.BaseKGE

Additive Convolutional ComplEx Knowledge Graph Embeddings

name = 'AConEx'

conv2d

fc_num_input

fc1

norm_fc1

bn_conv2d

feature_map_dropout

residual_convolution(C_1: Tuple[torch.Tensor, torch.Tensor], C_2: Tuple[torch.Tensor, torch.Tensor]) → torch.FloatTensor

Compute residual score of two complex-valued embeddings. :param C_1: a tuple of two pytorch tensors that corresponds complex-valued embeddings :param C_2: a tuple of two pytorch tensors that corresponds complex-valued embeddings :return:

forward_k_vs_all(x: torch.Tensor) → torch.FloatTensor

Score a (head, relation) batch against every entity.

Sub-classes must override this method. The default implementation raises ValueError to make missing overrides obvious at runtime.

Returns:

Shape (batch_size, num_entities) score matrix.

Return type:

torch.FloatTensor

forward_triples(x: torch.Tensor) → torch.FloatTensor

Score a batch of (head, relation, tail) index triples.

Parameters:

x (torch.LongTensor) – Shape (batch_size, 3) integer tensor where each row is [head_idx, relation_idx, tail_idx].

Returns:

Shape (batch_size,) triple scores.

Return type:

torch.FloatTensor

forward_k_vs_sample(x: torch.Tensor, target_entity_idx: torch.Tensor)

Score a (head, relation) batch against a sampled subset of entities.

Used by KvsSample and 1vsSample datasets. Sub-classes that support sample-based labelling must override this method.

Returns:

Shape (batch_size, k) score matrix where k is the number of sampled target entities.

Return type:

torch.FloatTensor

class dicee.models.complex.ComplEx(args)

Bases: dicee.models.base_model.BaseKGE

Base class for all Knowledge Graph Embedding models.

Inherits the Lightning training loop from BaseKGELightning and adds the embedding tables, normalisation / dropout layers, and the routing logic that dispatches forward() calls to the appropriate scoring method.

Sub-classes must implement at minimum:

• forward_triples() — score a batch of (h, r, t) triples.

• forward_k_vs_all() — score a (h, r) batch against every entity.

Parameters:

args (dict) – Flat configuration dictionary produced by vars(argparse.Namespace). Required keys: embedding_dim, num_entities, num_relations, learning_rate (or lr), optim, scoring_technique.

name = 'ComplEx'

static score(head_ent_emb: torch.FloatTensor, rel_ent_emb: torch.FloatTensor, tail_ent_emb: torch.FloatTensor)

static k_vs_all_score(emb_h: torch.FloatTensor, emb_r: torch.FloatTensor, emb_E: torch.FloatTensor)

Parameters:

• emb_h

• emb_r

• emb_E

forward_k_vs_all(x: torch.LongTensor) → torch.FloatTensor

Score a (head, relation) batch against every entity.

Sub-classes must override this method. The default implementation raises ValueError to make missing overrides obvious at runtime.

Returns:

Shape (batch_size, num_entities) score matrix.

Return type:

torch.FloatTensor

forward_k_vs_sample(x: torch.LongTensor, target_entity_idx: torch.LongTensor)

Score a (head, relation) batch against a sampled subset of entities.

Used by KvsSample and 1vsSample datasets. Sub-classes that support sample-based labelling must override this method.

Returns:

Shape (batch_size, k) score matrix where k is the number of sampled target entities.

Return type:

torch.FloatTensor

class dicee.models.complex.RotatE(args)

Bases: dicee.models.base_model.BaseKGE

RotatE: knowledge graph embedding by relational rotation in complex space.

Represents each entity as a complex vector in ℂ^(d/2), obtained by splitting its d-dimensional embedding into a real half and an imaginary half. Each relation is represented by d/2 phase angles θ that define a unit-modulus rotation r_i = e^{iθ_i}; since a phase angle needs only one real number, the relation embedding table is reinitialised at half of the entity embedding size. A true triple (h, r, t) should satisfy h ∘ r ≈ t under element-wise complex multiplication, giving the score:

f(h, r, t) = margin - ||h ∘ r - t||_2

Unlike TransE, RotatE can model symmetric, antisymmetric, inverse, and composition relation patterns.

References

Sun et al., RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space, ICLR 2019. https://arxiv.org/abs/1902.10197

name = 'RotatE'

margin = 6.0

half_dim

relation_embeddings

score(head_ent_emb: torch.FloatTensor, rel_ent_emb: torch.FloatTensor, tail_ent_emb: torch.FloatTensor) → torch.FloatTensor

Score a batch of triples using the RotatE margin-distance formula.

Parameters:

• head_ent_emb (torch.FloatTensor) – Each has shape (batch_size, embedding_dim).

• tail_ent_emb (torch.FloatTensor) – Each has shape (batch_size, embedding_dim).

• rel_ent_emb (torch.FloatTensor) – Shape (batch_size, d/2) relation phase angles θ.

Returns:

Shape (batch_size,) scores equal to margin - ||h ∘ r - t||_2.

Return type:

torch.FloatTensor

forward_k_vs_all(x: torch.Tensor) → torch.FloatTensor

KvsAll forward pass: score head/relation against all entities.

Computes margin - ||h ∘ r - e||_2 for every entity embedding e.

Parameters:

x (torch.Tensor) – Shape (batch_size, 2) integer tensor [head_idx, relation_idx].

Returns:

Shape (batch_size, num_entities) score matrix.

Return type:

torch.FloatTensor

forward_k_vs_sample(x: torch.Tensor, target_entity_idx: torch.Tensor) → torch.FloatTensor

KvsSample forward pass: score head/relation against a sampled entity subset.

Computes margin - ||h ∘ r - e||_2 for each of the k sampled entities e.

Parameters:

• x (torch.Tensor) – Shape (batch_size, 2) integer tensor [head_idx, relation_idx].

• target_entity_idx (torch.Tensor) – Shape (batch_size, k) indices of the k target entities per sample.

Returns:

Shape (batch_size, k) score matrix.

Return type:

torch.FloatTensor