KGEModel

model

class neuralkg_ind.model.KGEModel.model.Model(args)

Bases: Module

init_emb()

score_func(head_emb, relation_emb, tail_emb)

forward(triples, negs, mode)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

tri2emb(triples, negs=None, mode='single')

Get embedding of triples.

This function get the embeddings of head, relation, and tail respectively. each embedding has three dimensions.

Parameters:

• triples (tensor) – This tensor save triples id, which dimension is [triples number, 3].

• negs (tensor, optional) – This tenosr store the id of the entity to be replaced, which has one dimension. when negs is None, it is in the test/eval phase. Defaults to None.

• mode (str, optional) – This arg indicates that the negative entity will replace the head or tail entity. when it is ‘single’, it means that entity will not be replaced. Defaults to ‘single’.

Returns:

Head entity embedding. relation_emb: Relation embedding. tail_emb: Tail entity embedding.

Return type:

head_emb

training: bool

TransE

class neuralkg_ind.model.KGEModel.TransE.TransE(args)

Bases: Model

`Translating Embeddings for Modeling Multi-relational Data`_ (TransE), which represents the relationships as translations in the embedding space.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - ||h + r - t||_F\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

score_embedding(head_emb, relation_emb, tail_emb)

Calculating the triple embedding.

The formula for calculating the triple embedding is \(h + r - t\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

Returns:

The embedding of triple.

Return type:

embedding

mapping_embedding(head_emb, relation_emb)

Calculating the head and relation mapping embedding.

The formula for calculating the mapping embedding is \(h + r\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

Returns:

The embedding of mapping.

Return type:

embedding

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

training: bool

TransH

class neuralkg_ind.model.KGEModel.TransH.TransH(args)

Bases: Model

Knowledge Graph Embedding by Translating on Hyperplanes (TransH), which apply the translation from head to tail entity in a relational-specific hyperplane in order to address its inability to model one-to-many, many-to-one, and many-to-many relations.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

norm_vector

Relation-specific projection matrix, shape:[num_rel, emb_dim]

init_emb()

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - \|e'_{h,r} + d_r - e'_{t,r}\|_{p}^2\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

forward(triples, negs=None, mode='single')

The functions used in the training phase, same as TransE

get_score(batch, mode)

The functions used in the testing phase, same as TransE

training: bool

TransR

class neuralkg_ind.model.KGEModel.TransR.TransR(args)

Bases: Model

Learning Entity and Relation Embeddings for Knowledge Graph Completion`_ (TransR), which building entity and relation embeddings in separate entity space and relation spaces

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

transfer_matrix

Transfer entity and relation embedding, shape:[num_rel, emb_dim*emb_dim]

init_emb()

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - \| M_{r} {e}_h + r_r - M_{r}e_t \|_{p}^2\)

forward(triples, negs=None, mode='single')

The functions used in the training phase, calculate triple score.

get_score(batch, mode)

The functions used in the testing phase, predict triple score.

training: bool

RotatE

class neuralkg_ind.model.KGEModel.RotatE.RotatE(args)

Bases: Model

`RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space`_ (RotatE), which defines each relation as a rotation from the source entity to the target entity in the complex vector space.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim * 2].

rel_emb

Relation_embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - \|h \circ r - t\|\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

score_embedding(head_emb, relation_emb, tail_emb)

Calculating the triple embedding.

The formula for calculating the triple embedding is \(\|h \circ r - t\|\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

Returns:

The embedding of triple.

Return type:

embedding

mapping_embedding(head_emb, relation_emb)

Calculating the head and relation mapping embedding.

The formula for calculating the mapping embedding is \(h \circ r\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

Returns:

The embedding of mapping.

Return type:

embedding

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

training: bool

DistMult

class neuralkg_ind.model.KGEModel.DistMult.DistMult(args)

Bases: Model

Embedding Entities and Relations for Learning and Inference in Knowledge Bases (DistMult)

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation_embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(h^{ op} \operatorname{diag}(r) t\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

score_embedding(head_emb, relation_emb, tail_emb)

Calculating the triple embedding.

The formula for calculating the triple embedding is \(h^{ op} \operatorname{diag}(r) t\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

Returns:

The embedding of triple.

Return type:

embedding

mapping_embedding(head_emb, relation_emb)

Calculating the head and relation mapping embedding.

The formula for calculating the mapping embedding is \(h * r\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

Returns:

The embedding of mapping.

Return type:

embedding

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

training: bool

ComplEx

class neuralkg_ind.model.KGEModel.ComplEx.ComplEx(args)

Bases: Model

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is :math:`operatorname{Re}left(h^{ op} operatorname{diag}(r) overline{t}

ight)`

Args:

head_emb: The head entity embedding. relation_emb: The relation embedding. tail_emb: The tail entity embedding. mode: Choose head-predict or tail-predict.

Returns:

score: The score of triples.

score_embedding(head_emb, relation_emb, tail_emb)

Calculating the triple embedding.

The formula for calculating the score is :math:`operatorname{Re}left(h^{ op} operatorname{diag}(r) overline{t}

ight)`

Args:

head_emb: The head entity embedding. relation_emb: The relation embedding. tail_emb: The tail entity embedding.

Returns:

embedding: The embedding of triple.

mapping_embedding(head_emb, relation_emb)

Calculating the head and relation mapping embedding.

The formula for calculating the score is \(\operatorname{Re}\left(h^{ op} \operatorname{diag}(r))\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

Returns:

The embedding of mapping.

Return type:

embedding

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

training: bool

SimplE

class neuralkg_ind.model.KGEModel.SimplE.SimplE(args)

Bases: Model

SimplE Embedding for Link Prediction in Knowledge Graphs (SimpleE), which presents a simple enhancement of CP (which we call SimplE) to allow the two embeddings of each entity to be learned dependently.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_h_emb

Entity embedding, shape:[num_ent, emb_dim].

ent_t_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation_embedding, shape:[num_rel, emb_dim].

rel_inv_emb

Inverse Relation_embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(hh_emb, rel_emb, tt_emb, ht_emb, rel_inv_emb, th_emb)

Calculating the score of triples.

Parameters:

• hh_emb – The head entity embedding on embedding h.

• rel_emb – The relation embedding.

• tt_emb – The tail entity embedding on embedding t.

• ht_emb – The tail entity embedding on embedding h.

• rel_inv_emb – The tail entity embedding.

• th_emb – The head entity embedding on embedding t.

Returns:

The score of triples.

Return type:

score

l2_loss()

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

get_emb(triples, negs=None, mode='single')

training: bool

ConvE

class neuralkg_ind.model.KGEModel.ConvE.ConvE(args)

Bases: Model

Convolutional 2D Knowledge Graph Embeddings (ConvE), which use a 2D convolution network for embedding representation.

args

Model configuration parameters.

init_emb(args)

Initialize the convolution layer and embeddings .

Parameters:

• conv1 – The convolution layer.

• fc – The full connection layer.

• bn0 – The batch Normalization layer.

• bn1 – The batch Normalization layer.

• bn2 – The batch Normalization layer.

• inp_drop – The dropout layer.

• hid_drop – The dropout layer.

• feg_drop – The dropout layer.

• emb_ent – Entity embedding, shape:[num_ent, emb_dim].

• emb_rel – Relation_embedding, shape:[num_rel, emb_dim].

concat(head_emb, rela_emb)

score_func(head_emb, relation_emb, concat=<function ConvE.concat>, choose_emb=None)

Calculate the score of the triple embedding.

This function calculate the score of the embedding. First, the entity and relation embeddings are reshaped and concatenated; the resulting matrix is then used as input to a convolutional layer; the resulting feature map tensor is vectorised and projected into a k-dimensional space.

Parameters:

• head_emb – The embedding of head entity.

• relation_emb – The embedding of relation.

Returns:

Final score of the embedding.

Return type:

score

forward(triples)

The functions used in the training phase

Parameters:

triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

Returns:

The score of triples.

Return type:

score

get_score(batch, mode='tail_predict')

The functions used in the testing phase

Parameters:

batch – A batch of data.

Returns:

The score of triples.

Return type:

score

training: bool

BoxE

class neuralkg_ind.model.KGEModel.BoxE.BoxE(args)

Bases: Model

A Box Embedding Model for Knowledge Base Completion (BoxE), which represents the bump embedding as translations in the super rectangle space.

args

Model configuration parameters.

init_emb(args)

Initialize the entity and relation embeddings in the form of a uniform distribution.

Parameters:

• arity – The maximum ary of the knowledge graph.

• epsilon – Caculate embedding_range.

• order – The distance order of score.

• margin – Caculate embedding_range and loss.

• embedding_range – Uniform distribution range.

• ent_emb – Entity embedding, shape:[num_ent, emb_dim * 2].

• rel_emb – Relation_embedding, shape:[num_rel, emb_dim * arity * 2].

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

score_func(head_emb, relation_emb, tail_emb)

Calculate the score of the triple embedding.

Parameters:

• head_emb – The embedding of head entity.

• relation_emb – The embedding of relation.

• tail_emb – The embedding of tail entity.

Returns:

The score of triples.

Return type:

score

calc_score(ent_emb, box_low, box_hig, order=2)

Calculate the norm of distance.

Parameters:

• ent_emb – The embedding of entity.

• box_low – The lower boundaries of the super rectangle.

• box_hig – The upper boundaries of the super rectangle.

• order – The order of this distance.

Returns:

The norm of distance.

dist(ent_emb, lb, ub)

Calculate the distance.

This function calculate the distance between the entity and the super rectangle. If the entity is in its target box, distance inversely correlates with box size, to maintain low distance inside large boxes and provide a gradient to keep points inside; if the entity is not in its target box, box size linearly correlates with distance, to penalize points outside larger boxes more severely.

Parameters:

• ent_emb – The embedding of entity.

• lb – The lower boundaries of the super rectangle.

• ub – The upper boundaries of the super rectangle.

Returns:

The distance between entity and super rectangle.

training: bool

CrossE

class neuralkg_ind.model.KGEModel.CrossE.CrossE(args)

Bases: Model

Interaction Embeddings for Prediction and Explanation in Knowledge Graphs (CrossE), which simulates crossover interactions(bi-directional effects between entities and relations) to select related information when predicting a new triple

args

Model configuration parameters.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

rel_reverse_emb

Reverse Relation embedding, shape:[num_rel, emb_dim].

h_weighted_vector

Interaction matrix for head entities and relations, shape:[num_rel, emb_dim]

t_weighted_vector

Interaction matrix for tail entities and relations, shape:[num_rel, emb_dim]

hr_bias

Bias for head entity and relation

tr_bias

Bias for tail entity and relation

init_emb()

score_func(ent_emb, rel_emb, weighted_vector, mode)

Calculating the score of triples.

The formula for calculating the score is \(\sigma(tanh(c_r \circ h + c_r \circ h \circ r + b)t ^ T)\)

Parameters:

• ent_emb – entity embedding

• rel_emb – relation embedding

• weighted_vector – Interaction matrix for entities and relations

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

forward(triples, mode='single')

The functions used in the training phase, calculate hr_score and tr_score simultaneously

get_score(batch, mode)

The functions used in the testing phase, predict triple score

regularize_loss(norm=2)

Add regularization to loss

training: bool

HAKE

class neuralkg_ind.model.KGEModel.HAKE.HAKE(args)

Bases: Model

Learning Hierarchy-Aware Knowledge Graph Embeddings for Link Prediction (HAKE), which maps entities into the polar coordinate system.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim * 2].

rel_emb

Relation embedding, shape:[num_rel, emb_dim * 3].

phase_weight

Calculate phase score.

modules_weight

Calculate modulus score.

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - ||h_m \circ r_m- t_m||_2 - \lambda ||\sin((h_p + r_p - t_p)/2)||_1\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

training: bool

PairRE

class neuralkg_ind.model.KGEModel.PairRE.PairRE(args)

Bases: Model

`PairRE: Knowledge Graph Embeddings via Paired Relation Vectors`_ (PairRE), which paired vectors for each relation representation to model complex patterns.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim * 2].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - || h \circ r^H - t \circ r^T ||\)

Parameters:

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters:

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns:

The score of triples.

Return type:

score

get_score(batch, mode)

The functions used in the testing phase

Parameters:

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

training: bool

DualE

class neuralkg_ind.model.KGEModel.DualE.DualE(args)

Bases: Model

`Dual Quaternion Knowledge Graph Embeddings`_ (DualE), which introduces dual quaternions into knowledge graph embeddings.

args

Model configuration parameters.

ent_emb

Entity embedding, shape:[num_ent, emb_dim * 8].

rel_emb

Relation embedding, shape:[num_rel, emb_dim * 8].

init_weights()

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is :math:` <oldsymbol{Q}_h otimes oldsymbol{W}_r^{diamond}, oldsymbol{Q}_t> `

Parameters:

• head_emb – The head entity embedding with 8 dimensionalities.

• relation_emb – The relation embedding with 8 dimensionalities.

• tail_emb – The tail entity embedding with 8 dimensionalities.

• mode – Choose head-predict or tail-predict.

Returns:

The score of triples with regul_1 and regul_2

Return type:

score

forward(triples, negs=None, mode='single')

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_score(batch, mode)

The functions used in the testing phase :param batch: A batch of data. :param mode: Choose head-predict or tail-predict.

Returns:

The score of triples.

Return type:

score

quaternion_init(in_features, out_features, criterion='he')

Quaternion-valued weight initialization the initialization scheme is optional on these four datasets, random initialization can get the same performance. This initialization scheme might be useful for the case which needs fewer epochs.

training: bool