HAC onv GNN : Hierarchical Attention Based Convolutional Graph Neural Network for Code Documentation Generation in J upyter Notebooks (2021)

Abstract: Jupyter notebook allows data scientists to write machine learning code together with its documentation in cells. In this paper, we propose a new task of code documentation generation (CDG) for computational notebooks. In contrast to the previous CDG tasks which focus on generating documentation for single code snippets, in a computational notebook, one documentation in a markdown cell often corresponds to multiple code cells, and these code cells have an inherent structure. We proposed a new model (HAConvGNN) that uses a hierarchical attention mechanism to consider the relevant code cells and the relevant code tokens information when generating the documentation. Tested on a new corpus constructed from well-documented Kaggle notebooks, we show that our model outperforms other baseline models.

## introduction in recent years, computational notebooks such as jupyter have become popular programming platforms for data scientists and machine learning researchers to document ideas, write code, and visualize results, all in a single document @xcite . documentation in a notebook provides a rich medium for users to record not only what the code does, but also why they code it. this richness of content is one distinctive nature of code documentation in a notebook versus in traditional software source code. code documentation is found critical for data scientists to share or reuse code @xcite . however, research has shown that many data scientists still neglect to write appropriate documentation for their code in notebooks, as they feel writing documentation will slow down their coding process. @xcite report that among one million computational notebooks on github, 25% of them have no comment. equal contributions from the first authors: x827liu@uwaterloo.ca, dakuo.wang@ibm.com. as a first step towards building an automated documentation generation system for notebooks, in this paper we focus on the code documentation generation (cdg) task for jupyter notebooks. since there is no publicly available cdg dataset for notebooks, we construct a new dataset (notebookcdg) which contains around 28k processed code-documentation pairs extracted from 2,476 highly-ranked notebooks from kaggle competitions (details in section 3) part of work was done when xuye, april, and lingfei were at ibm. a few previous literature have explored techniques to generate documentation for software code snippet one at a time @xcite @xcite . however, in computational notebooks, one documentation (in a markdown cell) can cover more-than-one code cells after it. for instance, the ground truth text in to account for the above mentioned properties of documentation in computational notebooks, in this paper, we propose a graph-augmented encoderdecoder model to generate documentation for notebooks (section 4). in particular, our model consists of three parts: a code sequence encoder, an auxiliary documentation text encoder based on the already predicted documentation tokens, and a hierarchical attention-based convolutional graph neural network (haconvgnn) component. the first two sequence encoders encode the semantic information in code and documentation text, respectively. the graph encoder encodes the contextual abstract syntactic trees (i.e., ast extracted from the code sequence). in order to capture the relations between code sequences and the corresponding text documentations, we further employ a hierarchical attention mechanism consisting of a low-level attention module and a high-level attention module. the former attends to the token in a code sequence and the latter attends to the corresponding code cells in the ast tree. experiments show that our model achieves better performance on the notebookcdg dataset compared to baseline models on rouge scores, and in a muti-dimensional human evaluation study. base on this result, we integrated our approach into a user-facing downstream application @xcite to further explore the human-ai collaboration opportunity in the code documentation scenario. in the follow-up user study (reported seperately @xcite ), users found that the automatically generated documentation reminded them to document code they would have ignored, and improved their satisfaction with their computational notebooks. in summary, the main contributions of our work are: (1) a large-scale high quality dataset for the cdg task in the computational notebook context; in order to automate the machine learning and ai workflow, researchers have applied automation techniques on various code-related tasks @xcite , including code summarization @xcite @xcite , source code generation from natural language @xcite , and source code transformation @xcite . in this work, we focus on the code documentation generation(cdg) task. our work is closely related to code summarization. most existing datasets for code summarization contain one summary per one code snippet. for instance, code-searchnet @xcite ) contains two million function-documentation pairs across six programming languages (e.g., java, php, python). in contrast, our new dataset (notebookcdg) is designed for computational notebooks. the difference from previous cdg datasets is that in our dataset, a documentation text can correspond to several code snippets. ## notebookcdg dataset cdg for notebooks is a relatively new task. to our best knowledge, we could not find an appropriate dataset for this task. thus, we decided to construct a new dataset and share it with the community. publicly shared notebooks on github are often illdocumented @xcite , thus are not suitable for constructing the training dataset for cdg task. a recent work @xcite ) manually analyzed 80 publicly available notebooks on two kaggle challenges (i.e. out of 12,000 notebooks submitted to titanic and houseprice). kaggle allows community members to vote up and down on those notebooks, and @xcite 's findings show that the highly-voted notebooks are of good quality and quantity in code documentation. inspired by their work, we decided to utilize the top-voted and well-documented kaggle notebooks to construct the notebookcdg dataset 2 . we collected the top 10% highly-voted notebooks from the top 20 popular competitions on 2 we share the notebookcdg dataset with processed 28k code-document pairs at https://ibm.biz/bdfpk6 kaggle (e.g. titanic). we checked the data policy of each of the 20 competitions, none of them has copyright issues. we also contacted the kaggle administrators to make sure our data collection complies with the platform's policy. in total, we collected 3,944 notebooks as raw data. ## approach our model is built upon the standard encoderdecoder structure. to handle multiple code cells in computational notebooks, we propose a hierarchical attention mechanism based on convolutional graph neural network (haconvgnn) for capturing the relevant code cells during the decoding stage. the system architecture is illustrated in figure 1 . below, we describe each module in detail. ## automated evaluation we use rouge scores @xcite to evaluate our model's performance with regard to the groundtruth documentation content. we report rouge-1, rouge-2, and rouge-lcs (longest common sub-sequence). as shown in ablation study. in order to better understand the impact of the attention components in our model, we also perform an ablation study ( in general, we found that the hierarchical structure in our haconvgnn is proven to enhance our final performance. it is worth noting that the separated attention mechanism is essential in our model. remember that we use the attention mechanism for our four code cells separately. treating them as a single big code snippet leads to a considerable performance drop (see the last row in ## human evaluation we also conduct a human evaluation to further evaluate our model against the two baselines and the ground truth. participants. our human evaluation task involves reading code snippets and rating the generated documentation of the codes. we recruited participants with data science and machine learning backgrounds (@xmath0. task. we randomly selected 30 pairs of documentation and code(s) from our dataset. note that each pair has only one summary, but may have multiple code snippets. each participant is randomly assigned 10 trials, and the order of these 10 trials is also randomized. each pair is evaluated by 5 individuals. in each trial, a participant reads 4 candidate documentation for the same code snippet(s): three generated by the three models, and the other one is the groundtruth. participants do not know which documentation text is from which model. the participant is asked to rate the 4 documentation texts along three dimensions using a five-point likert-scale from -2 to 2. evaluation results. we conducted pairwise ttests to compare each model's performance. the result ( our model is also the only model that has a positive rating. for the informativeness dimension, groundtruth also has the best rating. our model (@xmath4.17) comes in second and outperforms code2seq (@xmath5.72, p<.01) and graph2seq (@xmath6.21, p<.01). for the readability dimension, in which we consider whether generated documentation is a valid english sentence or not, groundtruth outperforms all ml models again, but our model (@xmath7.67) also significantly outperforms baseline models code2seq (@xmath8.03 p<.01) and graph2seq (@xmath9.32 p<.01). our model can generate more readable documentation than baselines. model correctness informativeness readability groundtruth @xmath10.09, σ=0.95 @xmath11.85, σ=0.97 @xmath12.03, σ=1.01 our model @xmath13.21, σ=1.33 @xmath14.17, σ=1.18 @xmath15.67, σ=1.20 code2seq @xmath16.59, σ=1.29 @xmath17.72, σ=1.17 @xmath18.03, σ=1.35 graph2seq @xmath19.30, σ=1.40 @xmath20.21, σ=1.25 @xmath21.32, σ=1.35 all the results suggest that our model has abovezero ratings, which suggests it reaches an acceptable user satisfaction along all three dimensions. ## comparison with transformers we also carried out an additional experiment to compare our model with t5 @xcite , which is a state-of-the-art transformer encoderdecoder model. in order to fairly compare our model against t5, we do not use any pre-trained embeddings for the t5 model. also, t5 input has limitation for the input token length thus we did not feed ast hierarchy into it. more specifically, we initialize a t5-small model 6 with random weights and train this model using our training data. our code adapts the transformer models from hugging-face @xcite . we use the dev dataset to choose the hyperparameters and evaluate the trained model on our test dataset. the rouge f1 scores for the trained t5-small model are as follows: rouge-1 = 17.55, rouge-2 = 4.57, rouge-@xmath22.53. we found that the trained t5-small model achieves slightly better results than our model in rouge-1 and rouge-l. in practice, we found that the t5-small model relies on a much more hyperparameters and tends to generate less informative content compared to other models (see the documentation generated from different models in but in our dataset, as reported in ast token sequence is 1,732, which is too long as t5 input (512) or bart input (1,024). that is why t5 in sec 8 can only take the raw code sequence as input, instead of the ast hierarchy. it is known that programming code has a tree-based hierarchy and leveraging such ast hierarchy can enhance the baseline model (e.g., @xcite ). our contribution is that we provide a hierarchical attention architecture that is well suited for the programming code nature and can generalize to a much longer length of code inputs. imagine in a scenario where we can feed a whole code repo as training input by treating each code file as a lower layer, and connecting them through function/variable referencing -our architecture can also handle that. in general, we think our model is orthogonal to the standard transformer models. one interesting future work is to integrate our hierarchical attention mechanism into the transformer-based structure instead of a gru-based structure. ## downstream user application to demonstrate the application of the hacon-vgnn model, we designed a jupyter notebook plugin to assist document writing in data science programming (as shown in figure 4 ). the plugin is triggered when detecting users focusing on a code cell (figure 4 .a). the plugin then reads the contents from the focused cell and its adjacent cells, and sends the content to the backend. the backend server first generates a code summarization using the haconvgnn model (figure 4.b). in addition, we implemented two other approaches to generate documentation that was intended for explaining a design decision or explaining a technical concept for educational purposes. we retrieved the relevant documentation from the api webpage for educational purposes (figure 4 .c) and we used prompts to nudge users to explain an output (figure 4 .d). if the user likes one of these three candidates, they can simply click on one of them, and the selected documentation candidate will be inserted into above the code cell (if it describes what and why for the code), or below it (if it interprets the result of the code). our plugin went through several rounds of pilot testing and iterative design. participants found it reminds them to document code they would have ignored, reduce the time for developing documentation while they were actively exploring the data science task. the implementation details and a formal evaluation of understanding the benefits of the human-ai collaborative effort for automatic documentation generation are reported separately in @xcite . ## conclusion and future work this work targets a new application that aims to automatically generate code documentation (cdg) for a computational notebook. this project is part of our longterm research initiative of designing ai to automated the various tasks in an ai project's lifecycle @xcite . the notebookcdg context imposes unique challenges to the current code documentation generation approaches which only consider a single code snippet. we construct a dataset from kaggle challenge notebooks, and present a novel haconvgnn model to encode the multiple adjacent code cells as a hierarchical ast graph to enhance a sequence model architecture. both automated evaluation and human evaluation show that our model outperforms the baseline models. we also incorporate our algorithm into a jupyter notebook plugin to assist code documentation creation.