model.py (5941B)
1 from device import device 2 from bidirectional import BertModel 3 from torch_geometric.nn import RGCNConv 4 from transformers import AutoModel 5 import torch 6 import torch.nn as nn 7 import torch.nn.functional as F 8 PLM = 'bert-base-uncased' 9 10 11 class RelationalGCN(torch.nn.Module): 12 def __init__(self, args, input_dim=768, hidden_channels=128, 13 layers=5): 14 super(RelationalGCN, self).__init__() 15 self.args = args 16 num_relations = 12 if self.args.unidirectional or self.args.bidirectional else 10 17 self.convs = torch.nn.ModuleList( 18 [RGCNConv(in_channels=input_dim, out_channels=hidden_channels, num_relations=num_relations)] + 19 [RGCNConv(in_channels=hidden_channels, out_channels=hidden_channels, num_relations=num_relations) 20 for _ in range(layers-2)] + 21 [RGCNConv(in_channels=hidden_channels, out_channels=hidden_channels, num_relations=num_relations)] 22 ) 23 24 def forward(self, x, edge_index, edge_type, batch, inds, special_node_idx=None): 25 for layer in self.convs[:-1]: 26 x = F.relu(layer(x, edge_index, edge_type.view(-1))) 27 x = self.convs[-1](x, edge_index, edge_type.view(-1)) 28 if self.args.unidirectional: 29 special_mask = special_node_idx.view(-1, 1) 30 z = x[special_mask[:, 0]] 31 inds = torch.stack(inds, dim=-1).tolist() 32 for i in range(1, len(inds)): 33 v_c = 0 34 for j in range(i, 0, -1): 35 v_c += batch.tolist().count(j-1) 36 inds[i][0] += v_c 37 inds[i][1] += v_c 38 mask = torch.tensor(inds) 39 head = x[mask[:, 0]] 40 tail = x[mask[:, 1]] 41 avg_n = torch.cat([z, head, tail], dim=-1).view(special_mask.shape[0], 3, -1) 42 pool = torch.mean(avg_n, dim=1) 43 return pool, z 44 else: 45 inds = torch.stack(inds, dim=-1).tolist() 46 for i in range(1, len(inds)): 47 v_c = 0 48 for j in range(i, 0, -1): 49 v_c += batch.tolist().count(j-1) 50 inds[i][0] += v_c 51 inds[i][1] += v_c 52 mask = torch.tensor(inds) 53 head = x[mask[:, 0]] 54 tail = x[mask[:, 1]] 55 x = torch.cat([head, tail], dim=-1) 56 return x 57 58 59 class TextRGCN(torch.nn.Module): 60 def __init__(self, args, graph_input_dim=768): 61 super(TextRGCN, self).__init__() 62 self.hidden_channels = args.gnn_hidden 63 self.GNN_only = args.GNN_only 64 self.LM_only = args.LM_only 65 self.gcn = RelationalGCN(args, graph_input_dim, self.hidden_channels, layers=args.gnn_layers).to(device) 66 self.dropout_t = nn.Dropout(args.dropout) 67 self.dropout_g = nn.Dropout(args.dropout) 68 self.unidirectional = args.unidirectional 69 self.bidirectional = args.bidirectional 70 if args.bidirectional: 71 self.text_encoder = BertModel.from_pretrained(PLM, return_dict=True, output_hidden_states=True).to(device) 72 else: 73 self.text_encoder = AutoModel.from_pretrained(PLM, return_dict=True, output_hidden_states=True).to(device) 74 75 gnn_output_dim = self.hidden_channels if self.unidirectional else self.hidden_channels * 2 76 77 if self.GNN_only: 78 self.output = nn.Linear(gnn_output_dim, gnn_output_dim) 79 self.output2 = nn.Linear(gnn_output_dim, 1) 80 elif self.LM_only: 81 self.output = nn.Linear(self.text_encoder.config.hidden_size, self.text_encoder.config.hidden_size,) 82 self.output2 = nn.Linear(self.text_encoder.config.hidden_size, 1) 83 elif self.unidirectional: 84 self.output = nn.Linear(self.text_encoder.config.hidden_size + gnn_output_dim * 2, self.text_encoder.config.hidden_size,) 85 self.output2 = nn.Linear(self.text_encoder.config.hidden_size, 1) 86 else: 87 self.output = nn.Linear(self.text_encoder.config.hidden_size + gnn_output_dim, self.text_encoder.config.hidden_size + gnn_output_dim) 88 self.output2 = nn.Linear(self.text_encoder.config.hidden_size + gnn_output_dim, 1) 89 90 def forward(self, input_ids, attention_maskask, data, inds, special_node_indicies, special_token_indicies): 91 batch = data.batch.to(device) 92 if self.GNN_only: 93 enc = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds) 94 enc = self.dropout_g(enc) 95 elif self.LM_only: 96 enc = self.text_encoder(input_ids, attention_maskask)[1] 97 enc = self.dropout_t(enc) 98 else: 99 if self.bidirectional: 100 pooled, special_node_embedding = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds, special_node_idx=special_node_indicies) 101 z_text = self.text_encoder(input_ids, attention_maskask, special_node_embedding=special_node_embedding, spt=special_token_indicies)[1] 102 z_text = self.dropout_t(z_text) 103 enc = torch.cat([z_text, special_node_embedding, pooled], dim=-1) 104 elif self.unidirectional: 105 pooled, special_node_embedding = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds, special_node_idx=special_node_indicies) 106 z_text = self.text_encoder(input_ids, attention_maskask)[1] 107 z_text = self.dropout_t(z_text) 108 enc = torch.cat([z_text, special_node_embedding, pooled], dim=-1) 109 else: # The else statement is entered for the Disjoint and Grounded models 110 z_graph = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds, special_node_idx=None) 111 z_text = self.text_encoder(input_ids, attention_maskask)[1] 112 z_text = self.dropout_t(z_text) 113 enc = torch.cat([z_graph, z_text], dim=-1) 114 out = F.relu(self.output(enc)) 115 out = self.output2(out) 116 return out