Add AdvOC method files

core/AdvOC/ast_graph_encoder.py

@@ -0,0 +1,25 @@
+import os
+import torch
+import torch.nn as nn
+
+from constants import GNN_LAYER_TIMESTEPS, GNN_DROPOUT_RATE
+from gnn import GatedGraphNeuralNetwork, AdjacencyList
+
+class ASTGraphEncoder(nn.Module):
+    """Encoder which learns a representation of a method's AST. The underlying network is a Gated Graph Neural Network."""
+    def __init__(self, hidden_size, num_edge_types):
+        super(ASTGraphEncoder, self).__init__()
+        self.hidden_size = hidden_size
+        self.num_edge_types = num_edge_types
+        self.gnn = GatedGraphNeuralNetwork(self.hidden_size, self.num_edge_types,
+            [GNN_LAYER_TIMESTEPS], {}, GNN_DROPOUT_RATE, GNN_DROPOUT_RATE)
+
+    def forward(self, initial_node_representation, graph_batch, device):
+        adjacency_lists = []
+        for edge_type in range(self.num_edge_types):
+            adjacency_lists.append(AdjacencyList(node_num=graph_batch.num_nodes,
+                adj_list=graph_batch.edges[edge_type], device=device))
+        node_representations = self.gnn.compute_node_representations(
+            initial_node_representation=initial_node_representation, adjacency_lists=adjacency_lists)
+        hidden_states = node_representations[graph_batch.node_positions]
+        return hidden_states

core/AdvOC/clf_metrics.py

@@ -0,0 +1,53 @@
+from typing import Dict
+
+import numpy as np
+from sklearn.metrics import roc_auc_score, precision_recall_curve, auc
+
+
+class ClfMetric():
+    def __init__(self, threshold: float = 0.5):
+        """
+        :param threshold: The threshold for classification
+        """
+        super().__init__()
+        self.threshold = threshold
+
+    def eval(self, probs: np.array, labels: np.array) -> Dict:
+        """
+        :param probs: the prob of each class for each sample
+        :param labels: true labels
+        :return: precision, recall, f1, roc_auc, prc_auc
+        """
+        # calculate TPR, FPR, TNR, FNR & AUC
+        assert probs.shape[0] == labels.shape[0]
+        preds = (probs[:, 1] > self.threshold).astype(int)
+        TP, FP, TN, FN = 0, 0, 0, 0
+        for pred, label in zip(preds, labels):
+            if pred == label:
+                if pred == 1:
+                    TP += 1
+                else:
+                    TN += 1
+            else:
+                if pred == 1:
+                    FP += 1
+                else:
+                    FN += 1
+        roc_auc = roc_auc_score(labels, probs[:, -1])
+        precision = TP / (TP + FP)
+        recall = TP / (TP + FN)
+        f1 = (2 * TP) / (2 * TP + FP + FN)
+        pres, recs, thres = precision_recall_curve(labels, probs[:, -1])
+        prc_auc = auc(recs, pres)
+        return {
+            "precision": precision,
+            "recall": recall,
+            "f1": f1,
+            "roc_auc": roc_auc,
+            "prc_auc": prc_auc,
+            "TP": TP,
+            "FP": FP,
+            "TN": TN,
+            "FN": FN,
+            'score': recall * recall * (TP + TN + FP + FN) / (TP + FP)
+        }

core/AdvOC/constants.py

@@ -0,0 +1,53 @@
+from pathlib import Path
+
+START = '<sos>'
+END = '<eos>'
+NL_EMBEDDING_SIZE = 64
+CODE_EMBEDDING_SIZE = 64
+HIDDEN_SIZE = 64
+DROPOUT_RATE = 0.5
+NUM_LAYERS = 2
+LR = 0.0001
+BATCH_SIZE = 128
+ALPHA = 0.4
+BETA = 0.1
+MAX_EPOCHS = 100
+PATIENCE = 10
+
+LOG_EVERY_N_BATCHES = 200
+EVAL_EVERY_N_BATCHES = 4000
+
+GNN_HIDDEN_SIZE = 64  # not used
+GNN_LAYER_TIMESTEPS = 8
+GNN_DROPOUT_RATE = 0.0
+SRC_EMBEDDING_SIZE = 8
+NODE_EMBEDDING_SIZE = 64
+
+MULTI_HEADS = 4
+
+# by time
+DATA_PATH = 'core/AdvOC/data_process'  # TODO
+
+# by project
+# DATA_PATH = './'
+
+AST_TYPE_DICT = 'core/AdvOC/data_process/ast_type_dict.json'
+
+
+graph_flag = True
+if graph_flag:
+    VOCAB_FILE = str(Path(DATA_PATH) / 'vocab.json')
+    LOAD_EMBEDDINGS = False
+    EMBEDDING_PATH = str(Path(DATA_PATH) / 'embed.pkl')
+    FREEZE_EMBEDDING = False
+
+    NODE_EMBEDDING_SIZE = 128
+    MAX_SUBTOKENS = 8
+    NC_EDGE_TYPES = 3  # nl-code edge
+    MAX_CONTEXT_LENGHT = 160  # for module_manager
+    MAX_SCRIPT_LENGTH = 10
+
+debug_flag = False
+if debug_flag:
+    LOG_EVERY_N_BATCHES = 1
+    EVAL_EVERY_N_BATCHES = 10