Use node features for graph encoder

Author: emalgorithm

models/lstm_decoder.py

@@ -44,18 +44,6 @@ def forward(self, input, hidden, encoder_hiddens, input_seq=None):
             output = torch.cat((output, context), 2)
             output = self.attention_combine(output)
 
-        elif self.pointer_network:
-            # Create a matrix of shape [batch_size, seq_len, 2 * hidden_dim] where the last
-            # dimension is a concatenation of the ith encoder hidden state and the current decoder
-            # hidden
-            hiddens = torch.cat((encoder_hiddens, hidden[0].repeat(1, encoder_hiddens.size(1), 1)),
-                                dim=2)
-
-            # attention_coeff has shape [seq_len] and contains the attention coeffiecients for
-            # each encoder hidden state
-            attention_coeff = F.softmax(torch.squeeze(self.attention_layer(hiddens)), dim=0)
-            # TODO: This is the output already
-
         output = F.relu(output)
         output, hidden = self.gru(output, hidden)
         output = self.softmax(self.out(output[0]))

models/lstm_to_lstm.py

@@ -15,7 +15,7 @@ def __init__(self, encoder, decoder, device, graph_encoder=None, graph=False):
         assert encoder.hidden_size == decoder.hidden_size, "Hidden dimensions of encoder and decoder " \
                                                     "must be equal!"
 
-    def forward(self, sequence, target, adj=None):
+    def forward(self, sequence, target, adj=None, node_features=None):
         batch_size = 1
         max_len = target.shape[0]
         target_vocab_size = self.decoder.output_size
@@ -34,6 +34,7 @@ def forward(self, sequence, target, adj=None):
             n_tokens = sequence.size(0)
             x = torch.zeros(n_nodes, encoder_output.size(2)).to(self.device)
             x[:n_tokens, :] = encoder_output.view(encoder_output.size(1), encoder_output.size(2))
+            x[n_tokens:, :] = node_features
             graph_hidden = self.graph_encoder(x=x, adj=adj)
 
             # TODO: Combine the graph representation with the seq_encoder final layer using mlp

models/lstm_to_lstm_full_training.py

@@ -25,10 +25,11 @@ def evaluate(seq2seq_model, eval_pairs, criterion, eval='val', graph=False):
                 eval_pair = eval_pairs[i]
                 input_tensor = eval_pair[0][0].to(device)
                 adj_tensor = eval_pair[0][1].to(device)
+                node_features = eval_pair[0][2].to(device)
                 target_tensor = eval_pair[1].to(device)
 
                 output = seq2seq_model(sequence=input_tensor.view(-1), adj=adj_tensor,
-                                       target=target_tensor.view(-1))
+                                       target=target_tensor.view(-1), node_features=node_features)
             else:
                 eval_pair = eval_pairs[i]
                 input_tensor = eval_pair[0]
@@ -59,12 +60,13 @@ def evaluate(seq2seq_model, eval_pairs, criterion, eval='val', graph=False):
         return loss, f1, rouge_2, rouge_l
 
 
-def train(input_tensor, target_tensor, seq2seq_model, optimizer, criterion, graph, adj_tensor=None):
+def train(input_tensor, target_tensor, seq2seq_model, optimizer, criterion, graph,
+          adj_tensor=None, node_features=None):
     optimizer.zero_grad()
 
     if graph:
         output = seq2seq_model(sequence=input_tensor.view(-1), adj=adj_tensor,
-                               target=target_tensor.view(-1))
+                               target=target_tensor.view(-1), node_features=node_features)
     else:
         output = seq2seq_model(sequence=input_tensor.view(-1), target=target_tensor.view(-1))
 
@@ -83,9 +85,9 @@ def train_iters(seq2seq_model, n_iters, pairs, print_every=1000, learning_rate=0
     train_losses = []
     val_losses = []
 
-    # test_f1_scores = []
-    # test_rouge_2_scores = []
-    # test_rouge_l_scores = []
+    val_f1_scores = []
+    val_rouge_2_scores = []
+    val_rouge_l_scores = []
 
     print_loss_total = 0  # Reset every print_every
     plot_loss_total = 0  # Reset every plot_every
@@ -115,10 +117,11 @@ def train_iters(seq2seq_model, n_iters, pairs, print_every=1000, learning_rate=0
         if graph:
             input_tensor = training_pair[0][0].to(device)
             adj_tensor = training_pair[0][1].to(device)
+            node_features = training_pair[0][2].to(device)
             target_tensor = training_pair[1].to(device)
 
             loss, pred = train(input_tensor, target_tensor, seq2seq_model, optimizer,
-                               criterion, adj_tensor=adj_tensor, graph=graph)
+                               criterion, adj_tensor=adj_tensor, graph=graph, node_features=node_features)
         else:
             input_tensor = training_pair[0]
             target_tensor = training_pair[1]
@@ -168,11 +171,12 @@ def train_iters(seq2seq_model, n_iters, pairs, print_every=1000, learning_rate=0
             val_losses.append(val_loss)
             # test_losses.append(test_loss)
 
-            # test_f1_scores.append(test_f1)
-            # test_rouge_2_scores.append(test_rouge_2)
-            # test_rouge_l_scores.append(test_rouge_l)
+            val_f1_scores.append(val_f1)
+            val_rouge_2_scores.append(val_rouge_2)
+            val_rouge_l_scores.append(val_rouge_l)
 
-            pickle.dump([train_losses, val_losses],
-                        open(model_dir + 'res.pkl', 'wb'))
+            pickle.dump([train_losses, val_losses, val_f1_scores, val_rouge_2_scores,
+                         val_rouge_l_scores],
+                        open('results/res.pkl', 'wb'))
 
             plot_loss(train_losses, val_losses, file_path=model_dir + 'loss.jpg')

tokens_util.py

@@ -37,8 +37,23 @@ def to_tokens(self, idxs):
         return np.array([self.index2word[idx] for idx in idxs])
 
 
+def read_data():
+    data = pickle.load(open('data/methods_tokens_graphs2.pkl', 'rb'))
+    # data = pickle.load(open('data/methods_tokens_graphs.pkl', 'rb'))
+    methods_source = data['methods_source']
+    methods_graphs = data['methods_graphs']
+    methods_names = data['methods_names']
+
+    pairs = [((methods_source[i], methods_graphs[i]), methods_names[i]) for i in range(len(
+        methods_source))]
+    np.random.shuffle(pairs)
+
+    return pairs
+
+
 def read_tokens():
-    data = pickle.load(open('data/methods_tokens_data.pkl', 'rb'))
+    data = pickle.load(open('data/methods_tokens_graphs.pkl', 'rb'))
+    # data = pickle.load(open('data/methods_tokens_data.pkl', 'rb'))
     # data = pickle.load(open('../data/methods_tokens_data.pkl', 'rb'))
     methods_source = data['methods_source']
     methods_names = data['methods_names']
@@ -75,19 +90,6 @@ def prepare_data(num_samples=None):
     return lang, pairs
 
 
-def read_data():
-    data = pickle.load(open('data/methods_tokens_graphs.pkl', 'rb'))
-    methods_source = data['methods_source']
-    methods_graphs = data['methods_graphs']
-    methods_names = data['methods_names']
-
-    pairs = [((methods_source[i], methods_graphs[i]), methods_names[i]) for i in range(len(
-        methods_source))]
-    np.random.shuffle(pairs)
-
-    return pairs
-
-
 def indexes_from_sentence_tokens(lang, sentence):
     return [lang.word2index[word] for word in sentence]
 
@@ -117,9 +119,10 @@ def sparse_adj_from_edges(edges):
 
 def tensors_from_pair_tokens_graph(pair, lang):
     input_tensor = tensor_from_sentence_tokens(lang, pair[0][0])
-    input_adj = sparse_adj_from_edges(pair[0][1])
+    input_adj = sparse_adj_from_edges(pair[0][1][0])
+    node_features = torch.tensor(pair[0][1][1])
     target_tensor = tensor_from_sentence_tokens(lang, pair[1])
-    return (input_tensor, input_adj), target_tensor
+    return (input_tensor, input_adj, node_features), target_tensor
 
 
 def plot_loss(train_losses, val_losses, test_losses=None, file_path='plots/loss.jpg'):