|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch_geometric.nn import RGCNConv |
|
|
|
|
| class ContextEnhancedRGCN(nn.Module): |
| def __init__( |
| self, |
| num_nodes, |
| num_relations, |
| event_node_idx_cpu, |
| event_x_cpu, |
| context_dim=768, |
| node_emb_dim=16, |
| hidden_dim=16, |
| dropout=0.2, |
| num_bases=2 |
| ): |
| super().__init__() |
|
|
| self.node_emb = nn.Embedding(num_nodes, node_emb_dim) |
| self.context_proj = nn.Linear(context_dim, node_emb_dim) |
|
|
| self.register_buffer( |
| "event_node_idx", |
| event_node_idx_cpu.to(torch.long) |
| ) |
|
|
| self.register_buffer( |
| "event_x", |
| torch.tensor(event_x_cpu, dtype=torch.float) |
| ) |
|
|
| self.rgcn1 = RGCNConv( |
| node_emb_dim, |
| hidden_dim, |
| num_relations=num_relations, |
| num_bases=num_bases |
| ) |
|
|
| self.rgcn2 = RGCNConv( |
| hidden_dim, |
| hidden_dim, |
| num_relations=num_relations, |
| num_bases=num_bases |
| ) |
|
|
| self.dropout_layer = nn.Dropout(dropout) |
|
|
| def initial_features(self): |
| x = self.node_emb.weight.clone() |
| event_feat = self.context_proj(self.event_x) |
| x[self.event_node_idx] = x[self.event_node_idx] + event_feat |
| return x |
|
|
| def encode(self, edge_index, edge_type): |
| x = self.initial_features() |
| x = self.rgcn1(x, edge_index, edge_type) |
| x = F.relu(x) |
| x = self.dropout_layer(x) |
| x = self.rgcn2(x, edge_index, edge_type) |
| return x |
|
|
| def forward(self, event_idx, candidate_seed_idx, edge_index, edge_type): |
| z = self.encode(edge_index, edge_type) |
|
|
| event_z = z[event_idx] |
| seed_z = z[candidate_seed_idx] |
|
|
| event_z = event_z.unsqueeze(1).expand_as(seed_z) |
|
|
| scores = (event_z * seed_z).sum(dim=-1) |
| return scores |
|
|
