Clustering Java Code

The Data

We downloaded 16Gb of Java code from Github, then divided it into code snippets of 2 kinds:

• Method snippets:  Java classes were parsed (using the tree-sitter library) to extract methods. Then, we applied a filtering procedure to make sure that there are no duplicates or methods that are too short (can’t be under 5 lines)
• Class snippets: They contain entire Java files

We are going to cluster both kinds of snippets separately. This will yield different results: method clusters are going to be more specific, whereas class clusters should cover entire technological domains.

BERTopic

BERTopic is a Python clustering library that will help us to analyze our data.

It takes a list of string documents and then calculates the corresponding embeddings using a BERT model (if custom embeddings are not provided). It then reduces the dimensionality of the embeddings with UMAP and performs clustering using the HDBSCAN algorithm. The most important feature of HDBSCAN is that it’s hierarchical - it also produces the binary tree of resulting clusters. 

With BERTopic, we can easily get valuable insights on clustering results such as:

• The list of clusters sorted by their size. Clusters are named by their top keywords put together
• Clusters hierarchy plot
• Top keywords associated with each cluster. which are words in the cluster documents with the highest TF-IDF metric (these words occur in the cluster documents way more often than in English texts in general)
• Some representatives of each cluster 

UniXCoder

Documents that we’re going to pass to BERTopic are code snippets. There are two options: either let BERTopic calculate the embeddings by itself (using a BERT model), or pass our own custom embeddings. In our research, we are dealing with the source code, not the natural language, so it would be much better to find an encoding model that is specifically designed to process the code and pass its custom embedding to the BERTopic.

One of the latest such models is UniXCoder, released in July 2022. The most important feature of UniXCoder is that it first builds up the syntax tree of the code input, and then rearranges the tree into the token sequence. All the details can be found in this paper. 

In our research, we’re going to only use UniXCoder encoding capabilities (encode a code snippet into an embedding). However, you can do some other cool stuff with it too. For instance,

• It can handle code comments as well. If you encode the code snippet and its corresponding comment, the embeddings should be somewhat similar
• In the decoder mode, UniXCoder can autocomplete the source code. However, please keep in mind that the model itself is quite small, so don’t expect miracles from it
• In the decoder-encoder mode, the model can perform the mask prediction (which can even be more than 1 token long)

Documents Preprocessing

BERTopic is designed to work with natural language documents, whereas the source code has a different structure which can break the BERTopic keyword extraction. Before passing documents into BERTopic, we have turned them into sequences of English words (our tokenization algorithm also splits camelCase identifiers into separate words).

For example, the following method:

public ArrayList<Player> AddPlayer(Player p){

        this.players.add(p);

        return this.players;

 }

turns into the sequence “public array list player add player player p this players add p return this players”

There is some vocabulary that is very typical for those code documents but not that frequent in English overall:  words like public, return, and string. In order for them to not dominate keyword lists, we tokenized all the code snippets and removed from them those words that appear in the top 100 frequency list.