The repository includes the source codes in our paper entitled "Improving Code Representation Learning via Multi-view Contrastive Graph Pooling for Abstract Syntax Tree". The model can be applied on code classification task with two dataset OJ-104 and OJ-DEEP according to our paper.
-
Download data from the data Source: https://sites.google.com/site/treebasedcnn/
-
Install all the dependent packages via pip according to your environment:
Python >= 3.8 pip install pytorch=1.9.1 torch-cluster=1.5.9 torch-geometric=2.0.2 torch-scatter=2.0.8 torch-sparse=0.6.11 - Modify the config json file to train and test the model.
└─ dataset ├─ OJDatasetLoader:for loading OJ dataset ├─ OJDeepDatasetLoader:for loading OJ-DEEP dataset └─ layers ├─ GNN_node: Graph Neural Network Layer ├─ GNN_virtual_node: Graph Neural Network Layer with virtual node ├─ encoders: Node and edge encoders module ├─ gat_layer: Graph Attention Layer ├─ gcn_layer: Graph Convolution Layer ├─ gin_layer: Graph Isomorphism Network Layer ├─ graphsage_layer: GraphSAGE Layer ├─ mlp_layer: Multi-Layer Perceptron Layer └─ models ├─ model: Model definition └─ pooling ├─ lastreadout_layer: Last Readout Layer ├─ pooling_layer: Pooling Layer, main inplement of TreePool └─ utils └─ file: File operation utilities └─ parameter: Parameter utilities └─ train_util:Training utilities └─ train_parallel: main function of train, val and test procedure