Add documentation

Author: emalgorithm

models/full_model.py

@@ -4,6 +4,9 @@
 
 
 class FullModel(nn.Module):
+    """
+    Complete methodNaming model.
+    """
     def __init__(self, encoder, decoder, device, graph_encoder=None, graph=False):
         super().__init__()
 
@@ -30,6 +33,7 @@ def forward(self, sequence, target, adj=None, node_features=None):
         # output contains the hidden states for all input elements
         encoder_output, hidden = self.encoder(sequence)
 
+        # graph encoder
         if self.graph:
             # graph_hidden has shape [1, 1, hidden_size] and contains a graph representation
             n_nodes = adj.size(0)
@@ -47,6 +51,7 @@ def forward(self, sequence, target, adj=None, node_features=None):
         # first input to the decoder is the <sos> tokens
         input = torch.tensor([[0]], device=self.device)
 
+        # sequence decoder
         for t in range(1, max_len):
             output, hidden = self.decoder(input, hidden, encoder_output)
             outputs[t] = output

models/gat_encoder.py

@@ -1,11 +1,12 @@
 import torch.nn as nn
-import torch.nn.functional as F
 from models.graph_convolutional_layer import GraphConvolution
 from models.graph_attention_layer import GraphAttentionLayer
-import torch
 
 
 class GATEncoder(nn.Module):
+    """
+    Graph encoder using a Graph Attention Network.
+    """
     def __init__(self, num_features, hidden_size, dropout=0):
         super(GATEncoder, self).__init__()
 

models/gcn_encoder.py

@@ -4,6 +4,9 @@
 
 
 class GCNEncoder(nn.Module):
+    """
+        Graph encoder using a Graph Convolutional Network.
+    """
     def __init__(self, num_features, hidden_size, dropout=0):
         super(GCNEncoder, self).__init__()
 

models/lstm_decoder.py

@@ -4,24 +4,28 @@
 
 
 class LSTMDecoder(nn.Module):
-    def __init__(self, hidden_size, output_size, device, attention=False, pointer_network=False):
+    """
+    Sequence decoder which makes use of a single-layer LSTM.
+    """
+    def __init__(self, hidden_size, output_size, device, attention=False):
         super(LSTMDecoder, self).__init__()
         self.hidden_size = hidden_size
         self.output_size = output_size
+        self.attention = attention
+        self.device = device
+
         self.embedding = nn.Embedding(output_size, hidden_size).to(device)
         self.lstm = nn.LSTM(hidden_size, hidden_size).to(device)
         self.out = nn.Linear(hidden_size, output_size).to(device)
         self.softmax = nn.LogSoftmax(dim=1)
-        self.attention = attention
-        self.pointer_network = pointer_network
         self.attention_layer = nn.Linear(hidden_size * 2, 1).to(device)
         self.attention_combine = nn.Linear(hidden_size * 2, hidden_size).to(device)
-        self.device = device
 
     def forward(self, input, hidden, encoder_hiddens, input_seq=None):
         # encoder_hiddens has shape [batch_size, seq_len, hidden_dim]
         output = self.embedding(input).view(1, 1, -1)
 
+        # Compute attention
         if self.attention:
             # 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

models/lstm_encoder.py

@@ -3,6 +3,9 @@
 
 
 class LSTMEncoder(nn.Module):
+    """
+        Sequence encoder which makes use of a single-layer LSTM.
+    """
     def __init__(self, input_size, hidden_size, device):
         super(LSTMEncoder, self).__init__()
         self.hidden_size = hidden_size