46 lines
1.7 KiB
Python
46 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
|
|
|
|
class CanonicalFactPrototype(nn.Module):
|
|
"""Compose a typed symbolic state into a learned canonical prototype."""
|
|
|
|
def __init__(
|
|
self,
|
|
state_dim: int,
|
|
num_state_types: int,
|
|
num_subjects: int,
|
|
num_predicates: int,
|
|
num_objects: int,
|
|
):
|
|
super().__init__()
|
|
self.type_embedding = nn.Embedding(num_state_types, state_dim)
|
|
self.subject_embedding = nn.Embedding(num_subjects + 1, state_dim, padding_idx=0)
|
|
self.predicate_embedding = nn.Embedding(num_predicates + 1, state_dim, padding_idx=0)
|
|
self.object_embedding = nn.Embedding(num_objects + 1, state_dim, padding_idx=0)
|
|
self.consequent_predicate_embedding = nn.Embedding(
|
|
num_predicates + 1, state_dim, padding_idx=0
|
|
)
|
|
self.consequent_object_embedding = nn.Embedding(
|
|
num_objects + 1, state_dim, padding_idx=0
|
|
)
|
|
self.projection = nn.Sequential(
|
|
nn.LayerNorm(state_dim),
|
|
nn.Linear(state_dim, state_dim),
|
|
nn.GELU(),
|
|
nn.Linear(state_dim, state_dim),
|
|
)
|
|
|
|
def forward(self, targets: torch.Tensor) -> torch.Tensor:
|
|
shifted = targets[:, 1:].clamp_min(-1) + 1
|
|
value = self.type_embedding(targets[:, 0])
|
|
value = value + self.subject_embedding(shifted[:, 0])
|
|
value = value + self.predicate_embedding(shifted[:, 1])
|
|
value = value + self.object_embedding(shifted[:, 2])
|
|
value = value + self.consequent_predicate_embedding(shifted[:, 3])
|
|
value = value + self.consequent_object_embedding(shifted[:, 4])
|
|
return F.normalize(self.projection(value), dim=-1)
|