Graph-Augmented Code Summarization in Computational Notebooks

1 Introduction

Documenting the story behind code and results is critical for data scientists to work effectively with others, as well as their future selves [Zhang et al., 2020; Kery and Myers, 2017]. The story, code, and computational results together construct a computational narrative. Data scientists find these computational notebooks particularly useful as they can support rapid exploration with code and explanation with natural language [Rule et al., 2018].

Unfortunately, the ease of use of computational notebooks also comes with a cost. Data scientists often write fragmented and drafty code in computational notebooks during their quick experimentation of testing hypotheses or alternatives. It is a tedious process for data scientists to then manually document and refactor the raw notebook into a more readable computational narrative, thus many people neglect to do so [Wang et al., 2021a].

Natural language processing (NLP) researchers have started exploring various ways to automatically generate code summarizations and documentation using source code as input [LeClair and McMillan, 2019; LeClair et al., 2020]. But the documentation in data science notebooks are particularly challenging for two reasons: 1) documentations in a notebook may explain the rationale of a piece of code, or may interpret the outcomes of a table or a chart, which is difficult for existing automated code summarization techniques; 2) even for the documentation simply about what code snippet does, it is still challenging than other coding contexts, due to the unique nature of the notebook that one documentation markdown cell may cover multiple code cells.

∗Contact Author

In this paper, we propose Themisto, an end-to-end automated documentation generation system that can support users to create various documentations in a computational notebook. It has a frontend user interface that is integrated into the Jupyter Notebook environment (Fig. 1) thus users do not need to leave their familiar coding environment; its backend is built on top of a three-approach framework to tackle the two aforementioned challenges: 1) to generate the documentation that is not easy for an automation algorithm, we devise a query-based approach and a prompt-based approach so that the users can semi-automatically write documentations themselves; 2) to generate documentation for multiple code blocks, we consider each code block as a graph, and together as a hierarchical multi-graph structure and propose an attention-based hierarchical ConvGNN component to augment a seq2seq network to solve this problem. The HAConvGNN algorithm framework is reported in [Liu et al., 2021].

In summary, we present a three-approach automated code documentation system to help Jupyter Notebook users to create documentation for their code and results. Builds on our previously reported HAConvGNN algorithm, the system has a human-in-the-loop design feature that can make the human users' work easier for the difficult-to-generate documentation.

2 Themisto

We design and implement Themisto as a Jupyter Notebook extension that supports data scientists to write betterdocumented computational narratives.

2.1 SystemArchitecture

The Themisto system has two components: the client-side UI implemented as a Jupyter Notebook plugin using TypeScript code, and the server-side backend implemented as a server using Python and Flask.

The client-side program is responsible for rendering the user interface and monitoring the user actions on the notebook to edits in code cells. When the user's cursor focusing on a code cell, the UI will send the current code cell content to the server-side program through HTTP requests.

The server-side program takes the code content and generates documentation using both the deep-learning-based approach and the query-based approach. For the deep-learningbased approach, the server-side program first tokenizes the code content and generates the AST. It then generates the prediction with the pre-trained model. For the query-based approach, the server-side program matches the curated API calls with the code snippets and returns the pre-collected descriptions. For the prompt-based approach, the server-side program sends different prompts (e.g., for interpreting results or for explaining reason) base on the output type of the code cell.

2.2 User Interface Design

Figure 1 shows the user interface of Themisto as a Jupyter Notebook plugin. Each time the user changes their focus on a code cell, as they may be inspecting or working on the cell, the plugin is triggered. The plugin sends the user-focused code cell's content to the backend. Using this content, the backend generates a code summarization using the model and retrieves a piece of documentation from the API webpage. When such a documentation generation process is done, the generated documentation is sent from the server-side to the frontend, and a light bulb icon appears next to the code cell, indicating that there are recommended markdown cells for the selected code cell (as shown in Figure 1.A).

When a user clicks on the light bulb icon, they can see three options rendered in the dropdown menu: (1) the HAConvGNN approach to generate documentation for source code (Figure 1.B); (2) a query-based approach to retrieve the online API documentation for source code (Figure 1.C); and (3) a user prompt approach to nudge users to write more documentation (Figure 1.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 why for the code), or below it (if it interprets the result of the code).

3 Three Approaches for Documentation Generation

In this section, we briefly describe the three different approaches for documentation generation.

3.1 HAConvGNNModel Approach

We proposed an attention-based hierarchical ConvGNN component to augment a seq2seq model (HAConvGNN), inspired by GNN architecture in [LeClair et al., 2020]. These GNN models can take both the source code's structure (extracted as AST) and the source code's content as input, in comparison to the traditional sequence-to-sequence model architectures, which only take the source code's content as an input sequence. Because a markdown cell can cover multiple, we augment LeClair et al. [LeClair et al., 2020] by encoding adjacent four code cells under a markdown cell to gain a highlevel attention weight.

To build a training dataset, we collect the top 10% of the publicly available notebooks from the top 20 popular Kaggle competitions 1 . Then we remove non-English notebooks and transform the data to follow the data structure in [LeClair et al., 2020]. After data preprocessing, our dataset has 28,625

1 We checked these challenges' data policy and contacted the Kaggle administrators to make sure our data collection process complies with the Kaggle policy.

code(s)-summary pairs. Following the best practice of model training, we split the dataset into training, testing, and validation subsets with an 8 to 1 to 1 ratio. We use the Adamax optimizer [Kingma and Ba, 2014] with a batch size of 20. The learning rate we set is 0.001. The code sequence embedding size is 100. In the encoder, we use GRU [Cho et al., 2014] with the hidden size of 256. The hop size of our GNN is 2. The dropout rate of our attention layer is 0.5.

In our previously published algorithm paper, we reported both quantitative and qualitative evaluations of our model's performance against two baseline models using ROUGE scores [Lin, 2004]. The two baseline models are Code2Seq model [Alon et al., 2018] and Graph2Seq model [Xu et al., 2018] . In this paper, we briefly present the partial quantitative evaluation result. As shown in Table 1, our HAConvGNN model outperforms the other two baselines. For more details, please refer to [Liu et al., 2021].

3.2 Query-Based Approach and Prompt-Based Approach

Well-documented Kaggle notebooks often have the description of frequently-used data science code functions for educational purposes. And sometimes data scientists directly paste in a link or a reference to the external API documentation for a code function [Wang et al., 2021a]. Thus, we implement a query-based approach that curates a list of API from commonly used data science packages, and short descriptions

from external documentation sites. In our system, we only cover Pandas , Numpy , and Scikit-learn these three libraries as a starting point to explore this approach. We collected both the API names and the short descriptions by building a crawling script with Python. When users trigger this query-based approach for a code cell, Themisto matches the API names with the code snippets and concatenate all the corresponding descriptions.

Lastly, a well-documented notebook not only documents the process of the code, but also interprets the output, and explains rationales. These types of documentation are hard to generate with automated solutions To achieve it, we implement a prompt-based approach. It detects whether the code cell has a cell output or not: if the cell outputs a result, Themisto assumes that the user is more likely to add interpretation for the output result, thus the corresponding prompt will be inserted below the code cell. Otherwise, the system assumes the user may want to insert a reason or some educational type of documentations, thus it changes its prompt message.

4 Conclusion

In conclusion, this demo paper presents a prototype system that can support Jupyter notebook users to utilize our stateof-the-art NLP algorithm [Liu et al., 2021] to automatically generate documentation for a computational notebook. In the future, we will continue to pursue the intertwined NLP and HCI research agenda [Wang et al., 2021b]: we will design new algorithm models to improve the fully automated documentation generation accuracy, and we will design and evaluate our human-centered AI system's usability with target users.