import inspect
import logging
from typing import Any, Dict, List, Optional, Tuple, Union, cast
import lightning as pl
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from dicee.losses.custom_losses import (
ACLS,
AdaptiveLabelRelaxationLoss,
AdaptiveLabelSmoothingLoss,
AggregatedLSandLR,
CombinedAdaptiveLSandAdaptiveLR,
CombinedLSandLR,
ConfidenceBasedAdaptiveLabelRelaxationLoss,
DefaultBCELoss,
LabelRelaxationLoss,
LabelSmoothingLoss,
WeightedBCELoss,
)
from .adopt import ADOPT
logger = logging.getLogger(__name__)
[docs]
class BaseKGELightning(pl.LightningModule):
"""Thin PyTorch Lightning wrapper shared by all KGE models.
Provides the standard Lightning training loop hooks (``training_step``,
``on_train_epoch_end``, ``configure_optimizers``) as well as a helper
for reporting model size. All concrete KGE models should extend
:class:`BaseKGE` rather than this class directly.
"""
args: Dict[str, Any]
loss: nn.Module
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.training_step_outputs = []
# Epoch counter for trainer backends that call training_step/loss_function
# directly instead of going through a real pl.Trainer (whose current_epoch
# property would otherwise stay 0). See _current_training_epoch().
self._native_current_epoch = 0
def _current_training_epoch(self) -> int:
"""Current training epoch, regardless of trainer backend.
Under the ``PL`` (Lightning) trainer this defers to the built-in
``current_epoch`` property (backed by ``self.trainer``). Native
trainers that call ``training_step`` directly - currently
``torchCPUTrainer`` - update ``_native_current_epoch`` themselves at
the start of each epoch instead.
"""
if getattr(self, "_trainer", None) is not None:
return self.current_epoch
return self._native_current_epoch
[docs]
def mem_of_model(self) -> Dict:
""" Size of model in MB and number of params"""
# https://discuss.pytorch.org/t/finding-model-size/130275/2
# (2) Store NumParam and EstimatedSizeMB
num_params = sum(p.numel() for p in self.parameters())
# Not quite sure about EstimatedSizeMB ?
buffer_size = 0
for buffer in self.buffers():
buffer_size += buffer.nelement() * buffer.element_size()
return {'EstimatedSizeMB': (num_params + buffer_size) / 1024 ** 2, 'NumParam': num_params}
[docs]
def training_step(self, batch, batch_idx=None):
"""Execute one optimisation step for the given mini-batch.
Handles two- and three-element batches produced by the different
dataset classes (``KvsAll`` / ``NegSample`` vs. ``KvsSample``).
Parameters
----------
batch : tuple
``(x, y)`` for standard scoring, or ``(x, y_select, y)`` for
sample-based labelling.
batch_idx : int, optional
Index of the current batch (unused, kept for Lightning API compat).
Returns
-------
torch.FloatTensor
Scalar loss value for this batch.
"""
if len(batch)==2:
# Default
x_batch, y_batch = batch
yhat_batch = self.forward(x_batch)
elif len(batch)==3:
# KvsSample or 1vsSample
x_batch, y_select, y_batch = batch
yhat_batch = self.forward((x_batch,y_select))
else:
raise RuntimeError("Invalid batch received.")
loss_batch = self.loss_function(yhat_batch, y_batch, current_epoch=self._current_training_epoch())
self.training_step_outputs.append(loss_batch.item())
# Only log when using PyTorch Lightning trainer
# Check private _trainer attribute to avoid RuntimeError from property getter
if hasattr(self, '_trainer') and self._trainer is not None:
self.log("loss",
value=loss_batch,
on_step=True,
on_epoch=True,
prog_bar=True,
sync_dist=True,
logger=False)
return loss_batch
[docs]
def loss_function(self, yhat_batch: torch.Tensor, y_batch: torch.Tensor,
current_epoch: Optional[int] = None) -> torch.Tensor:
"""Compute the loss between model predictions and targets.
Delegates to ``self.loss`` which is configured in
:class:`BaseKGE.__init__` based on the scoring technique
(``BCEWithLogitsLoss`` for entity/relation prediction,
``CrossEntropyLoss`` for classification, ``MarginRankingLoss``
for ``scoring_technique='NegSampleMargin'``) or on ``loss_fn``
(one of the classes in :mod:`dicee.losses.custom_losses`).
For ``NegSampleMargin``, *yhat_batch*/*y_batch* still arrive in the
flat ``NegSample`` layout (positives first, followed by
``neg_ratio`` tiled blocks of negatives, row-aligned by triple).
This method splits that flat batch into positive/negative score
pairs before delegating to ``MarginRankingLoss``.
When ``adversarial_temperature`` is set (grouped/strict negative
sampling), *yhat_batch*/*y_batch* instead arrive as ``(batch,
1 + neg_ratio)`` groups with the positive in column 0, produced by
:class:`~dicee.dataset_classes._negative_sampling.GroupedNegativeSamplingDataset`.
This delegates to :func:`~dicee.models.sampled_loss.grouped_adversarial_bce`
instead of ``self.loss``.
Some ``loss_fn`` classes (e.g. ``CombinedLSandLR``,
``AdaptiveLabelSmoothingLoss``) declare a ``current_epoch`` parameter
on their own ``forward`` to schedule their behavior over training.
``self._loss_needs_epoch`` (computed once in ``BaseKGE.__init__`` by
inspecting ``self.loss.forward``) says whether that argument should
be forwarded to ``self.loss`` here.
Parameters
----------
yhat_batch : torch.FloatTensor
Model output scores, shape ``(batch_size, *)``.
y_batch : torch.FloatTensor
Ground-truth labels of the same shape as *yhat_batch*.
current_epoch : int, optional
Current training epoch, forwarded to ``self.loss`` only when it
declares a ``current_epoch`` parameter. Callers that never use an
epoch-aware ``loss_fn`` may omit this.
Returns
-------
torch.FloatTensor
Scalar loss value.
"""
if self.args.get("scoring_technique") == "NegSampleMargin":
pos_mask = y_batch >= 0.5
pos_scores = yhat_batch[pos_mask]
neg_scores = yhat_batch[~pos_mask]
neg_ratio = neg_scores.numel() // pos_scores.numel()
pos_scores = pos_scores.repeat(neg_ratio)
target = torch.ones_like(pos_scores)
return cast(torch.Tensor, self.loss(pos_scores, neg_scores, target))
temperature = self.args.get("adversarial_temperature")
if temperature is not None:
from .sampled_loss import grouped_adversarial_bce
return grouped_adversarial_bce(yhat_batch, y_batch, temperature)
if getattr(self, "_loss_needs_epoch", False):
if current_epoch is None:
current_epoch = self._current_training_epoch()
return cast(torch.Tensor, self.loss(yhat_batch, y_batch, current_epoch))
return cast(torch.Tensor, self.loss(yhat_batch, y_batch))
[docs]
def on_train_epoch_end(self, *args, **kwargs):
if len(args) >= 1:
raise RuntimeError(f"Arguments must not be empty:{args}")
if len(kwargs) >= 1:
raise RuntimeError(f"Keyword Arguments must not be empty:{kwargs}")
self.loss_history.append(sum(self.training_step_outputs) / len(self.training_step_outputs))
self.training_step_outputs.clear()
[docs]
def test_epoch_end(self, outputs: List[Any]):
""" """
[docs]
def test_dataloader(self) -> None:
pass
[docs]
def val_dataloader(self) -> None:
pass
[docs]
def predict_dataloader(self) -> None:
pass
[docs]
def train_dataloader(self) -> None:
pass
[docs]
class BaseKGE(BaseKGELightning):
"""Base class for all Knowledge Graph Embedding models.
Inherits the Lightning training loop from :class:`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:
* :meth:`forward_triples` — score a batch of ``(h, r, t)`` triples.
* :meth:`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``.
"""
def __init__(self, args: dict):
super().__init__()
self.args = args
self.embedding_dim = None
self.num_entities: Optional[int] = None
self.num_relations: Optional[int] = None
self.num_tokens = None
self.learning_rate = None
self.apply_unit_norm = None
self.input_dropout_rate = None
self.hidden_dropout_rate = None
self.optimizer_name = None
self.feature_map_dropout_rate = None
self.kernel_size = None
self.num_of_output_channels = None
self.weight_decay = None
if args.get("scoring_technique") == "NegSampleMargin":
self.loss = torch.nn.MarginRankingLoss(margin=args.get("margin", 4.0))
else:
self.loss = torch.nn.BCEWithLogitsLoss()
self.selected_optimizer = None
self.normalizer_class = None
self.normalize_head_entity_embeddings = IdentityClass()
self.normalize_relation_embeddings = IdentityClass()
self.normalize_tail_entity_embeddings = IdentityClass()
self.hidden_normalizer = IdentityClass()
self.param_init = IdentityClass
self.init_params_with_sanity_checking()
# Dropouts
self.input_dp_ent_real = torch.nn.Dropout(self.input_dropout_rate)
self.input_dp_rel_real = torch.nn.Dropout(self.input_dropout_rate)
self.hidden_dropout = torch.nn.Dropout(self.input_dropout_rate)
# average minibatch loss per epoch
self.loss_history = []
self.byte_pair_encoding = self.args.get("byte_pair_encoding", False)
self.max_length_subword_tokens = self.args.get("max_length_subword_tokens", None)
self.block_size=self.args.get("block_size", None)
self.defer_large_embeddings = bool(self.args.get("fsdp_sharded_entity", False))
if self.byte_pair_encoding and self.args['model'] != "BytE":
self.token_embeddings = torch.nn.Embedding(self.num_tokens, self.embedding_dim)
self.param_init(self.token_embeddings.weight.data)
# Reduces subword units embedding matrix from T x D into D.
self.lf = nn.Sequential(nn.Linear(self.embedding_dim * self.max_length_subword_tokens,
self.embedding_dim, bias=False))
if self.args["scoring_technique"] in ["AllvsAll", "KvsAll"]:
self.str_to_bpe_entity_to_idx = {str_ent: idx for idx, (str_ent, bpe_ent, shaped_bpe_ent) in
enumerate(self.args["ordered_bpe_entities"])}
self.bpe_entity_to_idx = {shaped_bpe_ent: idx for idx, (str_ent, bpe_ent, shaped_bpe_ent) in
enumerate(self.args["ordered_bpe_entities"])}
self.ordered_bpe_entities = torch.tensor(list(self.bpe_entity_to_idx.keys()), dtype=torch.long)
elif self.byte_pair_encoding and self.args['model'] == "BytE":
""" Transformer implements token embeddings"""
else:
self.init_entity_embeddings()
self.init_relation_embeddings()
if self.args.get("loss_fn") == "LS":
self.loss = LabelSmoothingLoss(smoothness_ratio=self.args.get("label_smoothing_rate", 0.0))
if self.args.get("loss_fn") == "LRLoss":
self.loss = LabelRelaxationLoss(alpha=self.args.get("label_relaxation_alpha", 0.1))
if self.args.get("loss_fn") == "BCELoss":
self.loss = DefaultBCELoss()
if self.args.get("loss_fn") == "WeightedBCELoss":
self.loss = WeightedBCELoss()
if self.args.get("loss_fn") == "CombinedLSandLR":
self.loss = CombinedLSandLR(smoothness_ratio=self.args.get("label_smoothing_rate", 0.0), alpha=self.args.get("label_relaxation_alpha", 0.1))
if self.args.get("loss_fn") == "AdaptiveLabelSmoothingLoss":
self.loss = AdaptiveLabelSmoothingLoss()
if self.args.get("loss_fn") == "AdaptiveLabelRelaxationLoss":
self.loss = AdaptiveLabelRelaxationLoss()
if self.args.get("loss_fn") == "ConfidenceBasedAdaptiveLabelRelaxationLoss":
self.loss = ConfidenceBasedAdaptiveLabelRelaxationLoss()
if self.args.get("loss_fn") == "CombinedAdaptiveLSandAdaptiveLR":
self.loss = CombinedAdaptiveLSandAdaptiveLR()
if self.args.get("loss_fn") == "AggregatedLSandLR":
self.loss = AggregatedLSandLR()
if self.args.get("loss_fn") == "ACLS":
self.loss = ACLS()
# Cache once whether self.loss.forward wants a current_epoch argument,
# so loss_function doesn't need to re-inspect it on every training step.
self._loss_needs_epoch = "current_epoch" in inspect.signature(self.loss.forward).parameters
[docs]
def init_entity_embeddings(self, embedding_dim: Optional[int] = None) -> None:
"""Create (or re-create) the entity embedding table.
This is the single place that honours :attr:`defer_large_embeddings`:
when entity rows are sharded across FSDP ranks the table must stay
``None`` until the trainer allocates the sharded adapter. Subclasses
that need an entity table of a non-default width must call this method
instead of assigning ``self.entity_embeddings`` directly, so the
deferral is never silently undone.
Parameters
----------
embedding_dim : Optional[int]
Width of the table. Defaults to :attr:`embedding_dim`.
"""
if self.defer_large_embeddings:
self.entity_embeddings = None
return
dim = self.embedding_dim if embedding_dim is None else embedding_dim
self.entity_embeddings = torch.nn.Embedding(self.num_entities, dim)
self.param_init(self.entity_embeddings.weight.data)
[docs]
def init_relation_embeddings(self, embedding_dim: Optional[int] = None) -> None:
"""Create (or re-create) the relation embedding table.
Relation tables are never deferred - they are small enough to be
replicated on every rank.
Parameters
----------
embedding_dim : Optional[int]
Width of the table. Defaults to :attr:`embedding_dim`.
"""
dim = self.embedding_dim if embedding_dim is None else embedding_dim
self.relation_embeddings = torch.nn.Embedding(self.num_relations, dim)
self.param_init(self.relation_embeddings.weight.data)
[docs]
def forward_byte_pair_encoded_k_vs_all(self, x: torch.LongTensor) -> torch.FloatTensor:
"""KvsAll scoring for BPE-encoded head entities and relations.
Retrieves subword-unit embeddings for the head entity and relation,
reduces them to fixed-size vectors via a linear projection, then
scores against all BPE entity embeddings.
Parameters
----------
x : torch.LongTensor
Shape ``(batch_size, 2, T)`` BPE token indices where dim 1
indexes ``[head, relation]`` and *T* is
``max_length_subword_tokens``.
Returns
-------
torch.FloatTensor
Shape ``(batch_size, num_bpe_entities)`` score matrix.
"""
# (1) Get unit normalized subword units embedding matrices: (B, T, D)
bpe_head_ent_emb, bpe_rel_ent_emb = self.get_bpe_head_and_relation_representation(x)
# Future work: Use attention to model similarity between subword units comprising head and relation
# attentive_head_rel_emb = self.attention_block(torch.cat((bpe_head_ent_emb, bpe_rel_ent_emb), 1))
# bpe_head_ent_emb = attentive_head_rel_emb[:, :self.max_length_subword_tokens, :]
# bpe_rel_ent_emb = attentive_head_rel_emb[:, self.max_length_subword_tokens:, :]
# (2) Reshaping (1) into row vectors.
B, T, D = bpe_head_ent_emb.shape
# Multi-node GPU setting.
device_r = bpe_head_ent_emb.get_device()
if device_r >= 0:
self.ordered_bpe_entities = self.ordered_bpe_entities.to(device_r)
else:
self.ordered_bpe_entities = self.ordered_bpe_entities.to("cpu")
# (3) Get unit normalized subword units embedding matrices of all entities : (E, T, D)
E = self.token_embeddings(self.ordered_bpe_entities)
# (4) Reshaping (3) into row vectors (E, T*D) .
E = E.reshape(len(E), T * D)
# (5) Reshape and Reduce from (_, T*D) into row vectors.
bpe_head_ent_emb = self.input_dp_ent_real(bpe_head_ent_emb.reshape(B, T * D))
bpe_rel_ent_emb = self.input_dp_rel_real(bpe_rel_ent_emb.reshape(B, T * D))
bpe_head_ent_emb = self.lf(bpe_head_ent_emb)
bpe_rel_ent_emb = self.lf(bpe_rel_ent_emb)
E = self.lf(E)
return self.k_vs_all_score(bpe_head_ent_emb, bpe_rel_ent_emb, E)
[docs]
def forward_byte_pair_encoded_triple(self, x: Tuple[torch.LongTensor, torch.LongTensor]) -> torch.FloatTensor:
"""NegSample scoring for BPE-encoded ``(head, relation, tail)`` triples.
Retrieves subword-unit embeddings for all three elements and reduces
them to fixed-size vectors via a linear projection before computing
the triple score.
Parameters
----------
x : torch.LongTensor
Shape ``(batch_size, 3, T)`` BPE token indices.
Returns
-------
torch.FloatTensor
Shape ``(batch_size,)`` triple scores.
"""
bpe_head_ent_emb, bpe_rel_ent_emb, bpe_tail_ent_emb = self.get_sentence_representation(x)
B, T, C = bpe_head_ent_emb.shape
bpe_head_ent_emb = bpe_head_ent_emb.reshape(B, T * C)
bpe_rel_ent_emb = bpe_rel_ent_emb.reshape(B, T * C)
bpe_tail_ent_emb = bpe_tail_ent_emb.reshape(B, T * C)
bpe_triple_score = self.score(self.lf(bpe_head_ent_emb), self.lf(bpe_rel_ent_emb), self.lf(bpe_tail_ent_emb))
return bpe_triple_score
[docs]
def init_params_with_sanity_checking(self) -> None:
"""Populate model hyper-parameters from ``self.args`` with safe defaults.
Reads embedding dimension, learning rate, dropout rates, normalisation
strategy, optimizer name, and parameter initialisation scheme from the
``args`` dict. Falls back to sensible defaults for any missing key so
that minimal ``args`` dicts (e.g. for unit tests) are still valid.
"""
if self.args.get('weight_decay'):
self.weight_decay = self.args['weight_decay']
else:
self.weight_decay = 0.0
if self.args.get('embedding_dim'):
self.embedding_dim = self.args['embedding_dim']
else:
self.embedding_dim = 1
self.num_entities = self.args.get('num_entities', None)
self.num_relations = self.args.get('num_relations', None)
self.num_tokens = self.args.get('num_tokens', None)
if self.args.get('learning_rate'):
self.learning_rate = self.args['learning_rate']
else:
self.learning_rate = .1
if self.args.get("input_dropout_rate"):
self.input_dropout_rate = self.args['input_dropout_rate']
else:
self.input_dropout_rate = 0.0
if self.args.get("hidden_dropout_rate"):
self.hidden_dropout_rate = self.args['hidden_dropout_rate']
else:
self.hidden_dropout_rate = 0.0
if self.args.get("model") in ['ConvQ', 'ConvO', 'ConEx', 'AConEx', 'AConvQ', 'AConvO']:
if self.args.get("kernel_size"):
self.kernel_size = self.args['kernel_size']
else:
self.kernel_size = 3
if self.args.get("num_of_output_channels"):
self.num_of_output_channels = self.args['num_of_output_channels']
else:
self.num_of_output_channels = 3
if self.args.get("feature_map_dropout_rate"):
self.feature_map_dropout_rate = self.args['feature_map_dropout_rate']
else:
self.feature_map_dropout_rate = 0.0
if self.args.get("normalization") == 'LayerNorm':
self.normalizer_class = torch.nn.LayerNorm
self.normalize_head_entity_embeddings = self.normalizer_class(self.embedding_dim)
self.normalize_relation_embeddings = self.normalizer_class(self.embedding_dim)
if self.args['scoring_technique'] in ['NegSample', 'FixedNegSample', 'KvsSample', 'FSDP1vsSample']:
self.normalize_tail_entity_embeddings = self.normalizer_class(self.embedding_dim)
elif self.args.get("normalization") == 'BatchNorm1d':
self.normalizer_class = torch.nn.BatchNorm1d
self.normalize_head_entity_embeddings = self.normalizer_class(self.embedding_dim, affine=False)
self.normalize_relation_embeddings = self.normalizer_class(self.embedding_dim, affine=False)
if self.args['scoring_technique'] in ['NegSample', 'FixedNegSample', 'KvsSample', 'FSDP1vsSample']:
self.normalize_tail_entity_embeddings = self.normalizer_class(self.embedding_dim, affine=False)
elif self.args.get("normalization") is None:
self.normalizer_class = IdentityClass
else:
raise NotImplementedError()
self.optimizer_name = self.args.get('optim',None)
if self.args.get("init_param") is None:
self.param_init = IdentityClass
elif self.args['init_param'] == 'xavier_normal':
self.param_init = torch.nn.init.xavier_normal_
else:
logger.warning(f'--init_param (***{self.args.get("init_param")}***) not found')
self.optimizer_name = IdentityClass
[docs]
def forward(self, x: Union[torch.LongTensor, Tuple[torch.LongTensor, torch.LongTensor]],
y_idx: Optional[torch.LongTensor] = None) -> torch.Tensor:
"""Route the forward pass to the appropriate scoring method.
Inspects the shape and type of *x* to decide which low-level scorer
to call:
* Tuple ``(x, y_idx)`` → :meth:`forward_k_vs_sample`
* ``(batch, 3)`` tensor → :meth:`forward_triples`
* ``(batch, 2)`` tensor → :meth:`forward_k_vs_all`
* BPE triple tensor → :meth:`forward_byte_pair_encoded_triple`
* BPE pair tensor → :meth:`forward_byte_pair_encoded_k_vs_all`
Parameters
----------
x : torch.LongTensor or Tuple[torch.LongTensor, torch.LongTensor]
Either a plain index tensor or a ``(triple_idx, target_idx)``
tuple for sample-based labelling.
y_idx : torch.LongTensor, optional
Target entity indices used by :meth:`forward_k_vs_sample`.
Ignored when *x* is a plain tensor.
Returns
-------
torch.FloatTensor
Score tensor whose shape depends on the selected scorer.
"""
if isinstance(x, tuple):
x, y_idx = x
return self.forward_k_vs_sample(x=x, target_entity_idx=y_idx)
else:
if not self.byte_pair_encoding and x.ndim == 3 and x.shape[-1] == 3:
flat_triples = cast(torch.LongTensor, x.reshape(-1, 3))
return self.forward_triples(flat_triples).reshape(x.shape[:2])
shape_info = x.shape
if len(shape_info) == 2:
batch_size, dim = x.shape
if dim == 3:
return self.forward_triples(x)
elif dim == 2:
# h, y = x[0], x[1]
# Note that y can be relation or tail entity.
return self.forward_k_vs_all(x=x)
else:
size_of_input_data = shape_info[1]
if size_of_input_data == 3:
# NegSample with BPE
return self.forward_byte_pair_encoded_triple(x=x)
elif size_of_input_data == 2:
# KvsAll with BPE
return self.forward_byte_pair_encoded_k_vs_all(x)
[docs]
def forward_triples(self, x: torch.LongTensor) -> torch.Tensor:
"""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
-------
torch.FloatTensor
Shape ``(batch_size,)`` triple scores.
"""
# (1) Retrieve embeddings & Apply Dropout & Normalization.
h_emb, r_emb, t_emb = self.get_triple_representation(x)
return self.score(h_emb, r_emb, t_emb)
[docs]
def forward_k_vs_all(self, *args, **kwargs):
"""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
-------
torch.FloatTensor
Shape ``(batch_size, num_entities)`` score matrix.
"""
raise ValueError(f'MODEL:{self.name} does not have forward_k_vs_all function')
[docs]
def forward_k_vs_sample(self, *args, **kwargs):
"""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
-------
torch.FloatTensor
Shape ``(batch_size, k)`` score matrix where *k* is the number
of sampled target entities.
"""
raise ValueError(f'MODEL:{self.name} does not have forward_k_vs_sample function')
[docs]
def get_triple_representation(self, idx_hrt) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
"""Retrieve and normalise embedding vectors for a triple index batch.
Parameters
----------
idx_hrt : torch.LongTensor
Shape ``(batch_size, 3)`` integer tensor with columns
``[head_idx, relation_idx, tail_idx]``.
Returns
-------
head_ent_emb, rel_ent_emb, tail_ent_emb : torch.FloatTensor
Each has shape ``(batch_size, embedding_dim)`` after applying the
configured dropout and normalisation.
"""
# (1) Split input into indexes.
idx_head_entity, idx_relation, idx_tail_entity = idx_hrt[:, 0], idx_hrt[:, 1], idx_hrt[:, 2]
# (2) Retrieve embeddings & Apply Dropout & Normalization
head_ent_emb = self.normalize_head_entity_embeddings(
self.input_dp_ent_real(self.entity_embeddings(idx_head_entity)))
rel_ent_emb = self.normalize_relation_embeddings(self.input_dp_rel_real(self.relation_embeddings(idx_relation)))
tail_ent_emb = self.normalize_tail_entity_embeddings(self.entity_embeddings(idx_tail_entity))
return head_ent_emb, rel_ent_emb, tail_ent_emb
[docs]
def get_head_relation_representation(self, indexed_triple) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
"""Retrieve and normalise embedding vectors for head entities and relations.
Parameters
----------
indexed_triple : torch.LongTensor
Shape ``(batch_size, 2)`` integer tensor with columns
``[head_idx, relation_idx]``.
Returns
-------
head_ent_emb, rel_ent_emb : torch.FloatTensor
Each has shape ``(batch_size, embedding_dim)`` after applying the
configured dropout and normalisation.
"""
# (1) Split input into indexes.
idx_head_entity, idx_relation = indexed_triple[:, 0], indexed_triple[:, 1]
# (2) Retrieve embeddings & Apply Dropout & Normalization
head_ent_emb = self.normalize_head_entity_embeddings(
self.input_dp_ent_real(self.entity_embeddings(idx_head_entity)))
rel_ent_emb = self.normalize_relation_embeddings(self.input_dp_rel_real(self.relation_embeddings(idx_relation)))
return head_ent_emb, rel_ent_emb
[docs]
def get_sentence_representation(self, x: torch.LongTensor) -> Tuple[
torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
"""Retrieve BPE subword-unit embeddings for a batch of triples.
Parameters
----------
x : torch.LongTensor
Shape ``(batch_size, 3, T)`` where *T* is
``max_length_subword_tokens``.
Returns
-------
head_ent_emb, rel_emb, tail_emb : torch.FloatTensor
Each has shape ``(batch_size, T, embedding_dim)``.
"""
h, r, t = x[:, 0, :], x[:, 1, :], x[:, 2, :]
head_ent_emb = self.token_embeddings(h)
rel_emb = self.token_embeddings(r)
tail_emb = self.token_embeddings(t)
return head_ent_emb, rel_emb, tail_emb
[docs]
def get_bpe_head_and_relation_representation(self, x: torch.LongTensor) -> Tuple[
torch.FloatTensor, torch.FloatTensor]:
"""Retrieve unit-normalised BPE embeddings for head entities and relations.
Each entity/relation is represented as a sequence of *T* subword
tokens. Their token embeddings are L2-normalised across the
sequence dimension so that the resulting matrix has unit Frobenius
norm.
Parameters
----------
x : torch.LongTensor
Shape ``(batch_size, 2, T)`` where dim 1 indexes
``[head, relation]`` and *T* is ``max_length_subword_tokens``.
Returns
-------
head_ent_emb, rel_emb : torch.FloatTensor
Each has shape ``(batch_size, T, embedding_dim)``, L2-normalised
over the ``(T, D)`` dimensions.
"""
# h: batchsize, T where T represents the maximum shaped token size
# h: B x T, r: B x T
h, r = x[:, 0, :], x[:, 1, :]
# B, T, D
head_ent_emb = self.token_embeddings(h)
# B, T, D
rel_emb = self.token_embeddings(r)
# A sequence of sub-list embeddings representing an embedding of a head entity should be normalized to 0.
# Therefore, the norm of a row vector obtained from T by D matrix must be 1.
# B, T, D
head_ent_emb = F.normalize(head_ent_emb, p=2, dim=(1, 2))
# B, T, D
rel_emb = F.normalize(rel_emb, p=2, dim=(1, 2))
return head_ent_emb, rel_emb
[docs]
def get_embeddings(self) -> Tuple[np.ndarray, np.ndarray]:
"""Return the entity and relation embedding matrices as numpy arrays.
Returns
-------
entity_embeddings : numpy.ndarray
Shape ``(num_entities, embedding_dim)``.
relation_embeddings : numpy.ndarray
Shape ``(num_relations, embedding_dim)``.
"""
return self.entity_embeddings.weight.data.data.detach(), self.relation_embeddings.weight.data.detach()
[docs]
class IdentityClass(torch.nn.Module):
"""No-op normalisation / dropout placeholder.
Used whenever no normalisation layer is requested (``--normalization None``).
All inputs are returned unchanged so that the rest of the model code does
not need conditional checks around normalisation calls.
"""
def __init__(self, args=None):
super().__init__()
self.args = args
[docs]
def __call__(self, x):
return x
[docs]
@staticmethod
def forward(x):
return x