def score_h(self, rt_batch: torch.LongTensor, **kwargs) -> torch.FloatTensor: # noqa:D102
return marginal_score(
entity_relation_batch=rt_batch.flip(1),
per_entity=self.head_per_tail,
per_relation=self.head_per_relation,
num_entities=self.num_entities,
)
def _get_feature_vector_attention_mask(self, feature_vector_length: int, attention_mask: torch.LongTensor):
output_lengths = self._get_feat_extract_output_lengths(attention_mask.sum(-1)).to(torch.long)
batch_size = attention_mask.shape[0]
attention_mask = torch.zeros(
(batch_size, feature_vector_length), dtype=attention_mask.dtype, device=attention_mask.device
)
# these two operations makes sure that all values before the output lengths idxs are attended to
attention_mask[(torch.arange(attention_mask.shape[0], device=attention_mask.device), output_lengths - 1)] = 1
attention_mask = attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()
return attention_mask
def forward(
self,
x_e: torch.FloatTensor,
x_r: torch.FloatTensor,
edge_index: torch.LongTensor,
edge_type: torch.LongTensor,
) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
r"""
Update entity and relation representations.
.. math ::
X_E'[e] = \frac{1}{3} \left(
X_E W_s
+ \left( \sum_{h,r,e \in T} \alpha(h, e) \phi(X_E[h], X_R[r]) W_f \right)
+ \left( \sum_{e,r,t \in T} \alpha(e, t) \phi(X_E[t], X_R[r^{-1}]) W_b \right)
\right)
:param x_e: shape: (num_entities, input_dim)
The entity representations.
:param x_r: shape: (2 * num_relations, input_dim)
The relation representations (including inverse relations).
:param edge_index: shape: (2, num_edges)
The edge index, pairs of source and target entity for each triple.
:param edge_type: shape (num_edges,)
The edge type, i.e., relation ID, for each triple.
:return: shape: (num_entities, output_dim) / (2 * num_relations, output_dim)
The updated entity and relation representations.
"""
# prepare for inverse relations
edge_type = 2 * edge_type
# update entity representations: mean over self-loops / forward edges / backward edges
x_e = (self.composition(x_e, self.self_loop) @ self.w_loop +
self.message(x_e=x_e,
x_r=x_r,
edge_index=edge_index,
edge_type=edge_type,
weight=self.w_fwd) +
self.message(x_e=x_e,
x_r=x_r,
edge_index=edge_index.flip(0),
edge_type=edge_type + 1,
weight=self.w_bwd)) / 3
if self.bias:
x_e = self.bias(x_e)
x_e = self.bn(x_e)
x_e = self.activation(x_e)
# Relation transformation
x_r = self.w_rel(x_r)
return x_e, x_r
