FeatureEnrichedEmbedding

class FeatureEnrichedEmbedding(tensor: Tensor | PretrainedInitializer, shape: None | int | Sequence[int] = None, **kwargs)

Bases: CombinedRepresentation

A combination of a static feature and a learnable representation.

In the following example, we show how to construct a feature-enriched embedding.

"""Demonstrate creating a model with a feature-enriched embedding."""

import torch

from pykeen.models import ERModel
from pykeen.nn import Embedding, FeatureEnrichedEmbedding
from pykeen.triples.generation import generate_triples_factory
from pykeen.typing import FloatTensor

n_entities = 15
n_relations = 3
n_triples = 100
embedding_dim = 32

# mock some triples
triples_factory = generate_triples_factory(n_entities, n_relations, n_triples)

# mock some feature tensor
features = torch.rand(n_entities, embedding_dim)

# this embedding is a combination of the features
# and a learnable embedding of the same shape
entity_representation = FeatureEnrichedEmbedding(features)

# we're going to use DistMult as the interaction, so
# we need a relation representation of the same size
relation_representation = Embedding(max_id=n_relations, embedding_dim=embedding_dim)

model = ERModel[FloatTensor, FloatTensor, FloatTensor](
    triples_factory=triples_factory,
    interaction="DistMult",
    entity_representations=entity_representation,
    relation_representations=relation_representation,
)

Initialize the feature-enriched embedding.

Parameters:

• tensor (FloatTensor | PretrainedInitializer) – the tensor of pretrained embeddings, or a pretrained initializer that wraps a tensor of pretrained embeddings.

• shape (tuple[int, ...]) – an explicit shape for the learned embedding. If None, it is inferred from the provided feature tensor.

• kwargs –

  Keyword arguments passed to pykeen.nn.CombinedRepresentation.__init__().

  For example, if you want to make sure that the dimensions of the output are the same as the input, set combination="ConcatProjection". to use pykeen.nn.ConcatProjectionCombination.