73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class CrossAttentionBlock(nn.Module):
|
|
def __init__(self, state_dim: int, num_heads: int, dropout: float):
|
|
super().__init__()
|
|
self.cross_attention = nn.MultiheadAttention(
|
|
state_dim, num_heads, dropout=dropout, batch_first=True
|
|
)
|
|
self.norm1 = nn.LayerNorm(state_dim)
|
|
self.ffn = nn.Sequential(
|
|
nn.Linear(state_dim, state_dim * 4),
|
|
nn.GELU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(state_dim * 4, state_dim),
|
|
)
|
|
self.norm2 = nn.LayerNorm(state_dim)
|
|
|
|
def forward(
|
|
self, slots: torch.Tensor, tokens: torch.Tensor, token_mask: torch.Tensor
|
|
) -> torch.Tensor:
|
|
update, _ = self.cross_attention(
|
|
slots, tokens, tokens, key_padding_mask=~token_mask, need_weights=False
|
|
)
|
|
slots = self.norm1(slots + update)
|
|
return self.norm2(slots + self.ffn(slots))
|
|
|
|
|
|
class SlotExtractor(nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
state_dim: int,
|
|
num_slots: int,
|
|
num_layers: int,
|
|
num_heads: int,
|
|
dropout: float,
|
|
num_subjects: int,
|
|
num_predicates: int,
|
|
num_objects: int,
|
|
num_state_types: int,
|
|
):
|
|
super().__init__()
|
|
self.input_projection = nn.Linear(input_dim, state_dim)
|
|
self.slot_queries = nn.Parameter(torch.randn(num_slots, state_dim) * 0.02)
|
|
self.blocks = nn.ModuleList(
|
|
[CrossAttentionBlock(state_dim, num_heads, dropout) for _ in range(num_layers)]
|
|
)
|
|
self.subject_head = nn.Linear(state_dim, num_subjects)
|
|
self.predicate_head = nn.Linear(state_dim, num_predicates)
|
|
self.object_head = nn.Linear(state_dim, num_objects)
|
|
self.type_head = nn.Linear(state_dim, num_state_types)
|
|
self.consequent_predicate_head = nn.Linear(state_dim, num_predicates)
|
|
self.consequent_object_head = nn.Linear(state_dim, num_objects)
|
|
|
|
def forward(self, tokens: torch.Tensor, token_mask: torch.Tensor) -> dict[str, torch.Tensor]:
|
|
tokens = self.input_projection(tokens)
|
|
slots = self.slot_queries.unsqueeze(0).expand(tokens.shape[0], -1, -1)
|
|
for block in self.blocks:
|
|
slots = block(slots, tokens, token_mask)
|
|
return {
|
|
"states": slots,
|
|
"subject_logits": self.subject_head(slots),
|
|
"predicate_logits": self.predicate_head(slots),
|
|
"object_logits": self.object_head(slots),
|
|
"type_logits": self.type_head(slots),
|
|
"consequent_predicate_logits": self.consequent_predicate_head(slots),
|
|
"consequent_object_logits": self.consequent_object_head(slots),
|
|
}
|