[Paper Review] HAConvGNN: Hierarchical Attention Based Convolutional Graph Neural Network for Code Documentation Generation in Jupyter Notebooks

Xuye Liu, Dakuo Wang+3 more|arXiv (Cornell University)|Mar 31, 2021

Software Engineering Research25 references4 citations

TL;DR

This paper proposes HAConvGNN, a hierarchical attention-based graph neural network that generates documentation for multiple code cells in Jupyter notebooks by modeling their structural relationships via abstract syntax trees (ASTs). It outperforms baseline models on ROUGE scores and human evaluations, demonstrating improved coherence and relevance in code documentation generation for multi-cell notebook contexts.

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.

Motivation & Objective

• To address the gap in code documentation generation (CDG) for multi-cell computational notebooks, where a single markdown cell often documents multiple code cells with structural dependencies.
• To develop a model that captures both token-level and cell-level dependencies across code cells using hierarchical attention mechanisms.
• To construct a large-scale, high-quality dataset (notebookCDG) of 28,625 code-documentation pairs from highly-ranked Kaggle notebooks for training and evaluation.
• To evaluate the model's performance in real-world settings through integration into a Jupyter Notebook plugin for human-AI collaboration in documentation writing.
• To demonstrate that structured modeling of code cell relationships improves the informativeness and relevance of generated documentation.

Proposed method

• The model uses a code sequence encoder and an auxiliary documentation text encoder to encode input code and predicted documentation tokens.
• It constructs a hierarchical graph representation from abstract syntax trees (ASTs) of code sequences, capturing structural relationships between code tokens and cells.
• A two-level hierarchical attention mechanism is employed: low-level attention attends to individual code tokens, and high-level attention attends to entire code cells in the AST structure.
• The graph neural network component integrates both semantic and structural information from the AST to enhance sequence-to-sequence generation.
• The model is trained end-to-end using a sequence-to-sequence framework with cross-entropy loss, optimized for fluency and relevance.
• The approach is integrated into a Jupyter Notebook plugin that suggests documentation based on focused code cells and their neighbors, supporting user interaction.

Experimental results

Research questions

• RQ1 Can a graph-based neural network with hierarchical attention effectively model multi-cell code documentation in Jupyter notebooks?
• RQ2 How does modeling AST structure improve the quality of generated documentation compared to sequence-only models?
• RQ3 To what extent does the HAConvGNN model outperform baseline models in automated metrics like ROUGE and human evaluation?
• RQ4 Can the model be effectively deployed in a real-world user-facing application to support human-AI collaboration in documentation writing?
• RQ5 Does the model generate more informative and contextually relevant documentation than standard sequence models like T5-small?

Key findings

• HAConvGNN outperforms baseline models on ROUGE-1 (21.83), ROUGE-2 (7.21), and ROUGE-L (23.45) metrics on the notebookCDG dataset.
• Human evaluation shows that HAConvGNN-generated documentation is rated significantly higher in relevance, coherence, and informativeness compared to baselines.
• The T5-small model achieved higher ROUGE-1 and ROUGE-L scores but generated less informative content, indicating a trade-off between automatic metrics and human perception.
• The model effectively handles long input sequences (up to 1,732 AST tokens), demonstrating scalability beyond standard transformer input limits (512/1,024 tokens).
• Integration into a Jupyter plugin led to user feedback indicating that the AI suggestions reminded users to document code they would have otherwise omitted.
• The human-AI collaboration plugin reduced documentation authoring time and improved user satisfaction during data science workflows.