class KGEModel(*args: Any, **kwargs: Any)

Bases: Module

An implementation of several knowledge graph embedding models.

The following models were …

ComplEx(head: torch.Tensor, relation: torch.Tensor, tail: torch.Tensor, mode: str) → torch.Tensor

Apply ComplEx model

Parameters:

  • head – (Tensor) heads of the triples

  • relation – (Tensor) relations of the triples

  • tail – (Tensor) heads of the triples

  • mode – (str) mode of processing

Returns:

(Tensor) calculated score

DistMult(head: torch.Tensor, relation: torch.Tensor, tail: torch.Tensor, mode: str) → torch.Tensor

Apply DistMult model

Parameters:

  • head – (Tensor) heads of the triples

  • relation – (Tensor) relations of the triples

  • tail – (Tensor) heads of the triples

  • mode – (str) mode of processing

Returns:

(Tensor) calculated score

RotatE(head: torch.Tensor, relation: torch.Tensor, tail: torch.Tensor, mode: str) → torch.Tensor

Apply RotatE model

Parameters:

  • head – (Tensor) heads of the triples

  • relation – (Tensor) relations of the triples

  • tail – (Tensor) heads of the triples

  • mode – (str) mode of processing

Returns:

(Tensor) calculated score

TransE(head: torch.Tensor, relation: torch.Tensor, tail: torch.Tensor, mode: str) → torch.Tensor

Apply TransE model

Parameters:

  • head – (Tensor) heads of the triples

  • relation – (Tensor) relations of the triples

  • tail – (Tensor) heads of the triples

  • mode – (str) mode of processing

Returns:

(Tensor) calculated score

forward(sample: torch.Tensor, mode: str = 'single') → torch.Tensor

Forward function that calculate the score of a batch of triples.

In the ‘single’ mode, sample is a batch of triple. In the ‘head-batch’ or ‘tail-batch’ mode, sample consists two part. The first part is usually the positive sample. And the second part is the entities in the negative samples. Because negative samples and positive samples usually share two elements in their triple ((head, relation) or (relation, tail)).

Parameters:

  • sample (_type_) – _description_

  • mode (str, optional) – _description_, defaults to ‘single’

Raises:

  • ValueError – _description_

  • ValueError – _description_

Returns:

_description_

Return type:

_type_

static test_step(model: torch.nn.Module, test_triples: torch.Tensor, args: Any, random_sampling: bool = False) → Dict[str, float]

Evaluate the model on tests or valid datasets

Parameters:

  • model (_type_) – _description_

  • test_triples (_type_) – _description_

  • args (_type_) – _description_

  • random_sampling (bool, optional) – _description_, defaults to False

Returns:

_description_

Return type:

_type_

static train_step(model: torch.nn.Module, optimizer: torch.optim.Optimizer, train_iterator: BidirectionalOneShotIterator, model_parameters: Any) → Dict

A single train step. Apply back-propation and return the loss

Parameters:

  • model – model to train

  • optimizer – _description_

  • train_iterator – _description_

  • model_parameters – _description_

Returns:

Dict with steps logs